PaperlayerView exact revision on arXiv

Source-linked AI summary

CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization

Yitong Chen, Zuxuan Wu, Xipeng Qiu, Yu-Gang Jiang

arXiv:2603.06449v1cs.CV

TL;DR

Visual tokenizers lack a causal ordering that naturally supports autoregressive next-token prediction, and existing diffusion autoencoders either expose all tokens or bias toward early ones. CaTok uses interval-conditioned MeanFlow decoding with REPA-A regularization to learn causal 1D image tokens. It reports state-of-the-art ImageNet reconstruction metrics and generation performance comparable to leading approaches, while evaluation remains focused on ImageNet-1K and does not cover larger datasets or broader tasks.

  • Problem

    Visual tokenizers lack a natural causal order, while all-token conditioning removes causality and nested dropout biases representations toward early tokens.

  • Method

    CaTok selects 1D tokens within sampled time intervals for MeanFlow decoding and uses REPA-A to align encoder features with external visual representations.

  • Results

    0.75 rFID, 22.53 PSNR, and 0.674 SSIM are reported for ImageNet reconstruction, while generation reaches comparable performance with 2.95 gFID.

  • Takeaways & Limitations

    CaTok provides causal 1D visual representations that support one-step and multi-step sampling and are suitable for standard autoregressive modeling.

  • Takeaways & Limitations

    The study focuses on ImageNet-1K and leaves larger datasets and broader task evaluations for future work.

Abstract

from arXiv · show

Autoregressive (AR) language models rely on causal tokenization, but extending this paradigm to vision remains non-trivial. Current visual tokenizers either flatten 2D patches into non-causal sequences or enforce heuristic orderings that misalign with the "next-token prediction" pattern. Recent diffusion autoencoders similarly fall short: conditioning the decoder on all tokens lacks causality, while applying nested dropout mechanism introduces imbalance. To address these challenges, we present CaTok, a 1D causal image tokenizer with a MeanFlow decoder. By selecting tokens over time intervals and binding them to the MeanFlow objective, as illustrated in Fig. 1, CaTok learns causal 1D representations that support both fast one-step generation and high-fidelity multi-step sampling, while naturally capturing diverse visual concepts across token intervals. To further stabilize and accelerate training, we propose a straightforward regularization REPA-A, which aligns encoder features with Vision Foundation Models (VFMs). Experiments demonstrate that CaTok achieves state-of-the-art results on ImageNet reconstruction, reaching 0.75 FID, 22.53 PSNR and 0.674 SSIM with fewer training epochs, and the AR model attains performance comparable to leading approaches.

25 Steps

CaTok’s reconstruction samples compare one-step and multi-step MeanFlow decoding, token reduction, and separate token intervals to illustrate efficiency, fidelity, causality, and semantic diversity.

  • Sampling: One-step and multi-step sampling are compared using 256 tokens.The figure places fast one-step sampling beside high-quality multi-step sampling.
  • Causality: Reducing tokens from 256 to 16 produces a fine-to-coarse reconstruction trend.This progression is presented as evidence that the 1D tokens are causal.
  • Concept decoding: Different 16-token segments reconstruct diverse visual concepts across token intervals.The segment-based reconstructions indicate that visual concepts are distributed across intervals.

1 Introduction

CaTok addresses the difficulty of imposing causal structure on visual tokens by combining interval-conditioned MeanFlow decoding with a causal 1D tokenizer and REPA-A regularization. On ImageNet, it reports state-of-the-art reconstruction metrics and comparable class-conditional generation performance.

  • Motivation: Image tokenization lacks a natural causal order, while flattened 2D sequences and heuristic orderings do not directly match next-token prediction.This motivates learning causal 1D visual representations rather than imposing an ordering after tokenization.
  • Motivation: Naïve flow decoders condition on all tokens, whereas nested-dropout consistency decoders bias conditioning toward early tokens and introduce imbalance.Both designs create limitations for causal autoregressive learning.
  • Method: CaTok selects tokens within a sampled interval [r, t] and binds them to the corresponding MeanFlow interval.This models the average velocity along a subpath, capturing causality while supporting one-step sampling.
  • Method: REPA-A aligns encoder image features with high-quality external visual representations to stabilize and accelerate autoencoder training.The method is presented as a regularization technique inspired by REPA and REPA-E.
  • Results: 0.75 rFID, 22.53 PSNR, and 0.674 SSIM are reported for CaTok-L on ImageNet reconstruction, with 2.95 gFID for comparable generation performance.The reported reconstruction metrics are state-of-the-art, while generation performance is described as comparable to leading approaches.

2 Background

The background introduces rectified-flow paths and velocity fields, then motivates MeanFlow as direct average-velocity modeling over time intervals to reduce few-step sampling errors.

  • Rectified flows: Rectified flows connect data x and prior noise ϵ through the path z_t = (1 − t)x + tϵ.A neural network learns the corresponding marginal velocity field.
  • Rectified flows: Inference begins from prior noise at z_1 = ϵ and generates samples by solving the flow dynamics.The integral is approximated numerically in discrete time steps.
  • Rectified flows: Euler updates estimate the average velocity over [r, t] using only the instantaneous velocity at t.The passage identifies this approximation as a source of sampling inaccuracies.
  • MeanFlow: MeanFlow directly fits the average velocity u over the interval [r, t] to mitigate errors from using fewer sampling steps.Its formulation derives average velocity from the instantaneous velocity field.
  • MeanFlow: The MeanFlow objective uses stop-gradient operations, and one-step sampling can be performed as z_0 = ϵ − u_θ(ϵ, 0, 1).The stop-gradient avoids double backpropagation through the Jacobian–vector product.

3 CaTok

CaTok combines a causal ViT encoder with a MeanFlow DiT decoder, training token-conditioned flow objectives and representation alignment to support causal 1D image tokenization and autoregressive generation.

  • 3.1 Architecture: CaTok uses a causal ViT encoder and MeanFlow DiT decoder to produce and decode compressed 1D visual tokens.The encoder processes images with registers, while the decoder models image generation from token-conditioned flow states.
  • 3.1 Architecture: Causal attention lets image features attend to each other, while each 1D token attends to all image features and only preceding 1D tokens.This mask enforces the dependency structure required for causal tokenization.
  • 3.2 Training: The decoder samples r and t, conditions on tokens indexed by [r · K, t · K], and predicts average velocity over that interval.Instantaneous velocity modeling with r = t is also used to improve training stability.
  • 3.2 Training: CaTok jointly optimizes MeanFlow, Rectified Flow, REPA, and REPA-A objectives.The flow objectives model average and instantaneous velocity, while representation objectives align learned features with pretrained visual representations.
  • 3.3 Autoregressive modeling: After encoder freezing, an autoregressive model predicts continuous 1D tokens with teacher forcing and diffusion loss, which the MeanFlow decoder can sample in one step.Generation starts from a learned class token and feeds the predicted token sequence into the decoder.
  • 3.2 Training: REPA-A aligns encoder image features with Vision Foundation Model representations, producing more informative tokens and accelerating convergence.The paper reports that this alignment also enhances overall performance.

4 Related works

Visual tokenizers span 2D patch-based methods, 1D tokenizers, and diffusion autoencoders, with differing strategies for compression, reconstruction, and causality.

  • 2D visual tokenizers: 2D visual tokenizers apply vector quantization to image patches, with later methods adding adversarial losses, multiple quantization stages, or lookup-free quantization.Examples include VQ-VAE, VQGAN, RQ-VAE, and MAGVIT-v2.
  • 1D visual tokenizers: 1D visual tokenizers extract compact token sequences using causal Q-Formers, learnable registers, or ViT decoders for semantic or mask-to-patch reconstruction.Subsequent work explores textual conditioning and adaptive-length tokenization.
  • Diffusion autoencoders as 1D tokenizers: Diffusion autoencoders use 1D tokens as conditioning for denoising or rectified-flow decoders, but naïve designs condition on all tokens simultaneously and lose causality.Consistency decoders restore causality by conditioning on early tokens through nested dropout.

5 Experiments

Experiments evaluate CaTok on ImageNet reconstruction, autoregressive generation, and design ablations. Results show strong reconstruction quality, competitive AR generation with less tokenization training, and benefits from causal interval selection and REPA-A.

  • Reconstruction: Reconstruction is evaluated on ImageNet-1K at 256 × 256 using FID, PSNR, and SSIM across CaTok variants with 32–256 tokens.Comparisons use state-of-the-art methods with comparable latent spaces and model sizes.
  • Reconstruction: CaTok-L-256† achieves the best PSNR and SSIM among one-step 1D tokenizers while supporting both one-step and improved multi-step reconstruction.Its rFID surpasses VQGAN by three points but remains behind modern tokenizers.
  • Autoregressive generation: CaTok attains comparable gFID and IS to state-of-the-art tokenizers after 400 AR-training epochs, using 160 tokenization-training epochs versus 300+.The evaluation measures image quality and class diversity on ImageNet-1K.
  • Scope: The AR-modeling study excludes broader tokenizer–generator interactions, larger datasets, and wider task evaluations because of computational cost and scope.The paper specifically notes token dimension and accumulating errors from longer token sequences as beyond scope.
  • Ablation study: MeanFlow enables reasonable one-step results, while REPA and REPA-A accelerate convergence and improve performance; interval-selected tokens add causality with a slight reconstruction trade-off.The ablation roadmap compares the conventional diffusion autoencoder with successive CaTok design choices.
  • Ablation study: Causal interval selection achieves the best gFID, whereas non-causal tokens hinder AR modeling and token imbalance reduces later-token contributions.CaTok addresses imbalance without an additional re-weighting mechanism.
  • Ablation study: REPA-A makes encoder features more informative and discriminative and mitigates the loss spike when MeanFlow loss is introduced at 25K steps.These effects stabilize decoder training and improve overall performance.

6 Conclusion

CaTok combines causal 1D tokenization with a MeanFlow decoder and REPA-A regularization. It achieves state-of-the-art ImageNet reconstruction metrics and comparable class-conditional generation performance.

  • Conclusion: CaTok binds MeanFlow’s average velocity field to 1D token segments, enabling visual causality along the flow path and one-step sampling.The method is designed to bridge autoregressive language models and vision models.
  • Conclusion: REPA-A stabilizes and accelerates diffusion autoencoder training by regularizing encoder representations.The regularization leverages external visual representations.
  • Conclusion: CaTok achieves state-of-the-art PSNR and SSIM on ImageNet reconstruction and comparable class-conditional generation results.The conclusion summarizes outcomes across reconstruction and generation.

7 More implementation details

CaTok-B and CaTok-L use detailed configurations for tokenization and autoregressive modeling, with training and sampling schedules specified for different training durations.

  • Table 4 provides detailed configurations for CaTok-B and CaTok-L tokenization and autoregressive modeling.
  • For reconstruction, one-step sampling disables CFG, whereas 25-step sampling uses CFG with scale 2.0.
  • For 80-epoch training, MeanFlow begins at epoch 10 and token selection at epoch 40; for 160 epochs, they begin at epochs 20 and 80.
  • CaTok generation does not use CFG, while the AR model follows the guidance setup of MUSE, MAR, and Semanticist to improve diversity for small-indexed tokens.

8 More experiments

Additional experiments assess CaTok with VQ, non-autoregressive generation, alternative REPA teachers and latent spaces, higher resolutions, and COCO data, showing broad compatibility and generalization.

  • CaTok-VQ: CaTok with FSQ is compared against LlamaGen without modifying the original training recipe, while Table 5 reports its reconstruction and generation results.
  • Extensions on the REPA encoder, latent space, and image resolution: CaTok shows consistent reconstruction across REPA teachers and latent spaces, while smaller DiT patch sizes alleviate reconstruction drops caused by latent re-compression.
  • Extensions on the REPA encoder, latent space, and image resolution: A training-free patchwise diffuse-and-blend strategy enables the DiT to generalize naturally to higher resolutions.
  • Additional dataset: CaTok achieves consistent performance on COCO-val-5K, indicating that its learned representations generalize across datasets with different distributions.
  • Non-autoregressive generator: CaTok remains compatible with non-autoregressive ϵMaskGiT generation and achieves better performance than TiTok under 8-step sampling.
  • Non-autoregressive generator: A CaTok variant without token dropout also yields improved performance, supporting applicability beyond the default token-dropout setting.
Loading 2603.06449v1…