Disentangled Representation Learning via Modular Compositional Bias

AI summary

Overview

• Research area: Unsupervised disentangled representation learning (DRL) — factorizing data into interpretable components — combined with diffusion-based generative modeling.
• Technical level: Advanced. The paper assumes familiarity with variational autoencoders, total correlation objectives, slot attention, diffusion models, score-based likelihood bounds, and InfoNCE-style contrastive losses.
• Scope: The paper proposes a single framework in which factor-specific inductive biases are placed entirely in a "mixing strategy" that recombines latents, rather than in the loss function or the network architecture, and evaluates it on attribute, object, and joint disentanglement benchmarks.

What This Paper Is About

Disentangled representation learning methods are currently split into two camps: attribute methods that bake assumptions like statistical independence into the learning objective, and object-centric methods that bake assumptions like spatial exclusivity into the architecture (for example slot attention or additive decoders). Because these biases are entangled with the objective or the model, handling a new factor of variation — or several factors at once — requires redesigning losses or architectures.

The paper's goal is to move the factor-specific bias into a small, swappable module called a compositional bias. The authors derive mixing strategies for attributes, objects, and their combination, and show that changing only the mixing strategy lets one set of objectives and one architecture handle each case.

Key Contributions

1. A disentanglement framework that decouples factor-specific inductive biases from both learning objectives and architectures, so attribute and object disentanglement share one set of objectives and one architecture.
2. Derivation of explicit mixing strategies as compositional biases — $\pi_{attr}$ for attributes, $\pi_{obj}$ for objects, and $\pi_{joint}$ for both — grounded in the factorized support assumption of the latent distribution.
3. A revised compositional consistency loss, inspired by InfoNCE, that minimizes the relative (rather than absolute) distance between the composite latent and its re-encoded counterpart, with Sinkhorn-Knopp soft assignment to handle permuted object latents.
4. Empirical demonstration that a frozen, pre-trained unconditional diffusion model can serve as the likelihood prior for composite images, avoiding the out-of-distribution conditioning problem the authors attribute to L2C, which uses a representation-conditioned diffusion model.

Main Findings

• Attribute disentanglement (FactorVAE and DCI metrics): On Shapes3D the method scores 0.975 FactorVAE and 0.837 DCI; on MPI3D 0.708 and 0.458; on Cars3D 0.877 and 0.365. It achieves the best DCI on all three datasets and the best FactorVAE score on Shapes3D and MPI3D, but not on Cars3D (FactorVAE score there is 0.877 versus the highest reported baseline value of 0.976 for DisDiff-VQ).
• Margin over second-best: The paper reports about 8% higher FactorVAE score and 15.7–21.4% higher DCI than the second-best methods on Shapes3D and MPI3D.
• Object property prediction: On CLEVR-Easy the method reports Shape 95.81, Color 95.38, Position 50.72; on CLEVR Shape 87.04, Color 93.93, Material 94.81, Position 0.032; on CLEVR-Tex Shape 70.90, Material 52.2, Position 0.133. It outperforms LSD on CLEVR-Tex and performs comparably on CLEVR and CLEVR-Easy; it does better than L2C on CLEVR and CLEVR-Easy and is competitive on CLEVR-Tex.
• Unsupervised object segmentation: On CLEVR the method reports FG-ARI 91.20, mIoU 26.54, mBO 26.65; on CLEVR-Tex FG-ARI 87.68, mIoU 58.88, mBO 59.12, beating both LSD and L2C on CLEVR-Tex and achieving the best mIoU and mBO on CLEVR.
• Joint disentanglement (CLEVR-Style): The method reports 96.50% style prediction accuracy and GRAM loss of 5.050, plus object scores of Shape 83.56, Color 90.48, Material 93.74, Position 0.053. Baselines LSD and L2C report roughly 10% style accuracy (9.500 and 12.20 respectively), and DisDiff is reported as N/A for object metrics.
• Ablation on objectives: The paper states that all three losses ($\mathcal{L}{Diff}$, $\mathcal{L}{Prior}$, $\mathcal{L}_{Con}$) are essential; adding each term sequentially improves attribute disentanglement, while object disentanglement improves clearly only when all terms are used together. The full ablation table is truncated in the content shown.
• Latent structure: On Shapes3D the method identifies all six ground-truth factors. It also produces empty latents (depicted with $\phi$), which the authors say lets it handle a varying number of objects.
• No built-in spatial clustering: Because the framework removes slot attention, segmentation masks must be obtained by training a Spatial Broadcast Decoder on frozen representations; the paper reports that this captures whole objects while LSD and L2C often split objects across multiple latents.

Methodology in Plain English

The method starts from an autoencoder in which an encoder maps an image to a set of $K$ latent vectors and a diffusion decoder reconstructs or denoises the image. The central trick is what happens between encoding and decoding.

Two images are encoded, and a stochastic mixing operator $\pi(\cdot,\cdot)$ recombines their latents into a composite set. How that recombinations is chosen is the entire inductive bias:

• For attributes, each latent slot has its own distinct support, so each slot must come from exactly one of the two images and never both — this enforces mutual exclusiveness (a face gets one nose).
• For objects, all slots share a support, so any slot from either image can go into any position — this permits unrestricted exchange.
• For joint disentanglement, the $K$ latents are split into the first $M$ for attributes and the remaining $K-M$ for objects, and the corresponding strategy is applied to each part.

Two objectives make the composite meaningful. The prior loss pushes the composite image toward realism by using a frozen, pre-trained unconditional diffusion model; the authors approximate its gradient rather than backpropagating through the whole diffusion chain. The compositional consistency loss re-encodes the generated composite image and forces the result to match the composite latent that produced it. Because absolute distances collapse (all latents cluster together), the loss is formulated as a relative, InfoNCE-style objective against negative samples, with Sinkhorn-Knopp soft assignment used to match permuted object slots.

The total objective adds a denoising reconstruction loss to these two terms. The prior and consistency losses update only the encoder, while the decoder stays fixed, to prevent the encoder and decoder from cooperating to produce pretty composites from bad latents.

Why This Matters

• Research impact: The paper argues that embedding inductive bias in objectives or architectures does not transfer across factor types. If the mixing-strategy idea generalizes, it offers a common interface for factor types beyond attributes and objects, and it removes slot attention from object-centric learning — a departure from a long-standing architectural default.
• Real-world applications:
  • Controllable image editing, where a single attribute (style, color, material) or a single object can be swapped or removed without disturbing the rest of the scene.
  • Compositional content creation, where objects are inserted into or removed from scenes, including handling scenes with a varying number of objects.
  • Scene understanding and image segmentation, since the latent representations can support unsupervised object mask extraction with a learned decoder.
  • Data-efficient learning and interpretability, which the introduction names as motivations for factorizing data into reusable concepts.
• Industry relevance: Models that support fine-grained, factor-level control of generated images are directly relevant to generative media tools, dataset augmentation pipelines, and any product where users want to edit one property of an image without regenerating everything. A single framework that covers attributes, objects, and both is cheaper to maintain than separate, factor-specific systems.

Future Directions

• Scaling beyond the reported setting: the authors preview the MSN-Style dataset, an augmentation of MultiShapeNet (MSN) with over 10k unique, realistic furniture shapes, to test scalability and robustness on more complex data in the appendix.
• Extending the mixing-strategy formulation to factor types beyond attributes and objects, since the paper's claim is that any factor-specific bias can be expressed as a recombination rule derivable from the factor's support structure.
• Closing the segmentation gap: without slot attention the method has no built-in mechanism to group pixels, so it must rely on a separately trained Spatial Broadcast Decoder; improving this is a natural next step.
• Understanding why the method underperforms on specific metrics — such as the FactorVAE score on Cars3D and the position metric on CLEVR-Tex — and whether these are consequences of the mixing strategy or of the decoders used.
• Removing the remaining dependence on a pre-trained unconditional diffusion model trained from scratch per dataset for the likelihood prior.

Target Audience

This paper is most useful to machine learning researchers working on disentangled and object-centric representation learning, diffusion-based generative models, and compositional image generation. It is also relevant to practitioners building controllable image editing or generative media systems who need factor-level manipulation, and to graduate students who already understand the basics of VAEs, slot attention, and diffusion models but want a clear statement of how inductive bias can be separated from architecture and loss.

Authors’ abstract

Recent disentangled representation learning (DRL) methods heavily rely on factor specific strategies-either learning objectives for attributes or model architectures for objects-to embed inductive biases. Such divergent approaches result in significant overhead when novel factors of variation do not align with prior assumptions, such as statistical independence or spatial exclusivity, or when multiple factors coexist, as practitioners must redesign architectures or objectives. To address this, we propose a compositional bias, a modular inductive bias decoupled from both objectives and architectures. Our key insight is that different factors obey distinct recombination rules in the data distribution: global attributes are mutually exclusive, e.g., a face has one nose, while objects share a common support (any subset of objects can co-exist). We therefore randomly remix latents according to factor-specific rules, i.e., a mixing strategy, and force the encoder to discover whichever factor structure the mixing strategy reflects through two complementary objectives: (i) a prior loss that ensures every remix decodes into a realistic image, and (ii) the compositional consistency loss introduced by Wiedemer et al. (arXiv:2310.05327), which aligns each composite image with its corresponding composite latent. Under this general framework, simply adjusting the mixing strategy enables disentanglement of attributes, objects, and even both, without modifying the objectives or architectures. Extensive experiments demonstrate that our method shows competitive performance in both attribute and object disentanglement, and uniquely achieves joint disentanglement of global style and objects. Code is available at https://github.com/whieya/Compositional-DRL.