Source-linked AI summary
WithinVAE: Compact Visual Autoencoding with Self-Distilled Semantics
Yuzhou Ji, Jiepeng Wang, Tianyu Jin, Lizhuang Ma, Haibin Huang, Chi Zhang, Xuelong Li
TL;DR
Visual tokenizers must balance reconstruction detail, compactness, and semantic structure, while pretrained recognition encoders can impose invariances mismatched to inversion. WithinVAE learns semantic targets from an EMA copy of the autoencoder and uses position-dependent token masking. It improves all four reconstruction metrics over RPiAE at half its scalar capacity, with frozen linear probing also improving over RPiAE.
Problem
Visual tokenizers must preserve detail, compactness, and semantics, but recognition-oriented invariances can suppress appearance variation needed for reversible generation.
Source: §1 · p. 1Method
WithinVAE combines a sequential variational bottleneck and position-dependent token masking with global and local self-distillation from the autoencoder’s EMA copy.
Source: §1 · p. 2Results
At half RPiAE’s scalar capacity, WithinVAE improves all four reconstruction metrics, including PSNR by 1.36 dB, and frozen linear probing improves by 6.482 points over RPiAE.
Source: §1 · p. 2, §4.2 · p. 6Takeaways & Limitations
A compact visual autoencoder can organize semantic structure without borrowing a recognition representation, while retaining reconstruction quality at reduced scalar capacity.
Source: §5 · p. 10Better reconstruction and frozen semantics did not produce better generation; the study also counts capacity in continuous scalars rather than entropy-coded bits.
Source: §4.8 · p. 9
Abstract
from arXiv · showhide
Visual autoencoders must balance reconstruction fidelity, compactness, and semantic structure. Existing semantic tokenizers typically inherit semantics from pretrained recognition encoders, whose invariances can discard the appearance detail required for inversion. We show that this inheritance is unnecessary: semantic structure can be learned inside an autoencoder's own reconstructive representation. We introduce WithinVAE, a compact variational visual autoencoder with a sequential bottleneck, DINO/iBOT-style self-distillation from an exponential-moving-average copy of the autoencoder itself, and position-dependent token masking to achieve a 4-fold compression ratio, which is free of pretrained representation encoder or external semantic target. On ImageNet-1K at $256\times256$, WithinVAE improves all four reported reconstruction metrics over RPiAE including PSNR by 1.36 dB, while using half its latent scalar capacity. Under frozen evaluation, the same latents improve linear probing by 6.5 points over RPiAE and 5.1 points over our reconstruction-only baseline, attributing the semantic gains to self-distillation rather than architecture or external supervision. An auxiliary latent-to-latent model reproduces the ordered-compression trend, and a pixel-to-latent model transfers the objective to a pretrained generative space. These results separate the origin of semantic supervision from its use: a compact visual autoencoder can organize semantic structure without borrowing a recognition representation.
1 INTRODUCTION
Visual tokenizers must balance detail, compactness, and semantics, but pretrained recognition representations may impose invariances mismatched to reversible reconstruction. WithinVAE asks whether semantic structure can instead emerge inside an autoencoder’s own reconstructive representation.
1 INTRODUCTION
- The motivating gap is that recognition-oriented invariance can suppress appearance variation needed for reversible generation, while stronger compression also removes information.Existing approaches often transfer representations from pretrained vision encoders, coupling the tokenizer to an understanding model trained for a different objective.Source: §1 · p. 1
- Frozen linear probing improves by 6.482 points over RPiAE and 5.072 points over the reconstruction-only baseline.Nearest-neighbor gains also rule out collapse, and the reconstruction-only comparison attributes the semantic gain to self-distillation.Source: §1 · p. 2
- WithinVAE produces semantic targets online from its own EMA copy rather than transferring them from a pretrained recognition encoder.Its global DINO-style targets come from pooled compact tokens, while local iBOT-style targets apply to masked image patches.Source: §1 · p. 2
- At half RPiAE’s scalar capacity, WithinVAE improves all four reconstruction metrics, including PSNR by 1.36 dB.The reported metrics are rFID, PSNR, LPIPS, and SSIM.Source: §1 · pp. 2–3
- The 128-token prefix exceeds RPiAE in PSNR and SSIM at one quarter of its scalar capacity.In the auxiliary latent-to-latent variant, K = 64 uses 6,144 scalars and already exceeds RPiAE’s 21.3 dB PSNR.Source: Figure 1 · p. 2, §1 · p. 3
2 RELATED WORK
Prior work explores both compact one-dimensional token sequences and semantics transferred from pretrained representation models. WithinVAE retains the query-based sequence interface while making its semantic targets originate within the autoencoder.
Visual tokenization and compact sequences
- Compression adds demands to visual representations spanning image, video, and structured 3D tasks.The cited work includes Xi et al. (2026b), An et al. (2026), Wang et al. (2025), and Chen et al. (2026b).Source: §2 · p. 3
- Compact visual tokenization ranges from discrete codes and spatial latents to query-based sequences, including TiTok and One-D-Piece, which orders tokens with Tail Token Drop.WithinVAE retains query-based sequence topology but uses continuous variational tokens with a position-dependent reliability profile.Source: §2 · p. 3
Representation priors for generation
- Many semantic tokenizers align to pretrained representations, while RPiAE uses a frozen pivot and variational bridge; WithinVAE instead generates targets online from its EMA copy.Related approaches include REPA, ReaLS, VA-VAE, PS-VAE, RAE, FAE, SVG, and AlignTok.Source: §2 · p. 3
Self-supervision inside an autoencoder
- WithinVAE borrows DINO and iBOT optimization forms, not their pretrained representations; its global branch supervises pooled compact tokens and its local branch regularizes the shared image-patch encoder.The local masked-prediction loss does not directly supervise individual compact tokens.Source: §2 · p. 3
3 METHOD
WithinVAE couples an image-reconstructing variational bottleneck to an EMA teacher–student semantic path. Position-dependent corruption orders token reliability, while global and local self-distillation apply complementary pressures to compact tokens and the shared encoder.
3.1 WITHINVAE: COMPACT VARIATIONAL AUTOENCODER
- A query encoder and token-wise variational bottleneck map an image to 256 continuous 32-dimensional tokens, which a pixel decoder reconstructs.The model uses ViT-L encoder and decoder backbones; shape-compatible weights are initialized from One-D-Piece, while boundary projections and the bottleneck are random.Source: §3.1 · p. 4, Figure 2 · p. 4
3.2 POSITION-DEPENDENT TOKEN RELIABILITY
- Increasing token masking probability with sequence position makes early tokens statistically more reliable for reconstruction.With pmax = 0.5, the surviving prefix reaches the pixel decoder; variable-length prefix decoding tests this ordering without retraining for each K.Source: §3.2 · p. 4
3.3 GLOBAL AND LOCAL SEMANTIC SELF-DISTILLATION
- Global DINO-style self-distillation matches mean-pooled compact tokens across views, while local iBOT-style prediction matches student masked patches to clean-view EMA-teacher targets.The local branch bypasses the compact-token mean pool and shapes compact tokens only indirectly through the shared encoder.Source: §3.3 · pp. 4–5
3.4 COMPLEMENTARY ROLES OF THE OBJECTIVES
- Position-dependent corruption sets decoder-facing token reliability, global self-distillation regularizes the downstream-facing compact representation, and local prediction supplies patch-level pressure through the shared encoder.The EMA teacher co-evolves with the reconstructive model; reconstruction retains image detail while self-distillation discourages a representation useful only to its paired decoder.Source: §3.4 · p. 5
4 EXPERIMENTS
The experiments test reconstruction and capacity, frozen semantic quality, prefix reliability, and transfer to generation. They use ImageNet-1K at 256 × 256 and evaluate trained autoencoders and EMA teachers.
4.1 IMPLEMENTATION DETAILS
- Experiments assess whether self-distillation preserves reconstruction at lower capacity, improves frozen semantic probes, supports prefix decoding, and transfers across interfaces.Each claim is tied to its corresponding metric rather than a single score.Source: §4.1 · p. 5
- Training uses ImageNet-1K at 256 × 256 with random crops and flips, scheduled warm-up for KL, adversarial, and semantic losses, and eight H100 GPUs.The main autoencoder trains for 10 epochs; LightningDiT trains for 80.Source: §4.1 · pp. 5–6
- The ViT-L encoder maps a 16 × 16 patch grid and 256 learned queries to a 32-dimensional token sequence that a ViT-L decoder reconstructs into pixels.The nonspatial interface uses compact tokens rather than a fixed h × w latent grid; evaluation uses posterior means.Source: §4.1 · pp. 5–6
- The EMA teacher duplicates the autoencoder’s encoder and semantic pathway, removing recognition-feature supervision while retaining generative architectural initialization.Shape-compatible parameters come from One-D-Piece; new parameters include the variational bottleneck and semantic head.Source: §4.1 · p. 6
- Reconstruction uses all 50,000 ImageNet validation images; continuous-latent scalar count means token count × channel width, a capacity proxy rather than bitrate.Frozen evaluations use mean-pooled compact tokens, with linear probing and exact cosine 20-NN alongside cross-view consistency.Source: §4.1 · p. 6
4.2 QUANTITATIVE RESULTS
WithinVAE’s main comparison with RPiAE shows improved reconstruction at half the scalar capacity, while its shorter prefixes also retain reconstruction quality. An auxiliary latent-to-latent model shows a similar ordered-compression trend.
4.2 QUANTITATIVE RESULTS
- At half RPiAE’s scalar capacity, WithinVAE improves all four metrics: rFID 0.46 vs. 0.50, PSNR +1.36 dB, LPIPS −0.057, and SSIM +0.068.The 128-token prefix uses 4,096 scalars, one quarter of RPiAE’s budget, and exceeds its full-sequence PSNR and SSIM.Source: §4.2 · p. 6
- 22.23 dB PSNR at 64 × 96 and 25.63 dB at 128 × 96 show useful prefix inference in the auxiliary latent-to-latent model.The 128 × 96 prefix is within 0.64 dB of the full sequence despite training without contiguous tail removal.Source: §4.2 · p. 6
4.3 ABLATION STUDY
Ablations separate architectural and semantic contributions and show that the complete objective improves reconstruction and frozen semantic performance together. The section also describes attention maps as a controlled but non-causal diagnostic.
4.3 ABLATION STUDY
- The full objective reaches 71.886% linear-probe accuracy, 5.072 points above the reconstruction-only variant and 6.482 above RPiAE.The table reports reconstruction and frozen-latent semantics; its full objective improves fidelity and linear-probe accuracy over every partial variant.Source: §4.3 · p. 6, Table 2 · p. 7
- Only the complete objective improves fidelity and semantics together, reaching 22.66 dB PSNR and 0.159 LPIPS (+0.76 dB, −0.016 versus ②) while adding a 5.072-point probe gain.It also improves 20-NN by 5.538 points over ① and 2.959 over ③, with a 99.125 view score.Source: §4.3 · p. 7
- The WithinVAE attention map averages latent-query-to-patch attention at EMA teacher layer 10, while RPiAE averages bridge attention over heads and queries.Inputs and map normalization are controlled, but the operators are architecture-specific and attention is not evidence of causal token contribution.Source: §4.3 · p. 7
4.4 DOWNSTREAM GENERATION
A 675M-parameter LightningDiT is trained directly on WithinVAE latents for 80 epochs. WithinVAE has higher unguided recall than RPiAE, but its other generation metrics lag, so reconstruction and frozen semantics do not ensure better generation.
4.4 DOWNSTREAM GENERATION
- Without guidance, WithinVAE exceeds RPiAE on recall (0.66 vs. 0.60), but trails on gFID (9.12 vs. 2.25).Guidance narrows but does not close the gap; the comparison uses the reported on/off settings, not necessarily identical guidance settings.Source: §4.4 · p. 7, Table 3 · p. 8
- Figure 3 pairs matched validation reconstructions and attention maps above with class-conditional samples from the 80-epoch LightningDiT below.The attention displays share crop, normalization, and color map, but use architecture-specific operators.Source: Figure 3 · p. 8
4.5 QUALITATIVE RESULTS
WithinVAE attention is more concentrated on foreground objects than RPiAE bridge maps in the shown examples, while the montage cannot establish coverage or malformed-sample frequency.
4.5 QUALITATIVE RESULTS
- In seven matched validation images, WithinVAE attention concentrates on foregrounds more clearly than RPiAE bridge maps, which appear more fragmented.Both are coarse, architecture-specific diagnostics that can extend into correlated background regions; neither is interpreted as segmentation or causal evidence.Source: §4.5 · p. 8
- The montage shows recognizable animal, object, and landscape structure, but cannot reveal coverage failures or malformed-sample frequency.The passage identifies 50K-sample metrics in Table 3 as the basis for comparing generators.Source: §4.5 · p. 8
- Table 4 reports end-to-end reconstruction at 256 × 256, with RGB metrics after decoding and unequal model-facing scalar capacities.Source: Table 4 · p. 9
4.6 TRANSFER ACROSS RECONSTRUCTION INTERFACES
Auxiliary experiments transfer the training objective across pixel-to-latent and latent-to-latent reconstruction interfaces. The latent-to-latent model retains a prefix-reliability trend, but the interfaces differ in what they test and in capacity.
4.6 TRANSFER ACROSS RECONSTRUCTION INTERFACES
- The latent-to-latent model loses only 0.64 dB from 512 to 128 tokens, reproducing the prefix trend across an architecture distinct from the main model.Its full bottleneck is 512 × 96, larger than the main model's.Source: §4.6 · p. 9
- The auxiliary systems test objective transfer, not architectural or capacity superiority, because they use different architectures and larger full bottlenecks.Source: §4.6 · p. 9
- A 4× sequence reduction in Qwen-latent-to-Qwen-latent reconstruction changes rFID from 0.40 to 0.42 and PSNR from 26.27 to 25.63 dB without retraining.At 64 tokens, PSNR and SSIM drop more sharply, so ordered reliability is useful but not lossless.Source: §4.6 · p. 9
- The pixel-to-Qwen-latent model tests learning compression and fixed-decoder geometry together, so its 23.91 dB PSNR cannot be attributed solely to self-distillation.The latent-to-latent interface removes the initial mapping and reaches 26.27 dB, but exposes 49,152 scalars at full length.Source: §4.6 · p. 9
4.7 ANALYSIS OF EXPERIMENTAL RESULTS
The analysis separates semantic attribution, objective transfer, and generation: removing self-distillation removes the measured semantic gain, while better reconstruction and frozen semantics do not ensure better generation. It also identifies concrete evaluation and rate-coding boundaries.
4.7 ANALYSIS OF EXPERIMENTAL RESULTS
- Removing self-distillation removes the semantic gain, while retaining it improves fidelity at the same time, supporting attribution to the objective rather than architecture or external targets.The analysis distinguishes this ablation evidence from transfer and generation results, which address different questions.Source: §4.7 · p. 9
4.8 LIMITATIONS
- Better reconstruction and frozen semantics did not produce better generation, making latent distribution geometry a remaining constraint on compact tokenizers.Source: §4.8 · p. 9
- Semantic evidence is global and classification-based; causal token interventions and dense prediction are proposed to establish what individual positions encode.Source: §4.8 · p. 9
- Capacity is counted in continuous scalars rather than entropy-coded bits, so ordered-prefix decoding is not yet a variable-rate codec on a rate–distortion frontier.Source: §4.8 · p. 9
5 CONCLUSION
WithinVAE learns semantic targets from the generative system that reconstructs the image, without a pretrained recognition encoder. Its compact latents improve reconstruction and frozen semantic evaluation, while position-dependent corruption orders information along prefixes.
5 CONCLUSION
- WithinVAE is a compact visual autoencoder that learns semantic targets from the same generative system that must reconstruct the image.Source: §5 · p. 9
- WithinVAE improves all four reconstruction metrics over the closest representation-pivoted baseline at half its scalar capacity.At one quarter of that capacity, its 128-token prefix still exceeds the baseline in PSNR and SSIM.Source: §5 · p. 10
- An EMA copy of the autoencoder supplies semantic targets, raising frozen linear probing by 5.1 points over a reconstruction-only baseline while reconstruction improves.Source: §5 · p. 10
- Position-dependent token corruption orders information along the sequence, allowing one trained model to decode variable-length prefixes without retraining.Source: §5 · p. 10