What We Don't C: Manifold Disentanglement for Structured Discovery

AI summary

Overview

Research area: Representation learning and generative modelling — specifically latent flow matching, classifier-free guidance, and disentanglement of learned latent representations for scientific discovery. The paper sits at the intersection of machine learning methodology and astrophysics (galaxy morphology).

Technical level: Advanced. The paper assumes familiarity with variational autoencoders, flow matching / optimal transport probability paths, ordinary differential equations, and classifier-free guidance.

Scope: The paper introduces a method for re-purposing frozen, pre-trained latent representations so that a known (already catalogued) feature can be suppressed, making residual, uncatalogued features more accessible — demonstrated on synthetic 2D Gaussians, coloured MNIST, and the Galaxy10 DECaLS dataset.

What This Paper Is About

Existing representation learning methods are validated on features already known to exist in the data, and disentanglement approaches typically try to separate all factors of variation, which fails on real data where ground-truth factors are unknown or entangled. The authors propose "What We Don't C" (WWDC), which instead takes a known feature as conditioning guidance and flows the latent representation backwards to a base distribution, deliberately removing that known feature so that whatever remains — the structure nobody has captured, considered, or catalogued — becomes easier to see. The goal is a cheap, reusable way to iterate: annotate what you know, strip it out, and inspect the residual manifold for what you missed.

Key Contributions

1. A new formulation, "manifold disentanglement." Rather than separating all features into individual dimensions, WWDC disentangles a known feature from an existing frozen representation, producing a rich residual representation. It requires an existing pre-trained manifold, so it avoids retraining representation models from scratch.

2. Theoretical arguments and geometric validation. The authors argue that guidance in the flow path necessarily represses information from the guiding conditional variables, and that because flow matching approximates optimal transport trajectories, the resulting Gaussian base distribution retains the global structure of the original data manifold.

3. Progressive empirical validation. The approach is verified on a fully synthetic 2D Gaussian dataset, then a coloured variant of MNIST where blue is withheld from conditioning, giving a controlled test of whether a non-guided feature surfaces.

4. A real-world astrophysics application. WWDC is applied to Galaxy10 DECaLS galaxy images to isolate morphological class features, producing residuals that separate guided features from the rest of the image, including background and imaging artefacts.

5. Native sample generation / style transfer. Because the flow is bijective, guided base embeddings can be re-initialised with different guidance, producing stylistically consistent samples — e.g. the same digit rendered as a different class while preserving stroke width, position, and colour.

Main Findings

• Structure survives the reverse flow when unguided. In the 2D Gaussian experiment, the class structure present at t = 1 is preserved in the base distribution at t = 0 for an unguided flow, because flow matching approximates optimal transport trajectories. Class retrieval is trivial in this base space, but distance retrieval is not.

• Guidance suppresses the conditioned feature. With class conditioning, classes show no discernible structure at t = 0. At guidance weight ω = 1 there is a turnover at t = 0.5: almost full mutual information with class beyond that time, progressively less before it, and no mutual information at t = 0. Weaker guidance preserves more class-wise mutual information.

• Suppressing one feature surfaces another. In the guided 2D Gaussian case, the Euclidean distance to each Gaussian's centre — a secondary feature — maps simply into the base distribution. A linear model explains increasingly more of the dimension-wise distance toward t = 0 under full guidance, whereas the unguided case gives R² ≈ 0.3. Because the raw problem is non-linear, the R² score is effectively zero at t = 1.

• Coloured MNIST shows the same effect at greater complexity. The VAE latent space is dominated by digit class, and a feature the model was not conditioned on (blue) shows no obvious structure there. After guided reverse flow, class structure almost entirely disappears and a gradient across the blue feature appears. The authors note they do not expect complete class removal, because confounding features informative over class exist (for example, "straightness" for predicting ones).

• Linear probes confirm selective suppression. Digit classification accuracy drops significantly with guidance compared to unguided and VAE representations, and regression to the red and green values used in conditioning is likewise repressed, especially with fewer training samples. Crucially, guidance does not adversely affect the sample size needed to recover blue — the withheld feature is consistently recovered across both flows. The authors attribute the incomplete drop to two factors: class information bound to features not unique to that class, and the intentionally weakly restricted latent space (chosen for high-quality samples and feature interpretability).

• Guided base embeddings support controlled generation. Flowing backwards to t = 0 using one digit's conditioning, then flowing forwards with different guidance, preserves stroke widths (e.g. 0 vs. 3), digit position (e.g. 4 vs. 7), and colour. The unguided base representation gives samples close to the original but with less correspondence.

• Galaxy images can be decomposed into guided and residual parts. Using ω = 3.5 with the 'round' class as guidance (chosen as the least semantically complex structure), the authors isolate class features in galaxy images. Background features remain unchanged, confirming the model identifies the galaxy of interest from simple class labels, and an imaging artefact (a yellow lower half in sample E, not a physical feature) is preserved while the galaxy structure changes. The residual images are a native output of WWDC.

Methodology in Plain English

The method reuses an existing latent space — in this work, one produced by a variational autoencoder — rather than training a new representation model. The pipeline works as follows:

1. Encode the data. A pre-trained VAE maps each data point to a latent vector.

2. Train a flow matching model on the latent space. This model learns a velocity field that transports samples between the latent distribution (target, at t = 1) and a chosen base distribution (at t = 0), typically a unit Gaussian. Because the base Gaussian is a reasonable match for the Gaussianity that a VAE's KL term already imposes, the optimal transport character of the flow minimally distorts the original manifold's structure.

3. Condition on what you already know. The velocity field is trained with classifier-free guidance: with probability p_cfg, the conditioning information (e.g. digit class, red and green values, galaxy class) is replaced by a null vector, so the model learns both conditional and unconditional paths. At inference, guided and unguided velocities can be combined with a weight ω.

4. Run the flow in reverse. Instead of generating new data, the authors take a real VAE sample at t = 1 and solve the ODE backwards to t = 0 using the midpoint method. Information belonging to the conditioning variable is repressed along the way; other structure is transported into the base distribution.

5. Inspect the residual. The resulting base representation is a "meaningful" space where the guided feature is suppressed and remaining features of interest become more detectable — testable with linear probes, t-SNE projections, or by generating residuals between guided and original images.

6. Optionally generate. Since the flow is bijective, re-running forwards from a guided base point with different conditioning swaps the guided attribute while retaining the rest.

Why This Matters

Impact on research: The paper offers a cheaper alternative to full retraining for controlling and repurposing representations — the frozen VAE is reused, and only a comparatively small flow model is trained. It directly addresses a documented impossibility result in the literature (Locatello et al., 2019) by relaxing the disentanglement objective: rather than requiring all factors to be separated, WWDC only attempts to separate one known feature from a manifold. It also enables an iterative discovery loop: annotate, strip out, inspect residuals, then use whatever new features are found as the next round of conditioning.

Real-world applications:

• Astronomy and survey science: Galaxy morphology analysis and separation of physical galaxy features from imaging artefacts, which the authors tie directly to upcoming large surveys such as LSST where high-dimensional complex features are abundant and exploration is limited largely by cost.
• Dataset annotation and filtering: Surfacing subtle features (such as a colour channel) in a latent space that an annotator would otherwise not notice, supporting labelling and data-selection workflows.
• Outlier and residual analysis: Isolation of what a model or measurement has not captured, potentially revealing classes, populations, or artefacts that were not anticipated.
• Controllable synthetic data generation: Cheap style transfer and class-swapping for inspection of hypothetical or counterfactual samples without retraining a generative model.

Industry relevance: Any discipline with high-dimensional datasets and expensive annotation stands to benefit — the paper lists filtering, search, clustering, labelling, outlier detection, and visualisation as the standard uses of learned representations that this method extends. The reliance on existing pre-trained models keeps computational cost low, which matters for iterative, application-driven workflows.

Future Directions

• Quantifying ODE error. Solutions to the ODE used for training and inference introduce an inherent source of error, and no quantification has been undertaken on how this affects representations along the chain, especially where conditioning information is lost. The roles of velocity-field network capacity, inference-time simulation, and optimisation procedure remain open.

• Systematic hyperparameter study. Computational resources limited the work, so the effects of the VAE latent size and the training dropout frequency on the quality of the unguided distributions have not been fully investigated.

• Better conditioning mechanisms. It is unclear which conditioning mechanisms most appropriately and efficiently approximate the guided velocity field, and further work is needed to find the most effective guidance mechanisms.

• Beyond Euclidean state spaces. The work currently considers only ℝ^d. Whether the approach transfers to discrete tokens and quantised VAEs is untested.

• Completing the discovery loop. Figure 1 envisions repeated cycles of annotation and residual inspection; demonstrating a full multi-round discovery pipeline, and how sufficient coverage of newly found features feeds back into conditioning, is left as a pathway rather than a demonstrated result.

Target Audience

Machine learning researchers working on representation learning, disentanglement, and flow matching will find the core methodological contribution and its connection to optimal transport most useful. Astrophysicists and other domain scientists working with high-dimensional observational data — particularly those using VAE latents for exploration and annotation — are the intended practical users, as the Galaxy10 case study shows. Practitioners who need to squeeze more insight out of existing pre-trained models without retraining them, and researchers interested in classifier-free guidance applied outside of traditional sampling, will also benefit. Readers should be comfortable with variational inference, ordinary differential equations, and generative modelling terminology.

Authors’ abstract

Accessing information in learned representations is critical for annotation, discovery, and data filtering in disciplines where high-dimensional datasets are common. We introduce What We Don't C, a novel approach based on latent flow matching that disentangles latent subspaces by explicitly removing information included in conditional guidance, resulting in meaningful residual representations. This allows factors of variation which have not already been captured in conditioning to become more readily available. We show how guidance in the flow path necessarily represses the information from the guiding, conditioning variables. Our results highlight this approach as a simple yet powerful mechanism for analyzing, controlling, and repurposing latent representations, providing a pathway toward using generative models to explore what we don't capture, consider, or catalog.