Technical Summary: Disentangled Representation Learning via Flow Matching

Problem Statement

Disentangled representation learning aims to encode high-dimensional observations into compact latent representations where underlying semantic factors (e.g., object shape, color, size) are explicitly separated. While recent generative modeling advances, particularly diffusion models, have offered new paradigms for this task, existing methods face significant limitations:

• VAE-based methods (e.g., $\beta$-VAE, FactorVAE) often struggle with a trade-off between reconstruction fidelity and disentanglement strength.
• GAN-based methods (e.g., InfoGAN) lack reversible inference mechanisms between data and latent spaces, limiting flexibility.
• Diffusion-based methods typically rely on inductive biases to encourage statistical independence among latent dimensions. However, they frequently lack strong semantic alignment. Latent dimensions may be statistically decorrelated yet still mix semantic factors, undermining interpretability and fine-grained controllability. Furthermore, the stochastic nature of diffusion trajectories and the complexity of iterative denoising can lead to high computational costs and limited geometric structure for precise alignment.

Methodology

The authors propose a framework based on Flow Matching that casts disentanglement as learning factor-conditioned flows in a compact latent space. The approach leverages the deterministic nature of Ordinary Differential Equations (ODEs) to provide a principled geometric basis for semantic alignment.

1. Framework Overview

• Latent Space Construction: Input images are encoded into a compact, semantically meaningful latent space using a pre-trained VQ-GAN encoder.
• Factor Extraction: A trainable factor encoder extracts $N$ factor embeddings ($S_\gamma(I)$), representing distinct, controllable semantic attributes.
• Conditional Flow Matching: The model learns a time-dependent vector field $v_\theta$ that transports probability mass from a source prior (standard Gaussian) to the target data distribution in the latent space. This transport is conditioned on the extracted factors via cross-attention mechanisms within a U-Net backbone.

2. Factorized Velocity and Orthogonality Regularization

A core innovation is the explicit decomposition of the learned vector field to enforce non-overlapping, factor-specific dynamics.

• Factorized Velocity: The aggregate velocity field is decomposed into $N$ factor-specific components: $v_\theta = \sum v^{(i)}_\theta$.
• Output Attention Routing: To implement this, the authors introduce an output-attention routing mask. For each spatial location, the predicted velocity is distributed across factors using a softmax-weighted attention mechanism. This ensures that the sum of factor-specific components equals the aggregate prediction (mass conservation).
• Orthogonality Regularizer ($L_{orth}$): To prevent cross-factor interference and information leakage, the authors introduce a regularizer that penalizes the cosine similarity between different factor-specific velocity components. This encourages distinct semantic factors to capture unique transport directions.
• Objective Function: The total loss combines the standard flow matching regression loss ($L_{FM}$) with the orthogonality regularizer:
  $$L(\theta, \gamma) = L_{FM}(\theta) + \lambda_{orth} L_{orth}$$

Key Contributions

1. Flow-Matching Formulation: The paper casts disentangled representation learning as learning a factor-conditioned flow in a compact latent space, enabling a deterministic and effective generative process via ODE-based sampling.
2. Explicit Semantic Alignment: The authors propose an alignment module that decomposes the learned vector field into factor-specific components. This provides a direct mapping from factors to latent-space dynamics, enabling fine-grained, factor-level control.
3. Orthogonality Regularization: A novel non-overlap (orthogonality) regularizer is introduced to suppress cross-factor interference, ensuring that different latent dimensions capture distinct semantic factors rather than redundant information.
4. Empirical Validation: Extensive experiments demonstrate consistent improvements over representative baselines across multiple datasets, yielding higher disentanglement scores, improved controllability, and better sample fidelity.

Experimental Results

The method was evaluated on synthetic datasets (Cars3D, Shapes3D, MPI3D-toy) and a real-world dataset (CelebA).

• Quantitative Performance:
  • On MPI3D-toy, the proposed method achieved a FactorVAE score of 0.907 and a DCI score of 0.649, significantly outperforming strong VAE baselines (e.g., DAVA: 0.410/0.300) and GAN baselines (e.g., ClosedForm: 0.523/0.318).
  • On CelebA, the method achieved the best TAD score (1.154) for attribute-level controllability and the lowest FID (8.1), indicating superior sample quality and reduced leakage compared to diffusion-based baselines like EncDiff and DyGA.
• Ablation Studies: Removing the orthogonality regularizer ($L_{orth}$) resulted in consistent performance drops, confirming its critical role in promoting factor separation.
• Downstream Tasks: In sample-efficient learning tasks using Gradient Boosted Trees (GBT) with limited labeled data (100 and 1,000 samples), the learned representations consistently outperformed baselines, suggesting that factor-wise disentanglement provides a beneficial inductive bias for downstream tasks.
• Visualization: Factor swapping experiments demonstrated that the model could transfer specific attributes (e.g., wall hue, object shape) between images without altering unrelated factors, confirming strong disentanglement.

Significance and Claims

The paper claims to be the first study to investigate flow matching for general disentangled representation learning. Its significance lies in:

• Geometric Basis: Offering a more principled geometric basis for semantic alignment compared to the stochastic trajectories of diffusion models.
• Explicit Control: Providing explicit, factor-level control within the flow-matching framework, addressing the "weak semantic alignment" issue prevalent in current diffusion-based methods.
• Alternative Paradigm: Proposing that viewing representation learning through the lens of structured latent transport offers a fundamentally different and promising alternative to noise-driven diffusion dynamics.

The authors maintain a modest stance, noting that while their method improves upon diffusion baselines, the gains over diffusion models are "smaller but consistent" compared to the substantial gains over VAE and GAN methods, acknowledging that diffusion models already provide strong generative backbones.