PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexTok: Resampling Images into 1D Token Sequences of Flexible Length

Roman Bachmann, Jesse Allardice, David Mizrahi, Enrico Fini, Oğuzhan Fatih Kar, Elmira Amirloo, Alaaeldin El-Nouby, Amir Zamir, Afshin Dehghan

arXiv:2502.13967v2cs.CVcs.LG

TL;DR

Fixed-length tokenizers do not adapt representation size to image or conditioning complexity. FlexTok introduces ordered, variable-length 1D token sequences with a rectified-flow decoder and nested dropout, achieving plausible reconstructions across lengths and coarse-to-fine autoregressive generation. Its experiments suggest generation can stop once a condition is fulfilled, while dense structured content requires more tokens.

  • Problem

    Existing 2D and 1D tokenizers use fixed-length representations, despite images and generation conditions having differing complexity.

  • Method

    FlexTok resamples images into ordered, discrete 1D token sequences and trains a rectified-flow decoder with nested dropout for variable-length reconstruction and autoregressive generation.

  • Results

    FlexTok supports high-fidelity reconstruction with few tokens and produces coarse-to-fine autoregressive generations whose specificity and conditioning alignment increase with token count.

  • Takeaways & Limitations

    Generation may stop early when a condition is fulfilled, whereas highly dense or structured content such as text requires more tokens.

Abstract

from arXiv · show

Image tokenization has enabled major advances in autoregressive image generation by providing compressed, discrete representations that are more efficient to process than raw pixels. While traditional approaches use 2D grid tokenization, recent methods like TiTok have shown that 1D tokenization can achieve high generation quality by eliminating grid redundancies. However, these methods typically use a fixed number of tokens and thus cannot adapt to an image's inherent complexity. We introduce FlexTok, a tokenizer that projects 2D images into variable-length, ordered 1D token sequences. For example, a 256x256 image can be resampled into anywhere from 1 to 256 discrete tokens, hierarchically and semantically compressing its information. By training a rectified flow model as the decoder and using nested dropout, FlexTok produces plausible reconstructions regardless of the chosen token sequence length. We evaluate our approach in an autoregressive generation setting using a simple GPT-style Transformer. On ImageNet, this approach achieves an FID<2 across 8 to 128 tokens, outperforming TiTok and matching state-of-the-art methods with far fewer tokens. We further extend the model to support to text-conditioned image generation and examine how FlexTok relates to traditional 2D tokenization. A key finding is that FlexTok enables next-token prediction to describe images in a coarse-to-fine "visual vocabulary", and that the number of tokens to generate depends on the complexity of the generation task.

1. Introduction

FlexTok addresses the fixed-length limitation of image tokenization with variable-length, ordered 1D sequences that encode images from coarse semantics to fine details. A rectified-flow decoder supports plausible reconstruction across token lengths, while autoregressive models use the ordering for coarse-to-fine generation.

  • Traditional 2D and recent 1D tokenizers use fixed-length representations regardless of image complexity.
  • FlexTok encodes images into ordered 1D sequences that can be truncated from 1 to 256 tokens and decoded into plausible images.Earlier tokens capture high-level semantic and geometric information, while later tokens add finer details.
  • Nested dropout and causal attention order tokens from coarse concepts to progressively finer details.This ordering supports high-quality reconstruction with few tokens and increasingly detailed representations as token count grows.
  • FlexTok’s visual vocabulary places semantic and geometric concepts before details, allowing generation length to depend on conditioning complexity.ImageNet-1k classes may require as few as 8 tokens, whereas complex text conditions can benefit from up to 256 tokens.
  • A rectified-flow decoder is used to maintain reconstruction quality across varying token lengths, including extreme compression rates.

2. Related Work and Background

Image tokenization compresses images into discrete latent representations for generative modeling, typically preserving a 2D spatial grid. Prior 1D approaches reduce spatial redundancy but generally retain fixed sequence lengths, motivating FlexTok’s adaptive design.

  • Tokenizers make generation more tractable by compressing images into compact, discrete latent representations.Compression removes imperceptible details, while discretization enables per-token categorical distributions.
  • VQ-VAE frameworks combine an image encoder, a quantizer mapping continuous latents to discrete codebook entries, and a decoder.
  • Finite scalar quantization replaces learned codebooks with low-dimensional projections and fixed bins while maintaining reconstruction quality.
  • TiTok replaces the 2D latent grid with a compact 1D sequence of learned register tokens and reconstructs images using masked patch tokens.
  • TiTok uses fixed token counts and two-stage training, while FlexTok supports one token and reports superior reconstruction and generation quality relative to TiTok.ElasticTok and ALIT support minimum lengths of 256 and 32 tokens, respectively.

3. Method

FlexTok combines a discrete 1D register bottleneck with nested dropout, causal attention, and a rectified-flow decoder. These components produce ordered variable-length representations suited to autoregressive coarse-to-fine generation.

  • FlexTok resamples 2D VAE latents into discrete 1D register tokens and conditions a rectified-flow decoder for reconstruction.Causal attention and nested dropout induce ordering, enabling decoding from nested token subsets.
  • The encoder uses learnable registers as read-write storage, discards encoded patches, and quantizes the registers into the bottleneck representation.
  • The rectified-flow decoder predicts a flow from partially noised VAE latents conditioned on quantized register tokens.The training objective minimizes the squared error between predicted flow and the noise-minus-clean-latent target.
  • Nested dropout randomly removes suffix tokens so the encoder and decoder learn ordered representations that reconstruct images from variable-length sequences.The ordering emerges from training rather than being handcrafted; simple images need fewer tokens than complex ones.
  • Causal attention restricts each register to attend only to earlier or equal registers, aligning the representation with left-to-right next-token prediction.
  • Autoregressive Transformers evaluate FlexTok through class-conditional ImageNet-1k generation and text-to-image generation on DFN-2B.

4. Implementation

The implementation uses a staged pipeline: continuous VAE compression, flexible-length FlexTok tokenization, and autoregressive generation. Experiments examine reconstruction tradeoffs and conditional generation using class- and text-conditioned Transformers.

  • The pipeline trains a VAE, resamples its 2D latent grids into flexible discrete 1D sequences, and then trains autoregressive generation models.
  • Stage 0 trains SDXL-style VAEs with 4, 8, and 16 channels, using the 16-channel, 8×-downsampled VAE thereafter.
  • Stage 1 uses up to 256 register tokens, a 6-dimensional FSQ bottleneck, nested dropout, rectified flow, and REPA conditioning at 256×256 resolution.The FSQ levels [8, 8, 8, 5, 5, 5] give an effective vocabulary size of 64 000.
  • Stage 2 uses learned absolute positional embeddings because FlexTok tokens lack a 2D grid structure.
  • Class-conditioned models add a learned class embedding to an [SOI] token and scale from 49M to 1.3B parameters.
  • Reconstruction comparisons show FlexTok supports semantic, ordered tokenization down to a single token, while rate-distortion results improve with more tokens.Increasing tokenizer size significantly improves reconstruction FID, but has less effect on MAE and DreamSim.
  • Text-conditioned models cross-attend to FLAN-T5-XL embeddings and scale to 3B parameters using μP.

5. Experiments

FlexTok supports reconstruction and conditional generation across variable token lengths, with tokens ordered from coarse semantic content to finer details. Experiments show that token requirements depend on conditioning complexity, while generation quality remains consistent across lengths.

  • Flexible-length tokenization: FlexTok produces plausible reconstructions from nested token sequences ranging from 1 to 256 tokens.Early tokens capture high-level semantic features, while additional tokens improve fine-grained alignment and reconstruction metrics.
  • Flexible-length tokenization: Class- and text-conditional generations become increasingly specific to their conditioning as more tokens are produced.DINOv2-L classification accuracy plateaus around 32 tokens, whereas text-image alignment continues improving with additional tokens.
  • Flexible-length tokenization: Simple prompts can achieve satisfactory results with 4 to 16 tokens, whereas detailed prompts benefit from the full 256-token sequence.Lower token counts produce greater variation across random decoder seeds, especially for detailed prompts.
  • Autoregressive model scaling: Long sequences exceeding 128 tokens become more challenging as autoregressive model size decreases, while the first 1-8 tokens are effectively independent of model size.Performance scales strongly with model size for long sequences in both gFID and CLIPScore.
  • Comparison with baselines: FlexTok achieves superior reconstruction and generation quality to previous 1D approaches at each token budget using a single model.The comparison is performed on ImageNet-1k class-conditional generation.

6. Discussion & Conclusion

FlexTok demonstrates flexible-length image reconstruction and coarse-to-fine generation through ordered 1D token sequences. The paper argues that adaptive token counts can support compact semantic representations or longer detailed representations depending on task complexity.

  • Discussion & Conclusion: FlexTok token sequences form a visual vocabulary that enables coarse-to-fine image generation.The approach supports high-fidelity reconstructions with very few tokens and progressively more detailed representations.
  • Discussion & Conclusion: Token counts can be adapted to generation-task complexity, with early stopping proposed when the conditioning is fulfilled.The paper notes that dense or structured content such as text requires more tokens and objectives emphasizing semantically meaningful concepts.
  • Future directions: FlexTok-like adaptive tokenizers may extend to redundant domains such as audio and video.The paper identifies long-horizon video generation, understanding, and visual reasoning as potential directions.
  • Design space: The design space includes VAE choice, token-ordering structure, decoder loss formulation, and inductive-bias losses.These components are investigated to obtain a compact, high-quality variable-length 1D tokenizer.
  • VAE choice ablation: A 16-channel VAE is selected because increasing latent channels above 4 significantly improves FlexTok reconstruction quality.The experiments report a strong correlation between VAE reconstruction quality and latent-channel count.

A.2. Resampling strategy ablation

Ablations show that rectified flow decoding, 1D tokenization, REPA, and ordering mechanisms each affect reconstruction and generation quality. The final design balances performance across token lengths while accounting for noise-schedule stability.

  • Resampling and decoder: Rectified flow decoding significantly improves rFID and gFID but worsens MAE.This establishes a perceptual-quality benefit with a trade-off in pixelwise error.
  • Resampling and decoder: Switching from 2D grid to 1D tokenization improves MAE, DreamSim, rFID, and gFID, with the largest improvement in rFID.These comparisons are reported in Table 3.
  • Inductive-bias losses: REPA significantly improves DreamSim, rFID, and gFID while accelerating convergence to high-quality reconstructions.Visual comparisons also report higher-fidelity and more semantic reconstructions with REPA.
  • Ordering strategies: A causal attention mask improves rFID and gFID over unstructured registers but does not by itself enable flexible-length tokenization.Nested dropout is introduced to support flexible-length use.
  • Ordering strategies: Pow2 provides a more balanced nested-dropout schedule across sequence lengths, while Unifpow2 is preferable for higher sequence lengths.The comparison is evaluated on ImageNet-1k validation data.

B. Evaluating the Representation Quality of FlexTok Tokenizers

Frozen-token linear probing shows that FlexTok’s quantized register-token representations are linearly separable. Classification performance generally improves with more tokens and exceeds the reported TiTok peak.

  • Representation quality: FlexTok’s quantized register-token representations are linearly separable despite no such guarantee from the tokenizer objective.The evaluation trains a linear classifier while keeping the tokenizer frozen.
  • Representation quality: Increasing the number of register tokens consistently improves linear-probing performance.The unified model retains earlier tokens unchanged, allowing additional tokens to provide more information.
  • Representation quality: 64.6% top-1 accuracy is achieved by the best FlexTok configuration on ImageNet.This peak performance significantly surpasses TiTok.
  • Representation quality: A larger encoder improves linear separability when all register tokens are activated, while a larger decoder slightly degrades it.These trends are reported across encoder and decoder sizes.

C. Implementation and Training details

FlexTok combines a Transformer-based 1D resampler, FSQ quantization, causal register attention, nested dropout, and a rectified-flow decoder. The implementation is evaluated across tokenizer sizes, datasets, reconstruction metrics, and inference settings.

  • VAE training: VAE models vary latent channel dimensions from 4 to 16 and use reconstruction, perceptual, KL, and Patch-GAN discriminator losses.They are trained on DFN with spatial downsampling factor 8 and evaluated on COCO using mean squared error and rFID.
  • VAE results: The 16-channel VAE achieves an rFID of 0.354, establishing a reconstruction-quality upper bound for subsequent tokenizer and autoregressive models.Increasing latent channel dimension significantly improves VAE reconstruction performance.
  • Tokenizer architecture: FlexTok uses Transformer encoders and decoders with hidden dimension w = 64 · d, d attention heads, FSQ quantization, and causal register masking.Patch tokens attend to each other; register token i attends only to register tokens j ≤ i.
  • Decoder: Rectified-flow decoders condition on encoded register tokens and reconstruct images from noised VAE latent patches.The decoder uses classifier-free guidance through randomly replacing encoded registers with a learned null condition.
  • Training and evaluation: Resamplers are trained on ImageNet-1k for class-conditional generation or DFN-2B for text-to-image generation, with 200B training tokens and 256 tokens per 256² image.Training uses random crops, horizontal flips, warmup, cosine decay, and AdamW; hyperparameters are not swept for every resampler setting.
  • Inference sweeps: More than 25 denoising steps yield diminishing reconstruction returns, while adaptive projected guidance produces shallower, better-aligned guidance basins than standard guidance.Standard guidance has narrow optima dependent on token count; APG improves overall rFID values.

F.4. 1D vs 2D tokenizer classifier-free guidance ablation

The 1D FlexTok tokenizer is compared with a controlled 2D grid tokenizer using APG guidance sweeps on COCO. FlexTok performs better on perceptual metrics across guidance scales, while the 2D baseline retains an MAE advantage.

  • Results: FlexTok outperforms the 2D baseline by a large margin in rFID and DreamSim across guidance scales.The 2D grid tokenizer performs better in MAE.
  • Experimental setup: The comparison sweeps APG guidance scales for FlexTok d18-d28 and a controlled 2D tokenizer baseline, with FlexTok evaluated using all 256 tokens.Both models are trained on DFN-2B and compared on the COCO 30.5k validation set.
  • Inference configuration: The selected inference configurations for ImageNet-1k and DFN-2B tokenizers are recorded in Tables 10 and 11.These settings are used for subsequent image-generation experiments.

G. Autoregressive Class-Conditional Image Generation Hyperparameter Sweeps

Class-conditional autoregressive experiments examine classifier-free guidance, sampling, token count, model scale, and tokenizer scale. Guidance is unnecessary for the tested FlexTok model, while larger models and tokenizers mainly improve longer-sequence generation.

  • Hyperparameter sweeps: For a 1.33B AR model with FlexTok d18-d28, no classifier-free guidance, CFG scale s = 1.0, gives the best gFID regardless of top-k sampling or generated token count.This result is shown in the top-k and token-count ablations.
  • AR model scaling: Generation quality with few tokens is largely independent of AR model size, whereas longer sequences benefit increasingly from larger AR models.Training loss consistently improves as AR model size increases.
  • Tokenizer scaling: Scaling the tokenizer from FlexTok d12-d12 to d18-d28 improves gFID with a fixed 1.33B AR model.For some tokenizer sizes, reconstruction quality limits generative-model performance.

I. Autoregressive Text-Conditional Image Generation Inference Hyperparameter Sweeps

Text-conditional experiments evaluate guidance, token count, AR model scale, and FlexTok versus 2D tokenization. CFG improves text-conditioned generation, while token count trades off image quality against text-image alignment.

  • Guidance sweeps: For a 3.06B text-conditional AR model with FlexTok, classifier-free guidance improves gFID at every generated sequence length.This differs from the class-conditional FlexTok setting, where CFG hurt gFID.
  • Token-count sweeps: A CFG scale of 2.5 optimizes gFID across token counts, while CLIPScore improves as more tokens are generated.At this scale, gFID drops sharply before gradually increasing with token count.
  • 2D comparison: The 2D grid tokenizer shows a similar dependence on CFG scale, so 2.5 is selected to balance gFID and CLIPScore in both approaches.Using the same scale enables a balanced comparison.
  • Model scaling: Scaling AR models from 113M to 3.06B lowers final training loss for both FlexTok and 2D tokenizers and significantly improves gFID at 256 tokens.The comparison uses tokenizers trained on DFN and generates 256 tokens.

K.1. FlexTok image reconstruction for different numbers of tokens – multiple samples per token sequence

FlexTok reconstructions become progressively more specific as token sequences lengthen, while early tokens preserve coarse semantic and geometric content. The figures also compare these flexible reconstructions with fixed-length baselines and show class- and text-conditioned generations across token budgets.

  • Multiple samples per token sequence: More tokens narrow the decoder’s image distribution, producing reconstructions increasingly specific to the original image.Five samples from the rectified flow decoder are drawn from the same token sequence at each length.
  • Token counts and model sizes: FlexTok reconstructions are shown across different token counts and model sizes on ImageNet-1k validation samples.The comparison includes d12-d12, d18-d18, and d18-d28 models.
  • Baseline comparison: FlexTok reconstructions are compared with three TiTok models and ALIT across different token counts.The comparison uses ImageNet-1k validation samples and the FlexTok d18-d28 model.
  • Class-conditional generation: Class-conditional generations vary token counts from 1 to 256 using a FlexTok d18-d28 tokenizer and a 1.33B-parameter autoregressive Transformer.The examples include multiple ImageNet-1k class labels and random seeds.
  • Text-conditional generation: Text-conditional generations vary token counts from 1 to 256 using a DFN-trained FlexTok tokenizer and a 3.06B-parameter autoregressive Transformer.The examples use text embeddings from fixed prompts or different PartiPrompts.
Loading 2502.13967v2…