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

Distribution Matching Variational AutoEncoder
Dec 8, 2025Sen Ye, Jianning Pei, Mengde Xu +4Variational Autoencoders Generative Modeling
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · University of Chinese Academy of Sciences · Tencent +1
Qwen-Image-VAE-2.0 Technical Report
May 13, 2026Zekai Zhang, Deqing Li, Kuan Cao +27Variational Autoencoders Learned Image Compression
Qwen Team
DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment
Mar 23, 2026Xin Cai, Zhiyuan You, Zhoutong Zhang +1Diffusion-Based Image Generation Diffusion Models
Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam +1

Dec 8, 2025J/K move · Enter open · S save

Distribution Matching Variational AutoEncoder

Sen Ye, Jianning Pei, Mengde Xu, Shuyang Gu+3

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Distribution Matching Variational AutoEncoder

Qwen-Image-VAE-2.0 Technical Report

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