A Theoretical Framework for Modular Learning of Robust Generative Models

Corinna Cortes; Mehryar Mohri; Yutao Zhong

cs.LGFeb 19, 2026

A Theoretical Framework for Modular Learning of Robust Generative Models

Authors: Corinna Cortes, Mehryar Mohri, Yutao Zhong

Organizations: Google Research, New York, NY · Courant Institute of Mathematical Sciences, New York, NY

Abstract

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 $\mathcal{G}_{1}$ 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

Panel of Figure 1 — 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.

Panel of Figure 3 — 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 (

\lambda\approx 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 $\pm$ 0.005	5.118 $\pm$ 0.011
Code Expert	4.722 $\pm$ 0.045	5.267 $\pm$ 0.788
FineWeb Expert	5.623 $\pm$ 0.004	5.623 $\pm$ 0.006
Retrained Model	5.133 $\pm$ 0.010	5.306 $\pm$ 0.257
Standard MoE (Scratch)	5.080 $\pm$ 0.009	5.245 $\pm$ 0.269
Gate Model (Ours)	4.994 $\pm$ 0.013	5.087 $\pm$ 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

\widehat{\pi}_{1}

and

\widehat{\pi}_{2}

. The target distribution (red dashed) exceeds this feasible capacity. The optimal gate

\pi_{g^{*}}

(solid blue) traces the target where possible but is “clipped” to the expert boundaries when the target falls outside the hull.

Panel of Figure 7 — Figure 7: Geometry of Prior Knowledge (Section C.1 ). The outer triangle represents the full probability simplex $\Delta$ . The green region $\Lambda$ represents the subset of valid mixtures defined by prior knowledge. The algorithm projects the adversary’s updates (red point) back onto $\Lambda$ (blue point), tightening the worst-case bound as per Theorem C.2 .

Panel of Figure 8 — Figure 8: Dynamics of the Primal-Dual Game (Algorithm 3 ). The optimization is modeled as a 3-player game. The $\lambda$ -player maximizes the mixture difficulty using Exponentiated Gradient. The $g$ -player minimizes the robust loss. The $\mu$ -player enforces the global normalization constraint ( $Z_{g}=1$ ) via Dual Ascent.

A(x)=\frac{\pi_{g^{*}}(x)}{M\,q(x)}=\frac{\pi_{g^{*}}(x)}{p\,q(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

\lambda

where the test distribution is predominantly made up by Domain B. As

\lambda

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

\lambda\sim 0

, but loses to the Robust Gate for larger values of

\lambda

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

\lambda

where the test distribution is predominantly made up by Domain B. As

\lambda

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

\lambda\sim 0

, but loses to the Robust Gate for larger values of

\lambda

$\lambda$ -test	Retrained, NLL $\pm$ std.dev.	Gate, NLL $\pm$ std.dev.
1/3, 1/3, 1/3	5.133 $\pm$ 0.010	4.994 $\pm$ 0.013
1/3, 1/2, 1/6	5.190 $\pm$ 0.080	5.068 $\pm$ 0.014
1/6, 1/3, 1/2	5.226 $\pm$ 0.011	5.099 $\pm$ 0.014
1/2, 1/3, 1/6	5.042 $\pm$ 0.011	4.890 $\pm$ 0.005
1/2, 1/6, 1/3	5.298 $\pm$ 0.006	5.117 $\pm$ 0.005
1/3, 1/6, 1/2	5.363 $\pm$ 0.054	5.187 $\pm$ 0.057

Appendix

Table 2: Robustness test results for different test distributions. The results are mean values

\pm

1 standard deviation, obtained over 5 runs with different initialization of the model training.

Explore similar work

Language Generation with Replay: A Learning-Theoretic View of Model Collapse
Mar 12, 2026Giorgio Racca, Michal Valko, Amartya SanyalModel Collapse Data Contamination in Language Models
University of Copenhagen · Isara Labs
CORTEX: Learning to Share and Specialize in Dense Language Models
Sep 28, 2026Chuiyang Meng, Ming Tang, Vincent W. S. WongDomain Adaptation for LLMs Language Modeling
Simon Fraser University · Southern University of Science and Technology · The University of British Columbia
Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts
Apr 20, 2026Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia +3Continual Learning for LLMs Mixture-of-Experts Language Models
University of Washington · Allen Institute for AI · University of California, Berkeley

Mar 12, 2026J/K move · Enter open · S save

Language Generation with Replay: A Learning-Theoretic View of Model Collapse

Giorgio Racca, Michal Valko, Amartya Sanyal

University of Copenhagen · Isara Labs

Model Collapse Data Contamination in Language Models LLM Training Text Generation

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Language Generation with Replay: A Learning-Theoretic View of Model Collapse

CORTEX: Learning to Share and Specialize in Dense Language Models

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts