Training large-scale generative models is resource-intensive and relies heavily on heuristic dataset weighting. We address two fundamental questions: Can we train Large Language Models (LLMs) modularly, combining small, domain-specific experts to match monolithic performance, and can we do so robustly for any data mixture, eliminating heuristic tuning? We present a theoretical framework for modular generative modeling where a set of pre-trained experts are combined via a gating mechanism. We define the space of normalized gating functions G1 and formulate the problem as a minimax game to find a single robust gate that minimizes divergence to the worst-case data mixture. We prove the existence of such a robust gate using Kakutani's fixed-point theorem and show that modularity acts as a strong regularizer, with generalization bounds scaling with the lightweight gate's complexity. Furthermore, we prove that this modular approach can theoretically outperform models retrained on aggregate data, with the gap characterized by the Jensen-Shannon Divergence. Finally, we introduce a scalable Stochastic Primal-Dual algorithm and a Structural Distillation method for efficient inference. Empirical results on synthetic and real-world datasets confirm that our modular architecture effectively mitigates gradient conflict and can robustly outperform monolithic baselines.
Figures & tables
Figure 1: Conceptual Architecture of the Modular Gated Solution.
Figure 2: Visualizing the JSD Gap. A gated model (blue) fits distinct modes perfectly by routing inputs. A retrained model (red) suffers from capacity interference, forcing an entropy increase proportional to the JSD.
Figure 3: Efficiency Strategies. (A) Monolithic Distillation trains a single large model to mimic the ensemble, discarding the original experts. (B) Structural Distillation trains a lightweight Causal Router to mimic only the gating decisions ( g∗ ), preserving the original experts. This maintains modularity: upgrading an expert in (B) improves the system immediately.
Figure 4: Modularity overcomes gradient conflict. Left: The Robust Gate (blue) outperforms Fixed monolithic baselines (green) across all mixtures. Right: Even compared to Oracle models (red) trained on the exact test distribution, the Gate wins in the high-interference region ( λ≈0.5 ). The concave shape of the Oracle curves indicates destructive interference.
Figure 5: Structured Distillation. The Causal Router (cyan) effectively recovers the performance of the non-causal Robust Gate (blue), outperforming the Larger Fixed baseline (green).
Model
NLL (Diff seed)
NLL (Diff data)
Wiki Expert
5.122 ± 0.005
5.118 ± 0.011
Code Expert
4.722 ± 0.045
5.267 ± 0.788
FineWeb Expert
5.623 ± 0.004
5.623 ± 0.006
Retrained Model
5.133 ± 0.010
5.306 ± 0.257
Standard MoE (Scratch)
5.080 ± 0.009
5.245 ± 0.269
Gate Model (Ours)
4.994 ± 0.013
5.087 ± 0.141
Table 1: NLL per token on Real-World Data.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Geometric Interpretation of Lemma B.1 . The shaded gray region represents the convex hull of the experts π1 and π2 . The target distribution (red dashed) exceeds this feasible capacity. The optimal gate πg∗ (solid blue) traces the target where possible but is “clipped” to the expert boundaries when the target falls outside the hull.
Figure 7: Geometry of Prior Knowledge (Section C.1 ). The outer triangle represents the full probability simplex Δ . The green region Λ represents the subset of valid mixtures defined by prior knowledge. The algorithm projects the adversary’s updates (red point) back onto Λ (blue point), tightening the worst-case bound as per Theorem C.2 .
Figure 8: Dynamics of the Primal-Dual Game (Algorithm 3 ). The optimization is modeled as a 3-player game. The λ -player maximizes the mixture difficulty using Exponentiated Gradient. The g -player minimizes the robust loss. The μ -player enforces the global normalization constraint ( Zg=1 ) via Dual Ascent.
A(x)=Mq(x)πg∗(x)=pq(x)πg∗(x)
Appendix
Algorithm 5 Exact Sampling via Rejection Sampling
Figure 9: Modularity overcomes gradient conflict at a 50-50% mix of Domain A and a pure Domain B. Left figure illustrates a comparison to the Fixed Smaller and Larger models (in Green), while the right figure illustrates the comparison to the Oracle models (in Red). The Robust Gate (blue) and the Fixed models in the left figure naturally obtain the best performance for small values of λ where the test distribution is predominantly made up by Domain B. As λ increases, the test distribution contains more data from Domain A and gets harder for all models. Yet the Robust Gate maintains its clear advantage. The Oracle models in the right figure still has an advantage for the really skewed distribution and λ∼0 , but loses to the Robust Gate for larger values of λ .
Figure 10: Modularity overcomes gradient conflict at a 25-75% mix of Domain A and Domain B. Left figure illustrates a comparison to the Fixed Smaller and Larger models (in Green), while the right figure illustrates the comparison to the Oracle models (in Red). The Robust Gate (blue) and the Fixed models in the left figure naturally obtain the best performance for small values of λ where the test distribution is predominantly made up by Domain B. As λ increases, the test distribution contains more data from Domain A and gets harder for all models. Yet the Robust Gate maintains its clear advantage. Results are illustrated with lines for the mean values over 5 runs and standard deviations indicated with shaded regions. The Oracle models in the right figure still has an advantage for the really skewed distribution and λ∼0 , but loses to the Robust Gate for larger values of λ .
λ -test
Retrained, NLL ± std.dev.
Gate, NLL ± std.dev.
1/3, 1/3, 1/3
5.133 ± 0.010
4.994 ± 0.013
1/3, 1/2, 1/6
5.190 ± 0.080
5.068 ± 0.014
1/6, 1/3, 1/2
5.226 ± 0.011
5.099 ± 0.014
1/2, 1/3, 1/6
5.042 ± 0.011
4.890 ± 0.005
1/2, 1/6, 1/3
5.298 ± 0.006
5.117 ± 0.005
1/3, 1/6, 1/2
5.363 ± 0.054
5.187 ± 0.057
Appendix
Table 2: Robustness test results for different test distributions. The results are mean values ± 1 standard deviation, obtained over 5 runs with different initialization of the model training.
As scaling laws push the training of frontier large language models (LLMs) toward ever-growing data requirements, training pipelines are approaching a regime where much of the publicly available online text may be consumed. At the same time, widespread LLM usage increases the volume of machine-generated content on the web; together, these trends raise the likelihood of generated text re-entering future training corpora, increasing the associated risk of performance degradation often called model collapse. In practice, model developers address this concern through data cleaning, watermarking, synthetic-data policies, or, in some cases, blissful ignorance. However, the problem of model collapse in generative models has not been examined from a learning-theoretic perspective: we study it through the theoretical lens of the language generation in the limit framework, introducing a replay adversary that augments the example stream with the generator's own past outputs. Our main contribution is a fine-grained learning-theoretic characterization of when replay fundamentally limits generation: while replay is benign for the strongest notion of uniform generation, it provably creates separations for the weaker notions of non-uniform generation and generation in the limit. Interestingly, our positive results mirror heuristics widely used in practice, such as data cleaning, watermarking, and output filtering, while our separations show when these ideas can fail.
Large language models are trained on heterogeneous data mixtures, where different knowledge domains require both shared knowledge and specialization. Existing modular approaches typically impose explicit components or discover modules through interpretability analysis after training. In this work, we propose CORTEX, a learning dynamics-inspired framework that learns internal modularization within dense language models. CORTEX partitions trainable matrices into parameter groups and learns module assignments from domain-conditioned gradient and cross-domain gradient similarity. We introduce the selective lesion score and module-domain mutual information to characterize the target-domain lesion effects and alignment, and analyze how module assignment affects the trade-off between assignment bias and update magnitude. Experiments with 160M, Qwen3-8B, and Qwen3-32B backbone models show that CORTEX achieves the highest synthetic-domain exact match and largest average perplexity reduction, while remaining competitive on real-domain evaluations and forming identifiable modules.
Chuiyang Meng, Ming Tang, Vincent W. S. Wong
Simon Fraser University · Southern University of Science and Technology · The University of British Columbia
Extending a fully post-trained language model with new domain capabilities is fundamentally limited by monolithic training paradigms: retraining from scratch is expensive and scales poorly, while continued training often degrades existing capabilities. We present BAR (Branch-Adapt-Route), which trains independent domain experts, each through its own mid-training, supervised finetuning, and reinforcement learning pipeline, and composes them via a Mixture-of-Experts architecture with lightweight router training. Unlike retraining approaches that mix all domains and require full reprocessing for any update (with cost scaling quadratically), BAR enables updating individual experts independently with linear cost scaling and no degradation to existing domains. At the 7B scale, with experts for math, code, tool use, and safety, BAR achieves an overall score of 49.1 (averaged across 7 evaluation categories), matching or exceeding re-training baselines (47.8 without mid-training, 50.5 with). We further show that modular training provides a structural advantage: by isolating each domain, it avoids the catastrophic forgetting that occurs when late-stage RL degrades capabilities from earlier training stages, while significantly reducing the cost and complexity of updating or adding a domain. Together, these results suggest that decoupled, expert-based training is a scalable alternative to monolithic retraining for extending language models.
Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia +3
University of Washington · Allen Institute for AI · University of California, Berkeley