SAGE: Salient Factor Discovery and Generation with Visual Foundation Representations

AI summary

Overview

Research area: Computer Vision — specifically representation learning, contrastive/disentangled factor analysis, and diffusion-based generative modeling built on frozen visual foundation models.

Technical level: Advanced. The paper assumes familiarity with self-supervised encoders (DINOv3), representation autoencoders (RAEv2), diffusion transformers, flow matching, adversarial training, and InfoNCE objectives.

Scope: SAGE is a two-stage method that splits the frozen spatial latent of a representation autoencoder into a common factor and a salient (target-specific) factor using only background/target dataset labels, then uses the salient factor for unsupervised subtype discovery, image editing, and salient-conditioned image generation.

What This Paper Is About

Given two datasets — for example, faces with eyeglasses (target) and faces without (background) — the paper asks whether a model can learn a representation that captures only what makes the target special (the shape, color, and position of a particular pair of glasses) while ignoring everything the two datasets share (identity, pose, scene). The goal is a salient representation that is detailed enough to reveal subtypes without subtype labels, decodable enough to reconstruct and edit images faithfully, and usable as a conditioning signal to generate new examples of a discovered subtype — even when that subtype has no name or text description.

Key Contributions

1. Salient factor learning inside a frozen foundation latent. SAGE factorizes the high-dimensional spatial latent of a frozen RAEv2 autoencoder (DINOv3-L encoder with multi-layer summation, plus a ViT decoder) into a common factor z_c and a salient factor z_s whose sum the frozen decoder maps back to pixels. No learned bottleneck is used, so instance-level detail such as a specific frame style or stroke shape is retained.

2. An objective set that resolves the underdetermined split. Reconstruction alone is satisfied by an empty salient factor, so the authors add swap adversarial training (erase a target's z_s to get a background-like image; add it to a background's z_c to get a target-like image), sparsity priors (ℓ2² on background salient factors, ℓ1 on target salient factors), image-space cycle consistency, latent-space swap consistency, and a cycle NCE loss that prevents common content from leaking into the salient factor.

3. Unsupervised subtype discovery and label-noise auditing. The salient space recovers ten digits, separates reading glasses from sunglasses, and further splits sunglasses into three styles, without ever using subtype labels. On FFHQ it also surfaced a cluster of mislabeled images, which prompted an audit of all 3,437 target validation images and 125 label corrections (3.64%), including 39 no-glasses images.

4. Salient-conditioned generation by example. A diffusion transformer initialized from the pretrained RAEv2 checkpoint is trained with flow matching on target images only, conditioned on a max-pooled version of the reference image's salient factor. This generates new images that keep the reference's digit or eyewear type while the face or scene varies — without text prompts or subtype labels.

Main Findings

• Near-lossless factorization: SAGE reaches rFID 1.78 on Digits-ImageNet and 1.65 on FFHQ, compared with above 120 for the generative contrastive-analysis baselines cVAE, SepVAE, and Double-InfoGAN. Decoding the unfactorized latent directly gives rFID 0.39 on Digits-ImageNet, so factorization adds little error.

• Clean factor separation: Common-only decoding D(z_c, 0) removes only the digit or glasses, and salient-only decoding D(0, z_s) renders the digit or eyewear alone at its original location and shape. Baselines blur reconstructions and leak faces or scenes into the salient-only output.

• Strong digit discovery: Linear-probe accuracy on Digits-ImageNet is 0.950 for SAGE versus at most 0.281 for baselines and 0.328 for the unfactorized RAEv2 latent; since SAGE and the raw latent share an encoder, the gain comes from factorization. Clustering gives ARI/NMI 0.337/0.472 versus at most 0.009/0.019.

• Content stays in the common factor: ImageNet category, a property of the shared scene, is decoded at 0.765 from the full latent and z_c but only 0.045 from z_s.

• Strong FFHQ discovery: SAGE reaches probe accuracy 0.983 and ARI/NMI 0.939/0.865, ahead of DINOv3 + SepCLR at 0.976 and 0.915/0.825. With k=3 after the audit (including no-glasses images), SAGE reaches 0.937/0.865 while all baselines remain below 0.46/0.56; DINOv3 + SepCLR drops to 0.453/0.550.

• Background/target separation: Silhouette scores of 0.820 (Digits-ImageNet) and 0.873 (FFHQ) are the highest among all compared methods.

• Label noise discovered: On FFHQ, cluster C contains images annotated as wearing glasses that have none, often due to face paint or hat-brim shadows; SAGE places 34 of the 39 audited no-glasses images there.

• Editing works by factor arithmetic: Adding a target's z_s to a background transfers that specific attribute (a digit 7 keeps its location and stroke shape; thick black frames appear on a new face), zeroing z_s removes the attribute, and interpolating salient factors traverses subtypes (digit 1 to 7; reading glasses darkening into sunglasses) on a fixed image.

• Generation follows the discovered subtype: On Digits-ImageNet, salient conditioning raises subtype accuracy from 27.7% to 90.5% and Vendi diversity from 2.65 to 12.50 (real images: 15.79), while lowering cosine similarity to the reference from 0.794 to 0.116. On FFHQ both conditions preserve eyewear type (96.6% versus 98.5%), but only salient conditioning yields varied samples (Vendi 6.96 versus 2.32; Cos-to-ref 0.470 versus 0.779). The raw DINOv3 condition attains lower conditional gFID in both cases, which the authors attribute to closer reference copying rather than subtype control; SAGE has lower unconditional gFID on both datasets.

• Medical transfer with only normal/disease labels: On OCT-Kermany, trained with only normal-versus-disease labels, the salient space separates three diseases (CNV, DME, DRUSEN) with linear-probe accuracy 0.960 and ARI/NMI 0.415/0.415. Adding a disease scan's salient factor to a normal scan introduces fluid-filled cavities; removing it yields a retina closer to normal. The authors describe these edits as qualitative, not clinical evidence.

• Ablations identify what each loss does: Removing the swap adversarial loss collapses digit probing from 0.950 to 0.123 and ARI/NMI to nearly zero, leaving z_s near its zero initialization. Removing target sparsity, Cycle NCE, or either consistency loss raises the salient ImageNet probe from 0.045 to between 0.086 and 0.302; removing the image-space cycle loss also worsens reconstruction (rFID 1.78 to 2.84).

Methodology in Plain English

Stage 1 — splitting the latent. A frozen pretrained encoder turns each image into a spatial grid of features, and a frozen decoder turns that grid back into an image. Two small trainable networks sit between them: one extracts the shared content, the other extracts the target-specific content. Adding the two gives back the original grid, which the frozen decoder renders.

Because a perfect reconstruction is possible with the salient network outputting nothing, the training adds pressures that force target content into it:

• Background images are reconstructed with the salient part set to zero, so the salient network starts empty and only grows when needed.
• Target salient factors are swapped onto background content and vice versa. Discriminators judge whether the results look like genuine background or genuine target images, so erasing the salient part must erase the attribute, and adding it must introduce the attribute.
• Sparsity penalties keep background salient factors at zero and keep target salient factors small, so shared content such as a face is not smuggled in.
• Cycle losses re-encode the swapped images and check that the original factors come back, which forces each salient factor to describe that image's glasses rather than a generic prototype of "glasses."
• A contrastive (InfoNCE) loss applied before and after the swap keeps the two factors from overlapping.

Stage 2 — generating from a discovered subtype. The learned salient factor of a reference image is pooled across space into one vector, projected to a single token, and fed to a diffusion transformer initialized from the pretrained RAEv2 checkpoint. The transformer is trained with flow matching on target images only, with a null token randomly substituted so classifier-free guidance can be used. At sampling time, different noise samples vary the shared content while the condition fixes the subtype — conditioning on only the salient part, rather than the full latent, is what stops the model from copying the reference person or scene.

Evaluation. Discovery is measured by linear probes and by k-means clustering (ARI/NMI) of the salient representations; reconstruction by rFID, PSNR, SSIM, and LPIPS; generation by gFID, the share of samples a classifier assigns the reference subtype, Vendi diversity, and cosine similarity to the reference. Subtype labels are used only for evaluation.

Why This Matters

The paper shows that a frozen foundation model's latent space can be surgically split using nothing more than which dataset an image came from — no attribute labels, no text descriptions — and that the resulting split is simultaneously good for reconstruction, discovery, editing, and generation. This matters most where the difference between two groups is real but unnamed, which is the normal situation in clinical imaging.

Real-world applications:

• Clinical imaging markers: contrasting patient and control scans to find undocumented imaging differences, as the paper illustrates on retinal OCT and motivates with the example of autism diagnosed behaviorally but not yet by brain MRI.
• Dataset curation and label auditing: the FFHQ result shows the salient space can flag mislabeled images automatically; the audit corrected 125 of 3,437 labels.
• Attribute-level image editing: adding, removing, or interpolating a factor lets a user transfer a specific accessory or digit between images while the person or scene is preserved.
• Example-driven content generation: producing new samples of a visual subtype specified by a single reference image, in settings where the subtype has no name a user could type into a prompt.

Industry relevance: the method reuses pretrained foundation encoders and diffusion transformers without retraining them, which is attractive for teams that already license or host such models; it gives a label-efficient route to fine-grained product or attribute variation (eyewear, accessories, digits, lesions) and a practical tool for cleaning large annotated datasets before training.

Future Directions

• Beyond a single binary split. SAGE takes one background and one target dataset and one dataset label per image; handling several target groups or nested attributes at once is not addressed.
• Clinical validation. The OCT disease edits are explicitly presented as qualitative and not clinical evidence; whether the discovered directions correspond to real pathology needs study by domain experts.
• Better conditional generation fidelity. SAGE's conditional gFID is higher than the raw-latent baseline's on both datasets (4.09 versus 2.40 on Digits-ImageNet; 16.03 versus 11.17 on FFHQ), so improving sample quality without reintroducing reference copying is an open problem.
• Preserving spatial layout in the condition. Generation conditions on a channel-wise max-pooled vector that discards location; retaining spatial structure could let generated attributes appear at controlled positions, similar to what the salient-only decodings already show.
• Combining with text or instruction conditioning. The method is positioned as an alternative when no description exists; merging the two, so a discovered subtype can also be named and steered linguistically, is a natural extension.

Target Audience

Researchers and graduate students working on disentangled and contrastive representation learning, latent diffusion and representation autoencoders, or generative models for scientific and medical imaging. It is also useful for practitioners who need controllable, example-conditioned image generation or automated dataset auditing, and for computational scientists in neuroscience and ophthalmology who have group labels but not yet a description of what distinguishes the groups. Readers without background in self-supervised encoders, adversarial losses, and flow matching will find the methodology dense.

Authors’ abstract

Given a target dataset, such as faces with eyeglasses, and a background dataset, such as faces without, contrastive analysis separates \textit{salient} factors specific to the target from \textit{common} content shared by both. We aim for salient representations that capture target-specific detail in each image, such as the shape, color, and position of the glasses, so that they reveal subtypes without subtype labels and guide the generation of new examples of a discovered subtype, even one with no name or text description. We introduce SAGE, which learns both factors directly in the high-dimensional spatial latent of a frozen representation autoencoder and conditions a diffusion transformer on the learned salient representation of a reference image. On Digits-ImageNet and FFHQ eyeglasses, SAGE combines high-fidelity \textit{reconstruction} (rFID below $2$) with unsupervised \textit{subtype discovery}, recovering the digits better than baselines (probe accuracy $0.950$ vs.\ at most $0.281$) and revealing eyewear types, finer sunglasses styles, and mislabeled images; salient-conditioned \textit{generation} raises Digits-ImageNet subtype accuracy over the unfactorized latent ($90.5\%$ vs.\ $27.7\%$) and diversity on both datasets. On retinal OCT, SAGE's salient space separates three diseases using only normal/disease labels.