PaperlayerView exact revision on arXiv

Source-linked AI summary

AlignTok: Aligning Visual Foundation Encoders to Tokenizers for Diffusion Models

Bowei Chen, Sai Bi, Hao Tan, He Zhang, Tianyuan Zhang, Zhengqi Li, Yuanjun Xiong, Jianming Zhang, Kai Zhang

arXiv:2509.25162v2cs.CV

TL;DR

Image diffusion tokenizers must provide both diffusion-friendly latent spaces and accurate reconstruction, but VAE training often leaves latents dominated by low-level details. AlignTok aligns a pretrained visual encoder through three stages, and the resulting tokenizer accelerates diffusion convergence and improves generation across ImageNet and LAION settings.

  • Problem

    VAE tokenizer training is dominated by reconstruction, often producing latent spaces dominated by low-level details rather than predictable semantic structure for diffusion.

  • Method

    AlignTok freezes a pretrained encoder while training an adapter and decoder, jointly optimizes all components with semantic preservation, and then refines the decoder.

  • Results

    On ImageNet and LAION, AlignTok accelerates diffusion convergence and improves generation compared with previous methods, including FLUX VAE and VA-VAE.

  • Takeaways & Limitations

    Aligning pretrained visual encoders provides a simple, scalable paradigm for semantically grounded continuous tokenizer design.

  • Takeaways & Limitations

    Reconstruction ability still lags behind FLUX VAE, and evaluation is limited to images up to 512 resolution.

Abstract

from arXiv · show

In this work, we propose aligning pretrained visual encoders to serve as tokenizers for latent diffusion models in image generation. Unlike training a variational autoencoder (VAE) from scratch, which primarily emphasizes low-level details, our approach leverages the rich semantic structure of foundation encoders. We introduce a three-stage alignment strategy called AlignTok: (1) freeze the encoder and train an adapter and a decoder to establish a semantic latent space; (2) jointly optimize all components with an additional semantic preservation loss, enabling the encoder to capture perceptual details while retaining high-level semantics; and (3) refine the decoder for improved reconstruction quality. This alignment yields semantically rich image tokenizers that benefit diffusion models. On ImageNet 256$\times$256, our tokenizer accelerates the convergence of diffusion models, reaching a gFID of 1.90 within just 64 epochs, and improves generation both with and without classifier-free guidance. Scaling to LAION, text-to-image models trained with our tokenizer consistently outperforms FLUX VAE and VA-VAE under the same training steps. Overall, our method is simple, scalable, and establishes a semantically grounded paradigm for continuous tokenizer design.

1 INTRODUCTION

VAE-based tokenizers often prioritize reconstruction, leaving diffusion latents dominated by low-level details and weakly semantic. AlignTok instead aligns a pretrained visual encoder through three stages, producing semantically grounded tokenizers that improve diffusion-model convergence and generation.

  • Motivation: VAE tokenizer training is dominated by reconstruction loss, so its latent space often develops an unpredictable structure dominated by low-level details.The encoder’s diffusion-friendly representation is learned only indirectly and weakly regularized by the KL prior.
  • Motivation: AlignTok aligns a pretrained visual foundation encoder with a tokenizer instead of requiring the encoder to learn semantic structure from scratch.The approach uses the encoder’s existing semantic representations to make diffusion-friendly latent-space learning easier.
  • Method: The three stages freeze the encoder while training an adapter and decoder, jointly optimize all components with semantic preservation, then refine only the decoder.These stages progressively add perceptual detail while retaining semantic structure and improving reconstruction fidelity.
  • Results: On ImageNet, the semantic tokenizer accelerates diffusion convergence and improves generation with and without classifier-free guidance, while LAION models converge faster than FLUX VAE and VA-VAE.The authors attribute these benefits to a more structured semantic latent space.
  • Contribution: The paper presents pretrained-encoder alignment as a simple, scalable paradigm for continuous visual tokenizer design.The proposed paradigm is intended to support future directions in generative modeling.

2 RELATED WORK

Prior tokenizer work uses reconstruction objectives, semantic regularization, or pretrained encoders in discrete-tokenizer settings. This paper targets continuous diffusion tokenizers with a simple autoencoder that directly aligns a pretrained encoder using semantic preservation.

  • Image Tokenizers for Generative Models: Continuous tokenizers map images into compact latent spaces where generative models can operate more efficiently, commonly using reconstruction losses.Typical losses include L1, perceptual, and adversarial losses.
  • Semantic Regularization: Semantic-regularization methods constrain tokenizer latents toward pretrained-encoder representations, whereas AlignTok aligns an encoder already capable of extracting high-level semantics.The paper positions this as an alternative to learning semantic structure through regularization.
  • Pretrained Encoders as Discrete Tokenizers: Previous pretrained-encoder approaches mainly study discrete tokenizers, using fixed features, architectural additions, or contrastive image–text supervision.The paper notes that contrastive approaches primarily assume language-aligned encoders such as SigLIP 2.
  • Positioning: AlignTok focuses on continuous diffusion tokenizers with a simple autoencoder and self-supervised semantic preservation, without extra architectural design or image–text supervision.The method is presented as generalizable to visual encoders and scalable for generative modeling.
  • Method Overview: The method progressively adapts a pretrained encoder through latent alignment, perceptual alignment, and decoder refinement.These stages separately establish semantic latents, add perceptual detail while preserving semantics, and improve reconstruction.

3 METHOD

The method aligns a pretrained visual encoder with a diffusion-compatible tokenizer through staged latent alignment, perceptual alignment, and decoder refinement. It preserves semantic structure while improving reconstruction fidelity, with the full three-stage model identified as the best reconstruction–generation balance in the ablation study.

  • 3 METHOD: Three alignment stages progressively adapt a pretrained encoder into a diffusion-friendly tokenizer while preserving semantic structure and improving reconstruction.The stages comprise latent alignment, perceptual alignment with semantic preservation, and decoder-only refinement.
  • 3 METHOD: A frozen pretrained encoder feeds a lightweight adapter that compresses high-dimensional features into a 32-channel latent code, which a decoder reconstructs into the input image.Only the adapter and decoder are trained during this first stage, while the encoder remains frozen.
  • 3 METHOD: Jointly optimizing the encoder, adapter, and decoder improves reconstruction but can sharply reduce semantic probing accuracy, motivating an L2 loss that keeps current latent codes close to the previous stage.The semantic preservation term constrains the evolving latent space while allowing the encoder to capture fine-grained perceptual details.
  • 3 METHOD: The preservation strategy maintains a semantically rich latent space while achieving comparable reconstruction performance.The paper reports this outcome for the blue curve in Figure 3.
  • 3 METHOD: The full three-stage model achieves the best balance between reconstruction and generation quality in the ImageNet 256×256 ablation at 80K training steps.The ablation uses 30 sampling steps and the CFG scale producing the lowest generation FID.
  • 3 METHOD: Decoder-only refinement improves reconstruction quality after the latent space is fixed, allowing the decoder to better exploit the existing representation without disturbing its semantics.The decoder is refined because earlier stages continually changed the latent representation and could leave it underfit.

4 EXPERIMENTS

Experiments show that AlignTok’s alignment design balances semantic structure, reconstruction, and generation, while improving diffusion efficiency across ImageNet settings and scaling to text-to-image generation.

  • Semantic Preservation Loss: A moderate semantic-preservation weight of 1 provides the best balance between generation and reconstruction in the Stage 1 + Stage 2 variant.Lower weights allow semantic drift, whereas higher weights improve generation but harm reconstruction fidelity.
  • High-Level Design: The three-stage design improves reconstruction and slightly improves generation, while stage 1 alone preserves semantics but reconstructs poorly.Stage 2 substantially improves both reconstruction and generation, and stage 3 further strengthens reconstruction.
  • Other Pretrained Encoders: DINOv2 provides the best balance of generation quality and reconstruction among the tested pretrained encoders.MAE achieves stronger reconstruction fidelity but the weakest generation performance.
  • Sampling Step: 50-step generations surpass VA-VAE outputs at 250 steps, while near-optimal performance is reached with 80 steps.The authors attribute this robustness to a smoother latent space that reduces the impact of discretization errors.
  • CFG Scale: Across CFG scales, the tokenizer consistently outperforms VA-VAE and maintains strong performance at low guidance values.The authors associate this reduced guidance dependence with well-separated class semantics in the latent space.
  • Convergence Speed: 5x faster training reaches comparable quality in ∼60K steps versus VA-VAE’s 300K steps with 30 sampling steps.The aligned tokenizer consistently achieves better gFID across training iterations.
  • Comparison with Other Tokenizers: The tokenizer outperforms VA-VAE across channel dimensions and remains stronger in unconditional generation, though reconstruction lags VA-VAE at 64 channels.The unconditional comparison uses models trained for only 80K steps, so performance may not be fully optimized.

5 LIMITATIONS AND DISCUSSIONS

The method remains limited by reconstruction quality and evaluation resolution, while suggesting broader tokenizer applications as future work.

  • Limitations: Reconstruction ability still lags behind FLUX VAE, although stronger decoders, larger channels, longer training, or greater scale may narrow the gap.The authors identify these as possible ways to improve reconstruction.
  • Limitations: Evaluation is limited to images up to 512 resolution, leaving higher-resolution exploration for future work.The authors mention DINOv3 as a possible resource for variable-resolution extensions.
  • Discussion: The alignment approach may extend beyond image diffusion tokenizers to video, discrete autoregressive tokenizers, and unified multimodal representations.These extensions are presented as promising future directions rather than evaluated results.

A IMPLEMENTATION DETAILS

The implementation uses pretrained DINOv2 encoders, lightweight adapters, convolutional decoders, and staged training across ImageNet and LAION experiments.

  • Tokenizer: ImageNet tokenizers use a DINOv2 ViT-L/14 encoder, a two-layer adapter, 32 latent channels, and a convolutional decoder.The encoder has approximately 304M parameters, while the decoder has approximately 42M.
  • Tokenizer: LAION tokenizer training uses a DINOv2 encoder, batch size 256, learning rate 1e-4, 32 H100 GPUs, and non-square images preserving aspect ratios.Other settings follow the ImageNet experiments.
  • Generative Models: ImageNet diffusion experiments use LightningDiT with 673M parameters, batch size 1024, learning rate 2e-4, and 80K steps or 64 epochs.Training takes approximately 12 hours on 8 H100 GPUs.
  • Generative Models: LAION text-to-image experiments use a 2B-parameter diffusion transformer, beginning at 256 resolution before continuing at 512 resolution.The initial phase runs for 200K steps and the continuation for an additional 90K steps.
  • Evaluation: Linear probing freezes the VAE encoder and adapter, then trains a single linear classifier to measure linearly accessible semantic information.The classifier predicts ImageNet-1K classes from adapter outputs.

B TRAINING AND INFERENCE COST

The staged pipeline has higher cost when the encoder is trainable, but its cumulative training cost is lower than both comparison tokenizers and its high-resolution inference memory can be substantially lower than VA-VAE.

  • Training Cost: Stage 2 consumes the most training memory because it fine-tunes DINOv2, while Stages 1 and 3 use less memory than VA-VAE because their encoder is frozen.Vanilla VAE uses the least memory because it has no pretrained encoder.
  • Training Cost: 576.15 cumulative A100 GPU hours are lower for the proposed method than for both Vanilla VAE and VA-VAE.The total is computed as 107.2 + 377.4 + 91.55 GPU hours across the stages.
  • Inference Cost: At 1024 resolution with batch size 8, peak memory is 3.015 GB for the proposed encoder versus 22.78 GB for VA-VAE.At 256 resolution with batch size 4, the proposed encoder instead uses 1.372 GB versus 0.838 GB for VA-VAE.
  • Inference Cost: Encoding latency and GFLOPS are generally higher than VA-VAE, except at 512 resolution where the proposed tokenizer has lower encoding latency.The comparison uses resolution-dependent inference measurements.

C MORE EXPERIMENTS

Among the evaluated pretrained encoders, DINOv2 provides the strongest balance between generation quality and reconstruction performance.

  • Pretrained Encoder Comparison: DINOv2 achieves the best balance, combining superior generation quality with competitive reconstruction performance.The comparison is reported in the additional reconstruction metrics of Table 9.

C.2 RECONSTRUCTION

The method improves reconstruction through encoder alignment and decoder refinement, while retaining stronger generation performance than competing tokenizers. A variant can reach competitive reconstruction quality by emphasizing perceptual details during stage 2.

  • C.2 RECONSTRUCTION: The model without stages 2 and 3 fails to reconstruct inputs accurately, whereas the other methods show comparable qualitative reconstruction quality.The failure is attributed to insufficient perceptual detail in the frozen pretrained encoder’s latent space.
  • C.2 RECONSTRUCTION: Our tokenizer outperforms all baselines in generative performance despite being quantitatively weaker in reconstruction quality.The comparison is evaluated on ImageNet 256×256 at 80K training steps with 30 sampling steps.
  • C.2 RECONSTRUCTION: Our method converges faster than VA-VAE on ImageNet, indicating that semantic alignment and preservation are more effective than semantic regularization.The comparison is shown across additional qualitative convergence results.
  • C.2 RECONSTRUCTION: Our method converges significantly faster than FLUX VAE, demonstrating the benefit of its learned semantically rich latent space.The comparison uses qualitative convergence results for text-to-image generation.

C.4 MORE QUANTITATIVE RESULTS

Across additional prompt sets and evaluation settings, the proposed tokenizer consistently achieves better generation quality than FLUX VAE and VA-VAE. This advantage holds both with and without classifier-free guidance.

  • C.4 MORE QUANTITATIVE RESULTS: The additional text-to-image evaluations compare models trained with FLUX VAE or the proposed tokenizer on two prompt sets at 256×256 resolution.Each model is trained for 100K steps in the reported FLUX VAE comparison.
  • C.4 MORE QUANTITATIVE RESULTS: The no-CFG FLUX VAE comparison evaluates 2B-parameter text-to-image models trained for 100K steps at 256×256 resolution.The tokenizer used for the proposed model is trained on LAION, and rFID uses 200K COYO-700M samples.
  • C.4 MORE QUANTITATIVE RESULTS: Across all configurations, our tokenizer consistently outperforms FLUX VAE and VA-VAE in generation quality with or without CFG.The comparisons cover additional prompt sets, GenEval, and settings with and without classifier-free guidance.

C.5 MORE QUALITATIVE RESULTS

Additional qualitative experiments show the proposed tokenizer across class-conditional and text-to-image generation settings, including multiple resolutions and aspect ratios. The evaluations also compare text-to-image generation against VA-VAE under CFG.

  • C.5 MORE QUALITATIVE RESULTS: The method is demonstrated on text-to-image generation at 256×256 after 200K training steps and at 512×512 after 290K training steps.The 512×512 results include multiple qualitative figures.
  • C.5 MORE QUALITATIVE RESULTS: The qualitative comparisons evaluate text-to-image generation with VA-VAE using 1B-parameter models trained for 50K steps with CFG scale 5 at 256×256 resolution.Both the proposed tokenizer and VA-VAE are trained on ImageNet.

C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE

The proposed latent space more closely matches DINOv2’s semantic structure than the compared tokenizer spaces, while the method still exhibits specific text-to-image failure cases. Qualitative convergence and generation comparisons further characterize its behavior.

  • C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE: Our latent space most closely resembles DINOv2 features, whereas Vanilla VAE can be overly smooth or sharp and VA-VAE is oversmoothed.The PCA comparison treats DINOv2 outputs as a semantic reference.
  • C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE: CKNNA and Total Variation place our latent space closer to DINOv2 than other tokenizers, while VA-VAE is closest on Gini Coefficient.The metrics capture different aspects of latent distributions, so they do not agree on one tokenizer across all measures.
  • C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE: The method’s failure cases include inaccurate clock numerals, incorrect object counts, inconsistent long text, and difficulty rendering fine details such as hands.These failures are reported for text-to-image generation at 512×512 resolution.
  • C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE: Without stages 2 and 3, the pretrained encoder lacks sufficient perceptual detail for accurate reconstruction, while the other methods are qualitatively comparable.This comparison concerns ImageNet 256×256 reconstruction quality.
  • C.6 ANALYSIS OF THE SEMANTICS OF THE LATENT SPACE: The convergence comparisons evaluate the method against VA-VAE on ImageNet and against FLUX VAE on text-to-image generation under specified CFG and sampling settings.The reported figures use best CFG for ImageNet comparisons and CFG 5 for text-to-image comparisons.
Loading 2509.25162v2…