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

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 LLMs, while requiring as few as three sampling steps. This yields a $40\times$ speedup over AR baselines and up to a $10<sup>3\times$ speedup over diffusion LLMs, demonstrating substantial efficiency advantages for language modeling.

• The paper introduces mixture-of-experts flow matching, which models multimodal latent transport with specialized vector fields and soft routing instead of collapsing velocity targets into one average field.
• YAN combines MoE flow matching with a latent non-autoregressive decoder, achieving strong generation and classification results at 200M parameters while requiring about three Euler sampling steps.
• The method delivers reported speedups of 40–50× over GPT-2 and BART and roughly 1,000× over LLaDA on long-document infilling, although scaling, sparse routing, and zero-shot performance remain untested.

Motivation and problem statement

Autoregressive (AR) LLMs incur one forward pass per generated token, while diffusion-based non-autoregressive (NAR) alternatives require hundreds to thousands of refinement steps to reach competitive quality, offsetting the benefit of parallel decoding. This paper investigates flow matching as a middle ground: a deterministic ODE-based generative paradigm whose trajectories can be trained to be nearly straight, enabling few-step sampling. The authors identify a concrete failure mode when vanilla flow matching (VFM) is applied to text latents: token representations exhibit anisotropy, multimodality, and manifold fragmentation, and a single global vector field regressed under an $\ell_2$ objective collapses multimodal conditional velocity distributions to their mean. The paper's central claim is that this limitation can be addressed by decomposing the transport into locally specialized expert vector fields.

Mixture-of-experts flow matching

The proposed mixture-of-experts flow matching (MoE-FM) replaces VFM's Gaussian approximation of the conditional velocity distribution $q_{\text{data}}(u^* \mid z_t, t)$ with a $K$-component Gaussian mixture whose means are expert vector fields and whose mixing weights come from a learned gating network $\pi^\psi(z_t, t)$. Training minimizes a negative log-likelihood loss rather than the $\ell_2$ regression loss. A theorem characterizes the conditional optima: each optimal expert field is a responsibility-weighted average of velocity targets, and the optimal routing is the expected posterior responsibility. The responsibilities $\gamma_k^\psi$ thus implement soft gating in vector-field space, allowing experts to specialize in distinct local transport geometries.

Two limiting cases are analyzed: as $\sigma \to 0$, routing converges to hard nearest-neighbor assignment; as $\sigma \to \infty$, the objective becomes uninformative and assignments non-identifiable. At sampling time, the method uses trajectory-level frozen routing—an expert is drawn once at $t=0$ and held fixed during ODE integration—which avoids expert switching and preserves geometric consistency along each trajectory. On synthetic grid and half-moon datasets, MoE-FM recovers disconnected and curved modes that VFM blurs across, and learns visibly straighter trajectories, enabling accurate generation in roughly four Euler steps.

The YAN model

YAN ("Flow Until You Almost Know") is a latent-variable NAR LLM, $$p(y \mid x) = \int p_\theta(y \mid z) p_\psi(z \mid x)\,dz$$, where the conditional latent generator $q_{\text{data}}(u^* \mid z_t, t)$0 is trained with MoE-FM. The design goal is a latent representation expressive enough that tokens become approximately conditionally independent given $q_{\text{data}}(u^* \mid z_t, t)$1, permitting fully parallel decoding through a lightweight decoder. An encoder contextualizes the source sequence and also provides target endpoints $q_{\text{data}}(u^* \mid z_t, t)$2 for self-supervised latent flow training; the authors report no benefit from teacher distillation in preliminary experiments.

Training proceeds in two stages. Stage 1 trains an asymmetric autoencoder with cross-entropy reconstruction plus two regularizers: an MMD penalty pushing the latent marginal toward isotropy (following the Wasserstein autoencoder rationale), an $q_{\text{data}}(u^* \mid z_t, t)$3 scale penalty, and Gaussian perturbation of encoder outputs during decoding. Visualizations confirm these regularizers yield more isotropic, origin-concentrated latents. Stage 2 trains the MoE-FM generator jointly with a cross-entropy term on decoded outputs. Notably, the authors observe that MoE-FM loss alone produces misaligned latents—well-modeled distributions that decode to incorrect tokens—and that the CE term is essential for anchoring the flow to the decoding objective. YAN is instantiated with both Transformer and bidirectional Mamba backbones (the latter augmented with explicit cross-attention), trained at the 200M-parameter scale on FineWiki/FineWeb.

Empirical results

Evaluation covers infilling (NarrativeQA), last-word completion (ROCStories, SimpleStories), QA (SQuAD, bAbI), and classification (AG News, DBpedia, SST-2), against GPT-2 (124M), BART (139M), and LLaDA-8B. Perplexity is deliberately excluded: NAR latent models lack tractable likelihoods, making cross-method comparisons unreliable.

Task / metric	GPT-2	BART	LLaDA-8B	YAN-M	YAN-TRF
NarrativeQA R-1/R-L/TPS	68.9/61.8/206	81.8/80.6/211	18.5/15.8/14	94.6/93.6/18.1k	94.9/93.9/20.6k
SimpleStories EM/BS-F1	42.8/85.1	46.5/90.3	21.6/20.6	59.7/91.1	65.5/93.9
ROCStories EM/BS-F1	28.5/79.7	21.3/70.6	5.1/-17.5	26.1/77.0	31.0/82.7
AG News acc.	93.8	91.2	92.1	95.1	97.2
DBpedia acc.	98.9	94.7	95.3	99.5	99.1
SST-2 acc.	90.1	88.0	90.7	87.4	91.0
SQuAD F1/BS-F1	48.0/41.0	78.9/76.7	88.8/87.2	70.8/71.3	80.4/78.2
bAbI F1/BS-F1	47.7/15.7	78.3/74.8	99.7/99.6	86.4/85.3	88.5/87.8

Three findings stand out. First, YAN achieves the best generation quality on infilling and completion despite being far smaller than LLaDA, which performs poorly on generation tasks at its reported operating points (e.g., negative BERTScore on ROCStories). Second, LLaDA retains the lead on QA tasks, which the authors attribute to its larger training scale and richer world knowledge—a candid acknowledgment that quality gaps remain where parametric knowledge matters. Third, the Transformer variant generally outperforms the Mamba variant, consistent with prior evidence that state-space models underperform attention on memory-intensive tasks.

On efficiency, YAN reaches high-quality long-document infilling with three Euler steps, yielding a 40–50× speedup over AR baselines (GPT-2, BART) and roughly $q_{\text{data}}(u^* \mid z_t, t)$4 over LLaDA, which requires approximately one step per token ($q_{\text{data}}(u^* \mid z_t, t)$5). Because the AR baselines are smaller than YAN, the speedup cannot be attributed to capacity differences. Sensitivity analysis shows quality degrades when $q_{\text{data}}(u^* \mid z_t, t)$6, whereas LLaDA improves monotonically with more steps. Diversity analysis shows the usual quality–diversity trade-off, but configurations exist where diversity exceeds baselines, adjustable via the number of sampling steps.

Limitations and open questions

Several caveats qualify the results. All experiments are at the 200M-parameter scale with task-specific fine-tuning; the paper does not demonstrate zero-shot or general-purpose capability, and scaling behavior of MoE-FM to larger models is untested. The current MoE formulation is dense, so inference cost grows linearly in the number of experts; sparse routing is proposed but not evaluated. The conditional-independence decoding assumption underlying parallel generation is an idealization whose degree of violation at scale is not quantified. The CE-anchoring term in stage 2 is motivated empirically rather than theoretically—the paper observes misalignment without it but offers only an interpretation. Finally, the efficiency comparison against LLaDA uses that model's best-quality setting ($q_{\text{data}}(u^* \mid z_t, t)$7); comparisons at matched quality–latency operating points across diffusion models would strengthen the claim.

Conclusion

The paper contributes a principled modification of flow matching—mixture-of-experts vector fields with provable conditional optima—that addresses the multimodality-induced failures of VFM on irregular text-latent geometries, and demonstrates that the resulting NAR model, YAN, matches or exceeds AR and diffusion baselines on generation tasks at 200M scale while requiring only three sampling steps. The headline efficiency result (40–50× over AR, ~$q_{\text{data}}(u^* \mid z_t, t)$8 over diffusion) rests on few-step ODE integration enabled by straighter trajectories, though its generality beyond the evaluated scale remains an open empirical question.