# A Theoretical Framework for Modular Learning of Robust Generative Models

arXiv: https://arxiv.org/abs/2602.17554 · HTML: https://mlsanity.com/paper/2602.17554/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 82 references that are in the map.
- First posted on arXiv 19 Feb 2026.

## Builds on (most cited first, 30 of 46)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Continual Lifelong Learning with Neural Networks: A Review](https://mlsanity.com/paper/1802.07569/) (2018)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [Green AI](https://mlsanity.com/paper/1907.10597/) (2019)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Cross-Entropy Loss Functions: Theoretical Analysis and Applications](https://mlsanity.com/paper/2304.07288/) (2023)
- [A vector-contraction inequality for Rademacher complexities](https://mlsanity.com/paper/1605.00251/) (2016)
- [Modular Deep Learning](https://mlsanity.com/paper/2302.11529/) (2023)
- [Algorithms and Theory for Multiple-Source Adaptation](https://mlsanity.com/paper/1805.08727/) (2018)
- [Principled Approaches for Learning to Defer with Multiple Experts](https://mlsanity.com/paper/2310.14774/) (2023)
- [Predictor-Rejector Multi-Class Abstention: Theoretical Analysis and Algorithms](https://mlsanity.com/paper/2310.14772/) (2023)
- [LLM Augmented LLMs: Expanding Capabilities through Composition](https://mlsanity.com/paper/2401.02412/) (2024)
- [Learning to Reject with a Fixed Predictor: Application to Decontextualization](https://mlsanity.com/paper/2301.09044/) (2023)
- [Theoretically Grounded Loss Functions and Algorithms for Score-Based Multi-Class Abstention](https://mlsanity.com/paper/2310.14770/) (2023)
- [Regression with Multi-Expert Deferral](https://mlsanity.com/paper/2403.19494/) (2024)
- [Ranking with Abstention](https://mlsanity.com/paper/2307.02035/) (2023)
- [Realizable H-Consistent and Bayes-Consistent Loss Functions for Learning to Defer](https://mlsanity.com/paper/2407.13732/) (2024)
- [A Universal Growth Rate for Learning with Smooth Surrogate Losses](https://mlsanity.com/paper/2405.05968/) (2024)

## Built on by (most cited first, 4 of 4)

- [Optimized Deferral for Imbalanced Settings](https://mlsanity.com/paper/2604.27723/) (2026)
- [Linear-Core Surrogates: Smooth Loss Functions with Linear Rates for Classification and Structured Prediction](https://mlsanity.com/paper/2604.27742/) (2026)
- [Online Learning-to-Defer with Varying Experts](https://mlsanity.com/paper/2605.12340/) (2026)
- [Principled Algorithms for Optimizing Generalized Metrics in Multi-Label Learning](https://mlsanity.com/paper/2605.28767/) (2026)

## Abstract (arXiv)

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.