Principled Out-of-Distribution Generalization via Simplicity

TL;DR

The paper investigates the theoretical underpinnings of out-of-distribution generalization in modern foundation models through a simplicity principle. By formalizing a simplicity metric $R(\beta)$ and the ground-truth $\beta^{\star}$ as the simplest training-minimizer, it analyzes a regularized maximum likelihood estimator under covariate shift in two regimes: a constant-gap regime with a fixed simplicity gap $\Delta$ and a vanishing-gap regime with a smooth proximity condition. The authors derive sharp non-asymptotic excess-risk bounds, achieving a fast $\tilde{O}(1/n)$ rate in the constant-gap setting and a tunable rate $\tilde{O}(n^{-1+2/(3\tau)})$ in the vanishing-gap regime, where $\tau$ governs the softness of the gap. Across theoretical development and illustrative experiments on diffusion models and a simplified MLP identity task, the work argues that the simplest model among source-minimizers generalizes best to the target distribution, offering a principled explanation for robust OOD behavior and guiding regularization strategies in practice.

Abstract

Modern foundation models exhibit remarkable out-of-distribution (OOD) generalization, solving tasks far beyond the support of their training data. However, the theoretical principles underpinning this phenomenon remain elusive. This paper investigates this problem by examining the compositional generalization abilities of diffusion models in image generation. Our analysis reveals that while neural network architectures are expressive enough to represent a wide range of models -- including many with undesirable behavior on OOD inputs -- the true, generalizable model that aligns with human expectations typically corresponds to the simplest among those consistent with the training data. Motivated by this observation, we develop a theoretical framework for OOD generalization via simplicity, quantified using a predefined simplicity metric. We analyze two key regimes: (1) the constant-gap setting, where the true model is strictly simpler than all spurious alternatives by a fixed gap, and (2) the vanishing-gap setting, where the fixed gap is replaced by a smoothness condition ensuring that models close in simplicity to the true model yield similar predictions. For both regimes, we study the regularized maximum likelihood estimator and establish the first sharp sample complexity guarantees for learning the true, generalizable, simple model.

Figures (7)

• Figure 1: Diffusion Model Image Generation Setting.
• Figure 2: generalizable vs. non-generalizable model weights. (a) Sum of squared Frobenius norms of weights in models trained on uniform mappings. (b) Sum of squared Frobenius norms of weights in models trained on permutation mappings. (c) Sum of squared Frobenius norms for models trained on interpolations between the identity and flipped maps. Here, $\alpha = 0$ corresponds to the identity map and $\alpha = 1$ to the flipped map. In all three plots, the model trained solely on $S$ using the identity map is shown in orange.
• Figure 3: Diffusion Training. (a) Training loss (MSE) per epoch, averaged over all 200,000 training examples. (b) Test loss (MSE) per epoch for each test class, averaged over all 2,000 test examples per class. The final test loss after all 400 epochs is plotted in dashed lines, with numerical values 2.00398e-4, 1.46671e-4, 1.67488e-4, and 1.32801e-4.
• Figure 4: Identity Mapping Training: (a) Training loss (MSE) per epoch averaged over all ten runs, plotted for the first $20,000$ of $40,000$ total epochs; final training loss: 6.02663e-4. (b) Test loss (MSE) per epoch for each test class, averaged over all ten models. We restrict the y-axis (loss) of the plot to make the differences between the test losses for different classes visible. The final test losses after all $20,000$ epochs are 3.60000e-3, 9.82313e-3, 5.24610e-3, 1.55543e-2.
• Figure 5: Uniform Mapping Training: (a) Histogram of training loss (MSE) in final epoch for all eighty trials. Training losses are averaged over ten independent models trained per trial. (b) An example of the training loss curve for a randomly selected trial averaged over all 10 runs.

Theorems & Definitions (37)

• Theorem 4.1
• Theorem 4.2
• Proposition C.1
• proof : Proof of Proposition \ref{['prop1']}
• Lemma C.2
• proof : Proof of Lemma \ref{['lem1']}
• Lemma C.3
• proof : Proof of Lemma \ref{['lem2']}
• Proposition C.4
• proof : Proof of Proposition \ref{['prop2']}