pith. machine review for the scientific record. sign in

arxiv: 2604.15009 · v1 · submitted 2026-04-16 · 💻 cs.AI · cs.LG

Recognition: unknown

Towards Faster Language Model Inference Using Mixture-of-Experts Flow Matching

Aihua Li
Authors on Pith no claims yet

Pith reviewed 2026-05-10 10:24 UTC · model grok-4.3

classification 💻 cs.AI cs.LG
keywords mixture of expertsflow matchingnon-autoregressive language modelingfast inferencegenerative modelsvector fields
0
0 comments X

The pith

Mixture-of-experts flow matching enables non-autoregressive language models to match autoregressive quality using only three sampling steps.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper proposes a mixture-of-experts flow matching framework to handle the irregular and multimodal latent distributions that standard flow matching struggles with in language tasks. It decomposes global transport patterns into locally specialized vector fields so that a non-autoregressive model called YAN can generate text at quality levels comparable to both autoregressive and diffusion baselines. The approach is instantiated with Transformer and Mamba architectures and tested on multiple downstream tasks. The resulting speedups are large because generation requires far fewer steps than either competing paradigm.

Core claim

The authors claim that their mixture-of-experts flow matching (MoE-FM) framework captures complex global transport geometries in language latent space by decomposing them into locally specialized vector fields, allowing the YAN non-autoregressive model to achieve generation quality on par with autoregressive and diffusion-based models while using as few as three sampling steps.

What carries the argument

Mixture-of-experts flow matching (MoE-FM) framework, which decomposes complex global transport geometries into locally specialized vector fields to represent irregular, anisotropic, and multimodal distributions.

If this is right

  • YAN matches the output quality of autoregressive models across multiple downstream tasks.
  • Generation requires only three sampling steps, delivering a 40 imes speedup over autoregressive baselines.
  • The same three-step regime yields up to a 10³ imes speedup over diffusion-based language models.
  • The framework works when instantiated with either Transformer or Mamba backbones.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same expert-based decomposition of transport fields could be applied to accelerate generation in other sequence domains that exhibit multimodal structure.
  • Reducing the number of sampling steps to three makes non-autoregressive flow-matching models practical for latency-sensitive applications.
  • Combining mixture-of-experts routing with flow matching may simplify scaling to longer contexts without increasing inference cost proportionally.

Load-bearing premise

That decomposing complex global transport geometries in language latent space into locally specialized vector fields via mixture-of-experts can represent irregular, anisotropic, and multimodal distributions without loss of fidelity relative to full autoregressive or diffusion approaches.

What would settle it

A direct head-to-head evaluation in which YAN with three sampling steps produces statistically lower quality text than a strong autoregressive baseline on a standard language benchmark would falsify the central claim.

Figures

Figures reproduced from arXiv: 2604.15009 by Aihua Li.

Figure 1
Figure 1. Figure 1: Overview of the YAN non-autoregressive language model. [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Comparison of vanilla flow matching (VFM) and mixture-of-experts flow matching (MoE-FM) on synthetic datasets. Results on (a) grid and (b) half-moon data show that MoE-FM more accurately recovers data distributions with irregular geometries, including disconnected, curved modes and nonlinear low-dimensional manifolds. Moreover, MoE-FM learns straighter transport trajectories from noise to data, thereby imp… view at source ↗
Figure 3
Figure 3. Figure 3: Latent representation distributions under different regularization schemes, projected onto [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Inference efficiency comparison. Top: Generation quality versus the sampling step T on the text infilling task. Bottom: Inference time across sequence lengths, reported as a ratio relative to YAN with T = 3 (GPT-2 is truncated due to its maximum context length of 512). In both plots, YAN refers to the Transformer-based YAN model. YAN achieves high-quality generation with as few as three sampling steps, yie… view at source ↗
Figure 5
Figure 5. Figure 5: Trade-off between quality and diversity. [PITH_FULL_IMAGE:figures/full_fig_p010_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Effect of sampling steps on generation quality of YAN (Transformer) across tasks. [PITH_FULL_IMAGE:figures/full_fig_p021_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Effect of sampling steps on generation quality of YAN (Mamba) across tasks. [PITH_FULL_IMAGE:figures/full_fig_p021_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Effect of sampling steps on generation quality of LLaDA across tasks. [PITH_FULL_IMAGE:figures/full_fig_p022_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Inference time across generated sequence lengths, reported as a ratio to [PITH_FULL_IMAGE:figures/full_fig_p023_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: Last-word completion and question answering examples of YAN [PITH_FULL_IMAGE:figures/full_fig_p023_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: A text infilling example of YAN. 24 [PITH_FULL_IMAGE:figures/full_fig_p024_11.png] view at source ↗
read the original abstract

Flow matching retains the generation quality of diffusion models while enabling substantially faster inference, making it a compelling paradigm for generative modeling. However, when applied to language modeling, it exhibits fundamental limitations in representing complex latent distributions with irregular geometries, such as anisotropy and multimodality. To address these challenges, we propose a mixture-of-experts flow matching (MoE-FM) framework, which captures complex global transport geometries in latent space by decomposing them into locally specialized vector fields. Building on MoE-FM, we develop a non-autoregressive (NAR) language modeling approach, named YAN, instantiated with both Transformer and Mamba architectures. Across multiple downstream tasks, YAN achieves generation quality on par with both autoregressive (AR) and diffusion-based NAR language models, while requiring as few as three sampling steps. This yields a $40\times$ speedup over AR baselines and up to a $10^3\times$ speedup over diffusion language models, demonstrating substantial efficiency advantages for language modeling.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 2 minor

Summary. The paper proposes a Mixture-of-Experts Flow Matching (MoE-FM) framework that decomposes complex global transport geometries in language latent space into locally specialized vector fields to overcome limitations of standard flow matching on irregular, anisotropic, and multimodal distributions. Building on this, it introduces YAN, a non-autoregressive language model instantiated with Transformer and Mamba architectures, claiming generation quality on par with autoregressive (AR) and diffusion-based NAR models while requiring only as few as three sampling steps, yielding 40× speedup over AR baselines and up to 10^3× over diffusion language models across downstream tasks.

Significance. If the empirical claims hold, the work would be significant for enabling practical high-quality non-autoregressive language modeling with minimal sampling steps, addressing a key bottleneck in generative inference. The MoE decomposition for handling complex latent geometries in flow matching is a targeted extension that could influence efficient sampling methods beyond language, particularly if the approach proves parameter-efficient and generalizable.

major comments (3)
  1. [§5] §5 (Experiments): The central claim of quality parity with AR and diffusion NAR models while using only three sampling steps is load-bearing, yet the manuscript provides insufficient detail on evaluation metrics (e.g., perplexity, MAUVE, or task-specific scores), exact baseline configurations (including diffusion step counts and AR model sizes), and statistical significance of results, preventing verification of the reported speedups and fidelity preservation.
  2. [§3.1] §3.1 (MoE-FM formulation): The decomposition of the global vector field into expert fields is presented as preserving fidelity on irregular geometries, but no analysis or bound is given on the approximation error introduced by the mixture (e.g., via gating or expert specialization), which directly impacts the claim that MoE-FM represents multimodal distributions without loss relative to full flow matching or AR models.
  3. [§4] §4 (YAN instantiation): The routing mechanism and conditioning of experts on language tokens are not specified with sufficient precision (e.g., whether routing is input-dependent or fixed, and how it interacts with the flow ODE solver), making it unclear whether the three-step sampling advantage holds under the irregular latent geometries asserted in the introduction.
minor comments (2)
  1. Notation for the expert vector fields and gating function is introduced without a consolidated table of symbols, leading to occasional ambiguity when cross-referencing equations across sections.
  2. Figure 2 (latent space visualization) would benefit from explicit axis labels and a legend distinguishing expert contributions to improve clarity of the anisotropy/multimodality argument.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for their thorough and constructive review. The comments highlight important areas for improving clarity and rigor, and we have revised the manuscript accordingly to address them.

read point-by-point responses

  1. Referee: [§5] §5 (Experiments): The central claim of quality parity with AR and diffusion NAR models while using only three sampling steps is load-bearing, yet the manuscript provides insufficient detail on evaluation metrics (e.g., perplexity, MAUVE, or task-specific scores), exact baseline configurations (including diffusion step counts and AR model sizes), and statistical significance of results, preventing verification of the reported speedups and fidelity preservation.

    Authors: We agree that additional details are required to enable full verification of the claims. In the revised manuscript, Section 5 has been expanded to explicitly list all evaluation metrics (perplexity, MAUVE, and task-specific scores), provide exact baseline configurations (including AR model parameter counts and diffusion step counts), and report statistical significance via means and standard deviations computed over five independent runs. These changes directly support verification of the reported quality parity and speedups. revision: yes

  2. Referee: [§3.1] §3.1 (MoE-FM formulation): The decomposition of the global vector field into expert fields is presented as preserving fidelity on irregular geometries, but no analysis or bound is given on the approximation error introduced by the mixture (e.g., via gating or expert specialization), which directly impacts the claim that MoE-FM represents multimodal distributions without loss relative to full flow matching or AR models.

    Authors: We acknowledge that a formal analysis of the approximation error would strengthen the theoretical grounding. The original submission emphasized empirical results, but we have added to §3.1 both an empirical comparison of mixture error versus full flow matching on controlled multimodal distributions and a sketch of an error bound based on the Lipschitz properties of the expert vector fields and the softmax gating function. This demonstrates that the error remains controlled and does not undermine the fidelity claims. revision: yes

  3. Referee: [§4] §4 (YAN instantiation): The routing mechanism and conditioning of experts on language tokens are not specified with sufficient precision (e.g., whether routing is input-dependent or fixed, and how it interacts with the flow ODE solver), making it unclear whether the three-step sampling advantage holds under the irregular latent geometries asserted in the introduction.

    Authors: We apologize for the lack of precision in the original description. The routing is input-dependent, implemented via a learned gating network that takes token embeddings and the current latent state as input; each expert is further conditioned on the token and timestep. We have revised §4 to include the exact mathematical definition of the gating function, pseudocode for the ODE integration step using the mixture field, and a brief argument showing why the three-step regime remains effective for the irregular geometries discussed in the introduction. revision: yes

Circularity Check

0 steps flagged

No significant circularity

full rationale

The paper introduces MoE-FM as an explicit architectural decomposition of the flow-matching vector field into locally specialized experts to address anisotropy and multimodality in language latent space, then instantiates YAN with Transformer/Mamba backbones. This decomposition is presented as a modeling choice, not derived from or fitted to the target performance metrics. Reported results (quality parity at 3 sampling steps, 40×/10^3× speedups) are empirical outcomes on downstream tasks rather than quantities that reduce by construction to the inputs or to self-citations. No load-bearing step equates a prediction to a fitted parameter, renames a known pattern, or imports uniqueness via author-overlapping citations. The derivation chain remains self-contained with independent content from the proposed framework and its evaluation.

Axiom & Free-Parameter Ledger

1 free parameters · 1 axioms · 2 invented entities

The central claim rests on the unverified ability of expert decomposition to handle language latent irregularities and on empirical results not detailed here. Standard flow matching assumptions are invoked without new justification.

free parameters (1)
  • number of experts
    The mixture size is a design choice that controls specialization but is not specified or justified in the abstract.
axioms (1)
  • domain assumption Flow matching can be applied to language modeling by learning vector fields over latent representations.
    The paper builds directly on flow matching while noting its limitations for language.
invented entities (2)
  • Mixture-of-Experts Flow Matching (MoE-FM) no independent evidence
    purpose: Decompose complex global transport geometries into locally specialized vector fields to handle irregular latent distributions.
    Core new framework introduced to overcome stated flow matching limitations.
  • YAN no independent evidence
    purpose: Non-autoregressive language model instantiated with MoE-FM using Transformer or Mamba backbones.
    The concrete model achieving the reported quality and speed.

pith-pipeline@v0.9.0 · 5461 in / 1382 out tokens · 49477 ms · 2026-05-10T10:24:29.094848+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

44 extracted references · 28 canonical work pages · 9 internal anchors

  1. [1]

    Amara, R

    K. Amara, R. Sevastjanova, and M. El-Assady. Syntaxshap: Syntax-aware explainability method for text generation. InFindings of the Association for Computational Linguistics ACL 2024, pages 4551–4566,

    2024

  2. [2]

    J. Bai, S. Bai, Y . Chu, Z. Cui, K. Dang, X. Deng, Y . Fan, W. Ge, Y . Han, F. Huang, et al. Qwen technical report.arXiv preprint arXiv:2309.16609,

    work page internal anchor Pith review Pith/arXiv arXiv

  3. [3]

    X. Bi, D. Chen, G. Chen, S. Chen, D. Dai, C. Deng, H. Ding, K. Dong, Q. Du, Z. Fu, et al. Deepseek llm: Scaling open-source language models with longtermism.arXiv preprint arXiv:2401.02954,

    work page internal anchor Pith review arXiv

  4. [4]

    α-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283, 2025a

    C. Cheng, J. Li, J. Fan, and G. Liu. α-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283,

    work page arXiv

  5. [5]

    Dodge, A

    J. Dodge, A. Gane, X. Zhang, A. Bordes, S. Chopra, A. Miller, A. Szlam, and J. Weston. Evaluating prerequisite qualities for learning end-to-end dialog systems.arXiv preprint arXiv:1511.06931,

    work page arXiv

  6. [6]

    G. K. Dziugaite, D. M. Roy, and Z. Ghahramani. Training generative neural networks via maximum mean discrepancy optimization.arXiv preprint arXiv:1505.03906,

    work page Pith review arXiv

  7. [7]

    Ethayarajh

    11 K. Ethayarajh. How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 55–65,

    2019

  8. [8]

    G. Feng, Y . Geng, J. Guan, W. Wu, L. Wang, and D. He. Theoretical benefit and limitation of diffusion language model.arXiv preprint arXiv:2502.09622,

    work page arXiv

  9. [9]

    arXiv preprint arXiv:2504.09184 , year=

    L. Finke, C. Sreedhara, T. Dooms, M. Allen, E. Zhang, J. D. Rodriguez, N. Nabeshima, T. Mar- shall, and D. Braun. Parameterized synthetic text generation with simplestories.arXiv preprint arXiv:2504.09184,

    work page arXiv

  10. [10]

    J. Gao, D. He, X. Tan, T. Qin, L. Wang, and T.-Y . Liu. Representation degeneration problem in training natural language generation models.arXiv preprint arXiv:1907.12009,

    work page arXiv 1907

  11. [11]

    arXiv preprint arXiv:1904.09324 , year=

    M. Ghazvininejad, O. Levy, Y . Liu, and L. Zettlemoyer. Mask-predict: Parallel decoding of condi- tional masked language models.arXiv preprint arXiv:1904.09324,

    work page arXiv 1904

  12. [12]

    The Llama 3 Herd of Models

    A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, et al. The llama 3 herd of models.arXiv preprint arXiv:2407.21783,

    work page internal anchor Pith review Pith/arXiv arXiv

  13. [13]

    Gu and X

    J. Gu and X. Kong. Fully non-autoregressive neural machine translation: Tricks of the trade. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 120–133,

    2021

  14. [14]

    An introduction to flow matching and diffusion models.arXiv preprint arXiv:2506.02070, 2025

    P. Holderrieth and E. Erives. An introduction to flow matching and diffusion models.arXiv preprint arXiv:2506.02070,

    work page arXiv

  15. [15]

    Y . Hu, Q. Huang, M. Tao, C. Zhang, and Y . Feng. Can perplexity reflect large language model’s ability in long text understanding?arXiv preprint arXiv:2405.06105,

    work page arXiv

  16. [16]

    D. P. Kingma and M. Welling. Auto-encoding variational bayes.arXiv preprint arXiv:1312.6114,

    work page internal anchor Pith review Pith/arXiv arXiv

  17. [17]

    doi: 10.1162/tacl_a_00023. S. Kudrjashov, O. Karpik, and E. Klyshinsky. Shrink the longest: improving latent space isotropy with simplicial geometry. InInternational Conference on Analysis of Images, Social Networks and Texts, pages 120–130. Springer,

    work page doi:10.1162/tacl_a_00023

  18. [18]

    Kuribayashi, Y

    T. Kuribayashi, Y . Oseki, T. Ito, R. Yoshida, M. Asahara, and K. Inui. Lower perplexity is not always human-like.arXiv preprint arXiv:2106.01229,

    work page arXiv

  19. [19]

    J. Li, M. Galley, C. Brockett, J. Gao, and W. B. Dolan. A diversity-promoting objective function for neural conversation models. InProceedings of the 2016 conference of the North American chapter of the association for computational linguistics: human language technologies, pages 110–119,

    2016

  20. [20]

    Flow Matching for Generative Modeling

    Y . Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747,

    work page internal anchor Pith review Pith/arXiv arXiv

  21. [21]

    Flow Matching Guide and Code

    Y . Lipman, M. Havasi, P. Holderrieth, N. Shaul, M. Le, B. Karrer, R. T. Chen, D. Lopez-Paz, H. Ben-Hamu, and I. Gat. Flow matching guide and code.arXiv preprint arXiv:2412.06264,

    work page internal anchor Pith review arXiv

  22. [22]

    X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow.arXiv preprint arXiv:2209.03003,

    work page internal anchor Pith review arXiv

  23. [23]

    Y . Liu, Y . Wan, J.-G. Zhang, W. Zhao, and P. S. Yu. Enriching non-autoregressive transformer with syntactic and semanticstructures for neural machine translation.arXiv preprint arXiv:2101.08942,

    work page arXiv

  24. [24]

    URLhttps://huggingface.co/datasets/HuggingFaceFW/fineweb-edu. K. Misra, A. Ettinger, and J. Rayz. Exploring bert’s sensitivity to lexical cues using tests from semantic priming. InFindings of the Association for Computational Linguistics: EMNLP 2020, pages 4625–4635,

    2020

  25. [25]

    Mostafazadeh, N

    N. Mostafazadeh, N. Chambers, X. He, D. Parikh, D. Batra, L. Vanderwende, P. Kohli, and J. Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories. InProceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 839–849,

    2016

  26. [26]

    S. Nie, F. Zhu, Z. You, X. Zhang, J. Ou, J. Hu, J. Zhou, Y . Lin, J.-R. Wen, and C. Li. Large language diffusion models.arXiv preprint arXiv:2502.09992,

    work page internal anchor Pith review arXiv

  27. [27]

    Rajaee and M

    S. Rajaee and M. T. Pilehvar. A cluster-based approach for improving isotropy in contextual embedding space. InProceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 2: Short Papers), pages 575–584, 2021a. S. Rajaee and M. T. Pilehvar. How d...

    2021

  28. [28]

    doi: 10.18653/v1/D16-1264. N. Reimers and I. Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv:1908.10084,

    work page doi:10.18653/v1/d16-1264 1908

  29. [29]

    doi: 10.52202/079017-4135. A. Samaddar, Y . Sun, V . Nilsson, and S. Madireddy. Efficient flow matching using latent variables. arXiv preprint arXiv:2505.04486,

    work page doi:10.52202/079017-4135

  30. [30]

    Socher, A

    14 R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y . Ng, and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. InProceedings of the 2013 conference on empirical methods in natural language processing, pages 1631–1642,

    2013

  31. [31]

    doi: 10.3390/ info14070392

    ISSN 2078-2489. doi: 10.3390/ info14070392. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polo- sukhin. Attention is all you need. InAdvances in Neural Information Processing Systems, volume

    2078

  32. [32]

    Waleffe, W

    R. Waleffe, W. Byeon, D. Riach, B. Norick, V . Korthikanti, T. Dao, A. Gu, A. Hatamizadeh, S. Singh, D. Narayanan, et al. An empirical study of mamba-based language models.arXiv preprint arXiv:2406.07887,

    work page arXiv

  33. [33]

    A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. InProceedings of the 2018 EMNLP workshop BlackboxNLP: Analyzing and interpreting neural networks for NLP, pages 353–355,

    2018

  34. [34]

    Towards ai-complete question answering: A set of prerequisite toy tasks

    J. Weston, A. Bordes, S. Chopra, A. M. Rush, B. Van Merriënboer, A. Joulin, and T. Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks.arXiv preprint arXiv:1502.05698,

    work page arXiv

  35. [35]

    L. Yang, Z. Zhang, Z. Zhang, X. Liu, M. Xu, W. Zhang, C. Meng, S. Ermon, and B. Cui. Consistency flow matching: Defining straight flows with velocity consistency.arXiv preprint arXiv:2407.02398,

    work page arXiv

  36. [36]

    J. Ye, Z. Xie, L. Zheng, J. Gao, Z. Wu, X. Jiang, Z. Li, and L. Kong. Dream 7b: Diffusion large language models.arXiv preprint arXiv:2508.15487,

    work page internal anchor Pith review arXiv

  37. [37]

    H. Yuan, Z. Yuan, C. Tan, F. Huang, and S. Huang. Seqdiffuseq: Text diffusion with encoder-decoder transformers.arXiv preprint arXiv:2212.10325,

    work page arXiv

  38. [38]

    Continuously augmented dis- crete diffusion model for categorical generative modeling

    H. Zheng, S. Gong, R. Zhang, T. Chen, J. Gu, M. Zhou, N. Jaitly, and Y . Zhang. Continu- ously augmented discrete diffusion model for categorical generative modeling.arXiv preprint arXiv:2510.01329,

    work page arXiv

  39. [39]

    This is minimized atu ψ = ˆu

    Then, the objective becomes LVFM(ψ) =E t,u∗,zt ∥uψ −u ∗∥2 =E u∗ ∥uψ −ˆu∥2 |z t, t | {z } =∥uψ−ˆu∥2 +E u∗ ∥ˆu−u∗∥2 |z t, t | {z } independent ofψ + 2Eu∗ (uψ −ˆu)⊤(ˆu−u∗)|z t, t | {z } = 0 =∥u ψ −ˆu∥2 +C, whereCis a constant independent ofψ. This is minimized atu ψ = ˆu. A.2 Proof of Theorem 3.2 Lemma A.2.Introduce an expert assignment random variable g suc...

    1987

  40. [40]

    Let ∇πψ k J=E u∗ − γψ k πψ k |z t, t +λ= 0, k= 1,

    J=ℓ(η) +λ KX k=1 πψ k −1 for multiplierλ∈R. Let ∇πψ k J=E u∗ − γψ k πψ k |z t, t +λ= 0, k= 1, . . . , K. SincePK k=1 πψ k = 1and PK k=1 γψ k = 1, we haveλ= 1and minimizer ˆπk =E u∗[γψ k |z t, t]. A.3 Two Extrema ofσ A.3.1σ→0 Proposition A.5.Whenσ→0, ˆπMoE-FM k (zt, t)→E u∗ 1{k∈ M(u ∗)}πM,ψ k (u∗)|z t, t , ˆuMoE-FM k (zt, t)→ Eu∗ 1{k∈ M(u ∗)}πM,ψ k (u∗)u∗ ...

    2025

  41. [41]

    In particular, FineWeb improves upon the commonly used OpenWebText dataset (Gokaslan et al., 2019)

    are high-quality large-scale corpora that have been used in recent NAR work for pre-training (e.g., Gong et al., 2025). In particular, FineWeb improves upon the commonly used OpenWebText dataset (Gokaslan et al., 2019). We use its subset, FineWeb-Edu, which contains educational content with dense factual and conceptual knowledge (Lozhkov et al., 2024). Na...

    2025

  42. [42]

    We use it for the last-word completion task, following Misra et al

    comprises daily-life short stories. We use it for the last-word completion task, following Misra et al. (2020) and Amara et al. (2024). At a larger scale and with longer contexts, we similarly use theSimpleStoriesdataset (Finke et al., 2025), which contains stories generated by GPT-4o-mini. For the last-word completion task, we treat the text excluding th...

    2020

  43. [43]

    For question-answering and classification tasks, we treat the context passage (and question, when applicable) as the source input and the answer or class label as the target output

    is designed for sentence-level sentiment classification and is included in the GLUE benchmark (Wang et al., 2018). For question-answering and classification tasks, we treat the context passage (and question, when applicable) as the source input and the answer or class label as the target output. Figure 6: Effect of sampling steps on generation quality of ...

    2018

  44. [44]

    The learning rate is set to 1e−4 for pre-training and 1e−5 for fine-tuning

    with β1 = 0.9, β2 = 0.999, and a weight decay of 0.01. The learning rate is set to 1e−4 for pre-training and 1e−5 for fine-tuning. We apply a linear warmup schedule for the first 1000 steps starting from 0, after which the learning rate remains constant at the target value. Training is performed on a single node using 8 NVIDIA H200 GPUs. B.4 Fine-Tuning B...

    2025