cs.CVOct 3, 2025
SaveWhat Drives Compositional Generalization in Visual Generative Models? The Importance of Continuous Training Objectives
Organizations: University of Freiburg · Bosch Center for Artificial Intelligence · Inria, École Normale Supérieure, CNRS, PSL Research University
Abstract
Compositional generalization, the ability to generate novel combinations of known concepts, is a key ingredient for visual generative models. Yet, not all mechanisms that enable or inhibit it are fully understood. In this work, we conduct a systematic study of which design choices critically determine compositional generalization in image and video generation. By isolating independent design axes, we identify two key factors strongly associated with compositional success: (i) whether the training objective operates on a discrete or continuous distribution, and (ii) the completeness of conditioning information about constituent factors during training. We also show that relaxing the discrete loss with an auxiliary continuous latent objective can partially recover compositional performance in discrete models like MaskGIT. Our findings, corroborated by diverse compositional tasks and preliminary evidence in world models and LLMs, motivate a shift toward continuous objectives for compositional generalization.
Figures & tables
Figure 1 : Compositional Generalization Analysis. We compare discrete masking (MaskGIT) and continuous diffusion (DiT) on novel compositions of three binary CelebA factors: gender, hair color, and smile. Models are trained on four observed combinations ( blue ) and evaluated on held-out Level-1 ( pink , one-factor change) and Level-2 ( red , two-factor change) compositions. DiT (row 1) generalizes compositionally, while MaskGIT (row 2) struggles.
| Model | Masking-based | Training loss | Latent Distribution |
|---|---|---|---|
| MaskGIT | ✓ | Categorical NLL | discrete |
| GIVT | ✓ | GMM-NLL | continuous |
| MAR | ✓ | Diffusion | continuous |
| DiT | ✗ | Diffusion | continuous |
Table 1: Summary of generative models.
Figure 2 : DiT exhibits compositional generalization regardless of the type of tokenizer used. While the training dynamics differ, DiT shows compositional generalization at the end of training. The blue , pink , and red curves show linear probe accuracies for the training data, level-1 compositions, or level-2 compositions, respectively. Results are consistent for MAR ( Figure 7 ) and video ( Figure 8 ).
Figure 3 : Compositional generalization performance on Shapes2D across different model architectures. Models that learn continuous distributions (DiT, MAR, and GIVT) consistently show better level-2 compositions than MaskGIT, with the decisive shift in performance occurring at the categorical-to-continuous intervention. The blue , pink , and red curves denote training, level-1, and level-2 compositions, respectively. Consistent results are observed for CLEVRER-Kubric ( Figure 9 ).
Figure 4 : Comparison of conditioning information levels and their impact on compositional generalization in DiT on Shapes2D. (a) Continuous (full-information) conditioning leads to uniform convergence across all compositions. (b) Label dropout conditioning leads to inconsistent generalization; several unseen compositions fail. (c) Discrete (quantized) conditioning leads to partial generalization, with some failing samples. (d) Discrete (quantized) conditioning with dropout, the most severe loss of information, leads to the most failure. Shaded areas indicate standard deviation across three different seeds. We provide additional results in Section C.5 . The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2.
Figure 5 : An overview of MaskGIT combined with the JEPA-based training objective. We apply the JEPA loss at specific layers ( ) on an intermediate masked token representation in the transformer and train a lightweight predictor to reconstruct target states using MSE as an error metric and a stop-gradient signal to avoid representation collapse.
Figure 6 : Comparison of linear probe accuracy for MaskGIT with ( 6(a) ) Standard and ( 6(b) ) JEPA-based training objectives for Shapes2D. The JEPA-based objective clearly enhances compositional abilities even though it cannot fully compensate the problems introduced by the discrete space. We also see lower polysemanticity between attention heads ( 6(c) ) and a decreasing Jaccard trend over shared circuits ( 6(d) ). The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions. Consistent videos results shown in Figure 13 .
Appendix figures & tables32 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7 : MAR exhibits generalization on both level-1 and level-2 compositions in Shapes2D regardless of whether (a) discrete, (b) continuous, or (c) no tokenizer is used. The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions.
Figure 8 : DiT exhibits generalization on both level-1 and level-2 compositions in CLEVRER-Kubric regardless of whether (a) discrete, (b) continuous, or (c) no tokenizer is used. The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions.
Figure 9 : Comparison of linear probe accuracy for CLEVRER-Kubric across ( 9(a) ) MaskGIT, ( 9(b) ) DiT, and ( 9(c) ) MAR. Models leveraging a continuous latent space (DiT, and MAR) show better level-2 compositions than MaskGIT. The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions.
Figure 10 : Compositional generalization across different training conditions for MaskGIT on Shapes2D. We see that MaskGIT struggles to compositionally generalize despite varying multiple different training parameters.
| Model | Seen Mean | Level-1 Novel Mean | Level-2 Novel |
|---|---|---|---|
| (011, 110, 101) | (111) | ||
| AR (discrete) | 86.7 | 38.5 | 6.2 |
| AR-GMM (continuous) | 90.2 | 64.1 | 46.9 |
Table 2 : Probe accuracy (%) of causal autoregressive models with a discrete (categorical) and a continuous (Gaussian mixture) objective on Shapes2D.
| Unmasking | 000 | 001 | 010 | 100 | Seen Mean | 011 | 101 | 110 | 111 | Novel Mean |
|---|---|---|---|---|---|---|---|---|---|---|
| Confidence | 87.8 | 94.3 | 82.9 | 77.3 | 85.6 | 64.9 | 54.0 | 38.4 | 0.7 | 39.5 |
| Random order | 87.8 | 95.0 | 86.1 | 74.3 | 85.8 | 45.7 | 41.2 | 26.3 | 2.0 | 28.8 |
Table 3 : Probe accuracy (%) of MaskGIT on Shapes2D with confidence-based vs. random unmasking order (256 samples per composition).
Figure 11 : Comparison of linear probe accuracy for DiT trained on CLEVRER-Kubric with ( 11(a) ) Explicit conditioning and ( 11(b) ) Implicit conditioning mechanisms. Training with explicit label information allows for better level-2 compositions. The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions.
Figure 12 : Qualitative Results of the concept dropout with DIT on Shapes2D. The visualization demonstrates that when contextual information about all concepts is incomplete, the model struggles to generalize compositionally. A typical failure mode is the substitution of novel concept combinations with the closest seen ones—for instance, the 110 configuration, which should correspond to a novel blue big triangle, is instead rendered as a red big triangle, likely due to its proximity in the training distribution. Similar behavior is happening to the 101, which should correspond to a novel red small triangle, sometimes result in a red small circle or red big triangle.
| Model | Conditioning | Seen | Novel L1 | Novel L2 | All Novel |
|---|---|---|---|---|---|
| DiT | Factors | 97.3 | 99.0 | 96.9 | 98.4 |
| Language | 74.2 | 71.4 | 48.4 | 65.6 | |
| Language (partial descriptions) | 76.6 | 62.5 | 1.6 | 47.3 | |
| MaskGIT | Factors | 85.6 | 52.4 | 0.7 | 39.5 |
| Language | 84.1 | 52.0 | 1.6 | 39.4 | |
| Language (partial descriptions) | 85.2 | 19.9 | 0.3 | 15.0 |
Table 4 : Probe accuracy (%) on Shapes2D with explicit factor conditioning vs. CLIP text conditioning. L1/L2 denote Level-1/Level-2 novel compositions; All Novel averages over all unseen compositions.
Figure 13 : Comparison of linear probe accuracy for MaskGIT with Standard ( 13(a) ) and JEPA-based ( 13(b) ) training objectives for CLEVRER-Kubric. The JEPA-based training objective clearly enhances compositional abilities, and lowers polysemanticity ( 13(c) ) even though it cannot fully compensate the problems introduced by the discrete space. The blue curves show performance on training data, pink curves depict level-1 compositions, and red curve denotes level-2 compositions.
| Layers | Accuracy |
|---|---|
| {6} | 27.61 |
| {8} | 26.35 |
| {6,8} | 38.16 |
| {7,9} | 36.62 |
| {6,8,10} | 53.52 |
| {7,9,11} | 56.27 |
Table 5 : Maximum linear probe accuracy after applying the JEPA loss on different layers and combinations of layers for a MaskGIT model trained on CLEVRER-Kubric
Figure 14 : Effect of weighting factor on probe accuracy
Figure 15 : Visualization of MaskGIT circuits for a large red cube and small blue sphere trained with the standard training objective versus the JEPA-based training objective.
| Color-ablated | Shape-ablated | |
| MaskGIT+JEPA | ||
| Color-varying | pp | 6pp |
| Shape-varying | 5pp | pp |
| Standard MaskGIT | ||
| Color-varying | 18pp | 14pp |
| Shape-varying | 15pp | 17pp |
Table 6 : Level-2 probe accuracy drops after ablating factor-specific circuits. Bold entries mark factor-aligned ablations. MaskGIT+JEPA shows selective degradation; standard MaskGIT shows uniform drops, consistent with entangled representations.
| Corruption | Strength | Shape | Color | Size | Joint |
|---|---|---|---|---|---|
| Blur ( ) | 1 | 98.9 | 100.0 | 89.5 | 88.4 |
| 2 | 99.2 | 100.0 | 91.7 | 90.9 | |
| 4 | 73.1 | 100.0 | 89.3 | 64.7 | |
| 8 | 0.0 | 100.0 | 50.0 | 0.0 | |
| Noise ( ) | 0.1 | 97.4 | 100.0 | 93.9 | 91.2 |
| 0.25 | 91.1 | 100.0 | 78.2 | 75.5 |
Table 7 : Probe accuracy (%) on real Shapes2D test images under increasing corruptions. Joint denotes all three factors being classified correctly.
| Model | (0,0,0) | (0,0,1) | (0,1,0) | (1,0,0) | (1,0,1) | (1,1,0) | (0,1,1) | (1,1,1) |
|---|---|---|---|---|---|---|---|---|
| DiT | 100% | 100% | 100% | 100% | 100% | 100% | 100% | 100% |
| MaskGIT | 94.7% | 97% | 94.7% | 94% | 79.49% | 62.6% | 58.85% | 30% |
Table 8 : Performance comparison of DiT and MG across different compositions.
| Model | 1-NN | 5-NN |
|---|---|---|
| DiT | 27 | 36 |
| MaskGIT | 21 | 16 |
| MaskGIT + JEPA | 31 | 27 |
Table 9 : Compositional Retrieval Accuracy (CRA, %) 100 generated samples.
| Method | Unseen FID | Unseen FDD | Seen FID | Seen FDD |
|---|---|---|---|---|
| DiT | 118.60 | 993.92 | 109.96 | 681.78 |
| MG | 162.38 | 1470.78 | 113.94 | 719.95 |
| MG+JEPA | 120.76 | 980.92 | 117.35 | 621.43 |
Table 10 : FID and FDD scores on seen and unseen splits on CelebA.
Figure 16 : Qualitative Results on Shapes2D
| Factor | Possible Values | Super Group 0 | Super Group 1 |
|---|---|---|---|
| Color | red, orange, yellow, lime green, green, cyan, blue, indigo, purple, pink | red, orange, yellow, lime green, green, cyan, blue | indigo, purple, pink |
| Shape | cube, cylinder, sphere | cube, cylinder | sphere |
| Size | 0–7 (ordinal) | 0, 1, 2, 3, 4, 5 | 6, 7 |
Table 11 : Factorization of the data into Super Groups . This partitioning is used to evaluate compositional generalization across disjoint super groups.
Figure 17 : Shapes3D results of DiT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 18 : Shapes3D results of DiT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 19 : Shapes3D results of DiT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 20 : Shapes3D results of MaskGIT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 21 : Shapes3D results of MaskGIT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 22 : Shapes3D results of MaskGIT. Dark blue : all factors from Supergroup 0; light blue : one factor from Supergroup 1; pale orange : level-1 compositions; and dark red : level-2 novel compositions.
Figure 23 : Qualitative Results on CelebA
| Real train | Real novel | ||||||
|---|---|---|---|---|---|---|---|
| Gen. novel | Model | | | | | | |
| | DiT | 0.470 | |||||
| MaskGIT | 0.180 | ||||||
| | DiT | 0.430 | |||||
| MaskGIT | 0.140 | ||||||
Table 12 : Nearest-neighbor compositional retrieval metric (CRA) comparing generated videos to real videos. Factors are time of day: Day , Night ; directions: Left , Straight , Right . Rows show a generated novel split evaluated against real train and real novel classes. The Hit@1 (correct target novel class) is bolded; black bold indicates better results.
| Method | FVD (Seen) | FVD (Unseen) |
|---|---|---|
| DiT | 730.30 | 2257.51 |
| MG | 2320.95 | 3197.74 |
Table 13 : FVD on seen and unseen splits on CoVLA.
Figure 24 : Qualitative results with Orbis on CoVLA. Two novel compositions are shown, and , each conditioned on a 7-frame initial context and predicting 5 frames ahead. Orbis-DiT follows the target compositions, whereas Orbis-MaskGIT often struggles.
Figure 25 : Example system prompt, unsuccessful model output, and verifier feedback for our language task with compositional rules: target restriction and red-suit value doubling.
Figure 26 : Example system prompt, successful model output, and verifier feedback for a single-rule training instance of our language task: target restriction without red-suit value doubling.