RoAR-VAE: Improving Latent Diffusability through Robust Prediction

RoAR-VAE. Improving Latent Diffusability through Robust Prediction. By Yang Yue and six others, from Tsinghua University and Microsoft Research. Published on the nineteenth of October. Its code repository is a placeholder: a readme, and nothing to run yet.

A visual tokenizer has two jobs that have to be balanced. It must preserve the fidelity of the image, and it must produce latents that a downstream diffusion model can learn efficiently, a property called diffusability. Earlier work improves diffusability in one of two ways: aligning the latents with semantic features, or training the tokenizer and the diffusion model jointly. Both routes bundle several ingredients, and the question the authors ask is which of them actually matters.

The authors revisit REPA-E, a representative approach, and take parts of it away. Remove its diffusion objective and its dedicated generative blocks, keep only the prediction of semantic features from corrupted latents, and downstream generation stays comparable under a matched setting.

That observation becomes a general principle they call robust prediction. Information that matters for generation should remain recoverable even when the latents are corrupted.

The principle is instantiated with two complementary objectives. RobustAlign predicts pretrained semantic features from corrupted latents, to preserve high-level structure. RobustRecon reconstructs the clean image, to retain fine-grained spatial and appearance information. Together they make up the Robust Alignment and Reconstruction VAE, or RoAR-VAE.

The authors also offer an account of why this helps. Their analysis suggests robust prediction makes useful information more accessible along the high-variance directions of the latent space, and those are the directions diffusion models tend to learn earlier.

Under the same training budget, RoAR-VAE improves average performance on generation benchmarks by two point one nine points over the strongest baseline tokenizer. On the reconstruction side, its accuracy on small text improves by sixteen point nine three percentage points over FLUX.1-VAE. The two numbers speak to the two jobs: one to how learnable the latents are, the other to how much detail they keep.

The method grew out of subtraction. Strip a joint training scheme down to one piece, see that the piece alone performs comparably, and name the principle behind it. RoAR-VAE then pairs that principle with reconstruction, for what the authors call a favourable balance between fidelity and generative learnability.