What Drives Compositional Generalization in Visual Generative Models? The Importance of Continuous Training Objectives

Compositional generalization, which is the ability for visual generative models to synthesize novel combinations of known concepts, remains inconsistent across different architectures. This work systematically investigates which design choices critically determine compositional success by isolating independent design axes in visual generative models. The key finding is that continuous training objectives are necessary for robust compositional generalization while discrete categorical objectives inhibit it, and that complete conditioning on the generating factors during training is critical for composition.

How it works

The research develops a systematic framework to isolate causal effects by categorizing modern visual generative models along three primary axes: (i) the Tokenizer, (ii) the Generative model, and (iii) the Conditioning signal. The study anchors itself by selecting canonical representatives—MaskGIT for discrete/masked modeling and DiT for continuous/diffusion modeling—and systematically interpolating between them by varying individual design axes while holding others fixed to isolate critical drivers.

The core research questions guiding this investigation include:

(i) Does the type of the tokenizer affect compositional generalization?

(ii) How does the generative model design affect compositionality? In particular, does it matter whether the modeled distribution is continuous or discrete (e.g., continuous latents vs. discrete tokens)? And is a denoising-based objective essential, or does a masking-based loss suffice?

(iii) Can generative models, without being conditioned on the full data-generating factors during training, generalize compositionally? Must conditioning have the exact factors, or can they rely on quantized or missing abstractions (e.g., “red” instead of the RGB value, or hiding “smile” in a set of other concepts)?

(iv) Guided by our findings, can we intervene on non-compositional models to endow them with compositional capabilities?

Key Findings and Drivers

The experiments consistently indicate that models trained to learn a continuous distribution by their training objective exhibit stronger compositional capabilities than models trained to model a categorical distribution. Furthermore, providing full conditioning information of the generating factors during training is critical; quantized or partial conditioning leads to weaker compositional generalization. Conversely, the study finds that training the tokenizer with a quantization bottleneck has no significant effect on the downstream compositional generalization.

Interventions and Enhancements

To address models trained with discrete objectives, researchers tested augmenting them with an auxiliary continuous objective. This involved extending MaskGIT’s training objective with a Joint Embedding Predictive Architecture (JEPA) objective, which learns to reconstruct target patch representations by using context patches from the same image in an abstract embedding space. This modification introduces continuous latent targets and yields clear improvements on held-out factor combinations, suggesting that predictive continuous objectives can shape the internal structure to support compositionality in discrete models. Mechanistic analysis further revealed that this objective induces more disentangled and semantically structured representations, and enables stronger compositional generalization, by reducing polysemanticity in attention heads.

Evaluation Across Modalities

The findings were validated across diverse datasets, including synthetic (Shapes2D, Shapes3D), real-world faces (CelebA), synthetic videos (CLEVRER-Kubric), and world models (CoVLA). In visual tasks like Shapes2D and CLEVRER-Kubric, models trained to learn continuous distributions—such as DiT—consistently show better level-2 compositions than MaskGIT. For real-world video datasets like CoVLA, the Orbis model trained with the continuous objective (Orbis-DiT) substantially outperforms MaskGIT on novel compositions.

Conclusion and Future Directions

The study concludes that the primary drivers are: (i) continuous training objectives are necessary for robust compositional generalization while discrete categorical objectives inhibit it; (ii) complete conditioning is essential, as quantized or incomplete signals lead to unstable or failed compositions. The introduction of a JEPA-based auxiliary loss successfully improves compositionality for discrete models. Mechanistic analysis suggests that this process promotes learning factor-specific circuits, providing actionable insights for designing compositional generative models that better preserve semantic structure and enable novel compositions beyond the training distribution. The results also suggest that the advantages of continuous reasoning mechanisms observed in visual generation may carry over to the language domain.

The gist

Continuous training objectives are necessary for robust compositional generalization while discrete categorical objectives inhibit it, and complete conditioning on the generating factors during training is critical for composition.

(Note: The summary above adheres strictly to the constraints provided, focusing only on extracted information and following the specified structure.)

How it works

The research develops a systematic framework to isolate causal effects by categorizing modern visual generative models along three primary axes: (i) the Tokenizer, (ii) the Generative model, and (iii) the Conditioning signal.

Based on the provided scientific paper, here are specific improvements for AI systems and what those improved systems can achieve:

1. Improve compositional generalization in visual generative models by shifting from discrete training objectives to continuous training objectives. This involves augmenting discrete models (like MaskGIT) with a continuous auxiliary objective, specifically a Joint Embedding Predictive Architecture (JEPA) loss, which operates on continuous latent representations rather than discrete tokens.

2. Enable visual generative models to robustly generate novel combinations of known concepts (e.g., novel gender/hair color/smile combinations in CelebA or novel object shapes/colors in Shapes3D).

3. Enhance the disentanglement of visual factors by using JEPA-based training objectives, which are shown to induce more disentangled and semantically structured representations compared to standard masking-based losses.

4. Improve the fidelity of compositional retrieval for unseen compositions in video generation tasks (e.g., CoVLA) by employing continuous generative models (like DiT) or JEPA-enhanced discrete models, leading to higher Compositional Retrieval Accuracy (CRA).

5. Mitigate representational entanglement within visual models by using JEPA-based objectives, which are shown to reduce the polysemanticity of attention heads and decrease the overlap of neurons across different circuits for different concepts. This leads to more factor-specific representations.

6. Develop multimodal AI systems capable of reasoning over complex compositional constraints, leveraging continuous reasoning mechanisms (like Continuous Chain-of-Thought) in Large Language Models, suggesting that the advantages of continuous objectives transfer to language modeling tasks like card arithmetic reasoning.

In summary, these improvements enable AI systems to move beyond simply interpolating training data toward a deeper understanding and synthesis of underlying causal factors, leading to:

• Robust generation of novel visual concepts.

• Better alignment between generated outputs and real-world or target distributions.

• More interpretable and disentangled internal representations that map clearly to semantic factors.