Distilling Latent Manifolds: Resolution Extrapolation by Variational Autoencoders

Jiaming Chu; Tao Wang; Lei Jin

cs.CVMar 15, 2026

Distilling Latent Manifolds: Resolution Extrapolation by Variational Autoencoders

Authors: Jiaming Chu, Tao Wang, Lei Jin

Organizations: School of Electronic Engineering, Beijing University of Posts and Telecommunications, Beijing, China

Abstract

Variational Autoencoder (VAE) encoders play a critical role in modern generative models, yet their computational cost often motivates the use of knowledge distillation or quantification to obtain compact alternatives. Existing studies typically believe that the model work better on the samples closed to their training data distribution than unseen data distribution. In this work, we report a counter-intuitive phenomenon in VAE encoder distillation: a compact encoder distilled only at low resolutions exhibits poor reconstruction performance at its native resolution, but achieves dramatically improved results when evaluated at higher, unseen input resolutions. Despite never being trained beyond $256^2$ resolution, the distilled encoder generalizes effectively to $512^2$ resolution inputs, partially inheriting the teacher model's resolution preference.We further analyze latent distributions across resolutions and find that higher-resolution inputs produce latent representations more closely aligned with the teacher's manifold. Through extensive experiments on ImageNet-256, we show that simple resolution remapping-upsampling inputs before encoding and downsampling reconstructions for evaluation-leads to substantial gains across PSNR, MSE, SSIM, LPIPS, and rFID metrics. These findings suggest that VAE encoder distillation learns resolution-consistent latent manifolds rather than resolution-specific pixel mappings. This also means that the high training cost on memory, time and high-resolution datasets are not necessary conditions for distilling a VAE with high-resolution image reconstruction capabilities. On low resolution datasets, the distillation model still could learn the detailed knowledge of the teacher model in high-resolution image reconstruction.

Figures & tables

Fig. 1: The figure shows low resolution image sample, including the reconstruction results of the Flux [ 3 ] VAE encoder as the teacher model (34MB), the distillation model (2MB) which only training on low resolution images, and the distillation model with the resolution augment. The second row also provides the MSE $\downarrow$ , LPIPS distance $\downarrow$ (Lower is better), and differences from the original image for the three.

Fig. 2: The figure includes the resolution generalization of manifolds during the distillation process of the model and the manifold distribution trained from scratch on low resolution datasets in general. Although distillation and training are carried out at around 256 resolutions, student models can learn the generalization knowledge about high-resolution from teacher models during the distillation process, while models trained from scratch can only achieve generalization around low resolutions. The proof of the distillation model’s cross resolution generalization ability and the degree of matching between manifolds can be referred to V and VI .

Fig. 3: From the figure, we can intuitively observe the consistency between the distillation model (2MB) and the teacher model (34MB) in terms of cross resolution performance trends, and the variational encoder (24MB) trained on low resolution images lacks the ability to generalize across resolutions.

Fig. 4: The above figures show the input original image, the reconstruction results of the teacher model, the reconstruction results of the student model, and the image reconstruction results after 1.5-fold upsampling enhancement. From the visualization results, we can see that the student model has little difference from the teacher model in terms of layout, color, brightness, etc., but lacks high-frequency detail information. After upsampling and reinforcement, the reconstruction results surpass the teacher model even in terms of details.

Scale Factor	Model	MSE (×10 -4 ) $\downarrow$	PSNR (dB) $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	rFID $\uparrow$
0.5	Flux VAE	$8.31_{\pm 0.81}$	$30.82_{\pm 0.42}$	$0.927_{\pm 0.005}$	$0.0126_{\pm 0.0009}$	0.069
0.5	Distilled VAE	$20.04_{\pm 1.57}$	$26.99_{\pm 0.34}$	$0.859_{\pm 0.007}$	$0.0478_{\pm 0.0023}$	1.031
1.0	Flux VAE	$15.05_{\pm 2.45}$	$28.28_{\pm 0.71}$	$0.914_{\pm 0.009}$	$0.0184_{\pm 0.0020}$	0.112
1.0	Distilled VAE	$32.11_{\pm 4.04}$	$24.97_{\pm 0.55}$	$0.832_{\pm 0.012}$	$0.0586_{\pm 0.0037}$	1.400
1.5	Flux VAE	$1.96_{\pm 0.39}$	$37.18_{\pm 0.89}$	$0.977_{\pm 0.002}$	$0.0081_{\pm 0.0009}$	0.033
1.5	Distilled VAE	$7.43_{\pm 0.97}$	$31.33_{\pm 0.56}$	$0.929_{\pm 0.006}$	$0.0395_{\pm 0.0024}$	0.822

TABLE I: Cross resolution test results of distillation model

Fig. 5: Resolution-dependent latent statistics of the teacher (Flux VAE) and the distilled student encoder. The empirical latent mean and standard deviation are computed over ImageNet validation images at resolutions ranging from $64^{2}$ to $1024^{2}$ . Both models exhibit highly consistent trends. Despite being trained only on low-resolution images, the student closely follows the teacher’s scaling behavior, indicating that distillation transfers resolution-aware latent parameterization rather than fixed-resolution representations.

Fig. 6: Latent distribution alignment between the teacher and student across resolutions, measured by Kullback–Leibler (KL) and Jensen–Shannon (JS) divergences. Divergence is minimized in the intermediate resolution range ( $256^{2}$ – $384^{2}$ ), which coincides with the peak reconstruction and perceptual performance of the distilled model. This observation is consistent with the fact that the distillation model is explicitly supervised on 256 resolution data.

Fig. 7: t-SNE visualization of latent embeddings produced by the teacher and distilled student encoders on identical image samples across different resolutions. Latent representations from both models exhibit nearly identical geometric structures, with resolution acting as a radial factor: low-resolution inputs cluster near the center, while higher resolutions progressively expand outward. The strong overlap between teacher and student embeddings demonstrates that distillation successfully transfers the global geometry of the teacher’s latent manifold.

Fig. 8: In this figure, we show the results of latent codes interpolation and cross resolution latent codes interpolation in the same resolution direction on the manifold of the student model. From the visualization of interpolation results, we can see that the results of linear interpolation at the same resolution are smooth, while cross resolution interpolation introduces interpolation artifacts.

Resolution	$C_{hidden}$	Params	MSE	PSNR	SSIM	LPIPS
256	16	0.1MB	78.21	21.56	0.754	0.1247
	32	2.2MB	32.11	24.97	0.832	0.0586
	64	6.2MB	21.42	25.73	0.844	0.0492

TABLE II: The influence of model capacity on distillation performance, MSE should times

10^{-4}

Interpolation Types	Position	MSE $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
-	-	32.11	24.97	0.832	0.0586
Bilinear	Post	24.62	26.15	0.847	0.0548
Bilinear	Pre	$7.43$	$31.33$	$0.929$	$0.0395$
Bicubic	Post	29.58	25.37	0.840	0.0567
Bicubic	Pre	8.18	31.23	0.920	0.0303

TABLE III: The Influence of Linear Interpolation Types and Positions on Model Performance. ”Pre” means the images are resized before input; ”Post” means the image are resized after reconstruction, which could observe the effect of only interpolation.

Loss	Training Hours	MSE	PSNR	SSIM	LPIPS
L1	2H	42.23	23.66	0.784	0.0862
Huber	2H	32.11	24.97	0.832	0.0586
Huber+LPIPS	16H	32.31	25.01	0.824	0.0602
Huber+LPIPS+Recon	20H	38.44	23.75	0.801	0.0634
Huber+LPIPS+KL	20H	67.02	20.31	0.730	0.1413

TABLE IV: The influence of different loss functions on model distillation training. Huber loss performs best.

Model	$T_{infer}$ (ms)	params(MB)	$Mem_{GPU}(MB)$	PSNR
LiteVAE	5.32	6.6	603.1	25.01
TinyVAE	3.67	2.4	522.2	24.68
FluxVAE	8.21	34.2	794.5	28.28
Ours	3.53	2.2	488.3	24.97
Ours(1.5x)	4.02	2.2	540.6	31.33

TABLE V: Encoding efficiency of distillation models compared to other lightweight VAEs (Only Encoder).

Explore similar work

Dec 8, 2025

Distribution Matching Variational AutoEncoder

Most visual generative models compress images into a latent space before applying diffusion or autoregressive modelling. Yet, existing approaches such as VAEs and foundation model aligned encoders implicitly constrain the latent space without explicitly shaping its distribution, making it unclear which types of distributions are optimal for modeling. We introduce Distribution-Matching VAE (DMVAE), which explicitly aligns the encoder's latent distribution with an arbitrary reference distribution via a distribution matching constraint. This generalizes beyond the Gaussian prior of conventional VAEs, enabling alignment with distributions derived from self-supervised features, diffusion noise, or other prior distributions. With DMVAE, we can systematically investigate which latent distributions are more conducive to modeling, and we find that SSL-derived distributions provide an excellent balance between reconstruction fidelity and modeling efficiency, reaching gFID equals 3.2 on ImageNet with only 64 training epochs. Our results suggest that choosing a suitable latent distribution structure (achieved via distribution-level alignment), rather than relying on fixed priors, is key to bridging the gap between easy-to-model latents and high-fidelity image synthesis. Code is avaliable at https://github.com/sen-ye/dmvae.

Sen Ye, Jianning Pei, Mengde Xu +4

State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · University of Chinese Academy of Sciences · Tencent +1

Variational Autoencoders Generative Modeling

Ask PDF

May 13, 2026

Qwen-Image-VAE-2.0 Technical Report

We present Qwen-Image-VAE-2.0, a suite of high-compression Variational Autoencoders (VAEs) that achieve significant advances in both reconstruction fidelity and diffusability. To address the reconstruction bottlenecks of high compression, we adopt an improved architecture featuring Global Skip Connections (GSC) and expanded latent channels. Moreover, we scale training to billions of images and incorporate a synthetic rendering engine to improve performance in text-rich scenarios. To tackle the convergence challenges of high-dimensional latent space, we implement an enhanced semantic alignment strategy to make the latent space highly amenable to diffusion modeling. To optimize computational efficiency, we leverage an asymmetric and attention-free encoder-decoder backbone to minimize encoding overhead. We present a comprehensive evaluation of Qwen-Image-VAE-2.0 on public reconstruction benchmarks. To evaluate performance in text-rich scenarios, we propose OmniDoc-TokenBench, a new benchmark comprising a diverse collection of real-world documents coupled with specialized OCR-based evaluation metrics. Qwen-Image-VAE-2.0 achieves state-of-the-art reconstruction performance, demonstrating exceptional capabilities in both general domains and text-rich scenarios at high compression ratio. Furthermore, downstream DiT experiments reveal our models possess superior diffusability, significantly accelerating convergence compared to existing high-compression baselines. These establish Qwen-Image-VAE-2.0 as a leading model with high compression, superior reconstruction, and exceptional diffusability.

Zekai Zhang, Deqing Li, Kuan Cao +27

Qwen Team

Variational Autoencoders Learned Image Compression Autoencoders Image Reconstruction

Ask PDF

Mar 23, 2026

DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment

Reducing token count is crucial for efficient training and inference of latent diffusion models, especially at high resolution. A common strategy is to build high-compression image tokenizers with more channels per token. However, when trained only for reconstruction, high-dimensional latent spaces often lose meaningful structure, making diffusion training harder. Existing methods address this with extra objectives such as semantic alignment or selective dropout, but usually require costly diffusion retraining. Pretrained diffusion models, however, already exhibit a structured, lower-dimensional latent space; thus, a simpler idea is to expand the latent dimensionality while preserving this structure. We therefore propose Detail-Aligned VAE, which increases the compression ratio of a pretrained VAE with only lightweight adaptation of the pretrained diffusion backbone. DA-VAE uses an explicit latent layout: the first

C

channels come directly from the pretrained VAE at a base resolution, while an additional

D

channels encode higher-resolution details. A simple detail-alignment mechanism encourages the expanded latent space to retain the structure of the original one. With a warm-start fine-tuning strategy, our method enables

1024 \times 1024

image generation with Stable Diffusion 3.5 using only

32 \times 32

tokens,

4\times

fewer than the original model, within 5 H100-days. It further unlocks

2048 \times 2048

generation with SD3.5, achieving a

6\times

speedup while preserving image quality. We also validate the method and its design choices quantitatively on ImageNet.

Xin Cai, Zhiyuan You, Zhoutong Zhang +1

Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam +1

Diffusion-Based Image Generation Diffusion Models Latent Diffusion Models Diffusion Model Inference Acceleration

Ask PDF