A Theoretical Framework for Modular Learning of Robust Generative Models

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.

Training a large-scale learning algorithm from scratch requires significant computing power. At the same time, when a single model tries to learn from very different domains—like computer code and natural language—the conflicting signals often cause internal interference, which reduces its accuracy on the different tasks. We addressed this problem with a modular approach: instead of training one giant network, we take smaller, pre-trained “domain expert" models and combine them using an adaptive gating mechanism. For any given input, this gate calculates the best way to mix the experts' outputs. We mathematically prove that this setup remains robust for any unpredictable mixture of the domains. By keeping the experts isolated, our method prevents distinct subjects from interfering with one another, allowing it to outperform standard unified models. We also provide a technique to streamline this gating process so the system generates text efficiently in practice. This framework makes training and updating learning algorithms much more practical. When new data emerges, researchers can simply plug in a new expert rather than retraining the entire system.