core

PS-VAE

· 2025-12 · arXiv:2512.17909

Why it matters

Starts from DINOv2 and SigLIP2 features and does general instruction editing with a Transfusion-style transformer diffusing in a 96-channel latent. Its importance to this atlas is the reason it does not keep the encoder frozen: the ablation states that a representation autoencoder's performance is strongly limited by weak reconstruction, producing details inconsistent with the input image. The paper unfreezes the encoder in its second stage to fix exactly that.

What this paper does

read: full text

Before this

Representation encoders such as DINOv2 and SigLIP2 make excellent understanding features but were not built to be generative latents. RAE was the first paper to use them directly, freezing the encoder and training only a decoder. That fixed reconstruction quality relative to using the raw features untouched, but this paper argues RAE's high-dimensional, unregularized feature space still lets diffusion models land on off-manifold latents with structurally distorted objects, and the frozen encoder's purely discriminative training still discards high-frequency pixel detail the decoder can never recover.

The problem

Adapt an understanding-oriented representation encoder's features into a generative latent that is simultaneously well regularized, so diffusion stays on-manifold, and capable of state-of-the-art pixel reconstruction, then build a unified text-to-image and editing model on top of it.

The idea

Compress DINOv2 or SigLIP2 features into a compact, KL-regularized 96-channel latent with a dedicated semantic encoder-decoder first, then in a second stage unfreeze the representation encoder itself and let a pixel-reconstruction loss push gradients back into it, holding it in place with a semantic-consistency loss against its original, frozen features.

How it works

Stage one, S-VAE, freezes the pretrained DINOv2 or SigLIP2 encoder and trains a semantic encoder and decoder mapping its 768-dimensional features into a 96-channel, 16x16-downsampled, KL-regularized latent, which addresses the off-manifold problem. Stage two, PS-VAE, unfreezes the representation encoder during pixel-decoder training by removing the detach operation that had blocked gradients, so the pixel-reconstruction loss now updates the encoder directly; a semantic-reconstruction loss against the original frozen encoder's features is added alongside the retained KL and pixel losses, so the encoder gains reconstruction detail without losing its semantic structure. This directly confirms the pattern already noted elsewhere in this atlas, that a frozen encoder's weak reconstruction fidelity limits editing quality, and the paper's own fix is to unfreeze it. The unified editing model uses a Transfusion-style shared transformer, in which clean latents of the input image, instruction text embeddings, and noisy target latents are concatenated into one sequence, with a full attention mask over the clean and noisy image latents and a causal mask over the instruction text. Editing happens directly in this 96-channel representation-model latent space. The fetched text describes no mask, attention injection, latent blending, or locality loss protecting regions the instruction did not target; consistency with the input depends entirely on what the model has learned, not on an architectural constraint.

Evidence

On ImageNet-1K reconstruction at stride 16, PS-VAE reaches rFID 0.203, PSNR 28.79, SSIM 0.817, against RAE's rFID 0.619, PSNR 19.20, SSIM 0.436, and MAR-VAE's rFID 0.534, PSNR 26.18, SSIM 0.715. On text-to-image, PS-VAE scores GenEval 76.56 and DPG-Bench 83.62 against RAE's 71.27 and 81.72. On instruction editing measured by EditingReward on GEdit-Bench, PS-VAE scores 0.222 against RAE's 0.059 and MAR-VAE's 0.056.

Limitations

The authors concede training is done exclusively on ImageNet-1K and suggest future work could use larger, more diverse data. They train and evaluate only at 256x256 resolution and leave higher resolutions to future work. They also report that under equal loss weighting, SigLIP2 reconstruction saturates earlier than DINOv2's, requiring per-encoder manual tuning of the semantic-to-pixel loss ratio.

Why it matters

It gives a precise, direct confirmation that a frozen representation encoder's weak reconstruction fidelity is a real bottleneck for using these features as a generative and editing latent, and that unfreezing the encoder with a semantic-consistency anchor is a working fix, the same conclusion reached independently by RPiAE.

Abstract, in the authors' own words

Modern Latent Diffusion Models (LDMs) typically operate in low-level Variational Autoencoder (VAE) latent spaces that are primarily optimized for pixel-level reconstruction. To unify vision generation and understanding, a burgeoning trend is to adopt high-dimensional features from representation encoders as generative latents. However, we empirically identify two fundamental obstacles in this paradigm: (1) the discriminative feature space lacks compact regularization, making diffusion models prone to off-manifold latents that lead to inaccurate object structures; and (2) the encoder's inherently weak pixel-level reconstruction hinders the generator from learning accurate fine-grained geometry and texture. In this paper, we propose a systematic framework to adapt understanding-oriented encoder features for generative tasks. We introduce a semantic-pixel reconstruction objective to regularize the latent space, enabling the compression of both semantic information and fine-grained details into a highly compact representation (96 channels with 16x16 spatial downsampling). This design ensures that the latent space remains semantically rich and achieves state-of-the-art image reconstruction, while remaining compact enough for accurate generation. Leveraging this representation, we design a unified Text-to-Image (T2I) and image editing model. Benchmarking against various feature spaces, we demonstrate that our approach achieves state-of-the-art reconstruction, faster convergence, and substantial performance gains in both T2I and editing tasks, validating that representation encoders can be effectively adapted into robust generative components.

Research line

Editing in a representation latentemerging

Problem

Semantic / foundation-model latents discard much of the high-frequency pixel detail (exact color, texture, fine structure) that faithful reconstruction — and, later, edit-region preservation — depends on.

Method note — the shared flow-matching interpolation

Every flow-matching / rectified-flow paper in this atlas trains toward a straight-line path between a noise sample x₀ and a data sample x₁:

x_t = (1-t)\,x_0 + t\,x_1, \qquad v_\theta(x_t, t) \approx x_1 - x_0

Builds on

Nothing recorded yet.

Built on by

Nothing recorded yet.