🤖 AI

Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

This paper introduces a novel variable-length tokenizer for Diffusion Transformers that achieves superior quality-compute trade-offs by employing learnable global merging to enable data-independent, cross-length representation alignment, overcoming the semantic misalignment issues inherent in traditional truncation-based methods.

Original authors: Dong Hoon Lee, Seunghoon Hong

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dong Hoon Lee, Seunghoon Hong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to send a high-definition movie to a friend. If you send the whole thing uncompressed, it takes forever to download and uses a massive amount of data. If you compress it too much, the picture becomes blurry and pixelated.

For a long time, AI image generators (like the ones that create pictures from text) faced this exact problem. They use a "tokenizer"—a translator that breaks an image down into small chunks called tokens.

High compression (few tokens) = Fast and cheap, but the image looks bad.
Low compression (many tokens) = Great quality, but slow and expensive.

Previously, if you wanted a different balance between speed and quality, you had to train a completely different AI model for each setting. It was like needing a different car engine just to drive 10 mph slower.

The Problem with "Cutting"

Some researchers tried to fix this by creating "variable-length" tokenizers. Their idea was simple: "Let's just cut off the end of the token list if we want to save space."

Think of this like a book where the most important plot points are on page 1, and the tiny details are on page 100. If you want a shorter version, you just tear off the last 50 pages.

The Catch: This changes the story. The "short" version focuses only on the big picture, while the "long" version includes the tiny details. Because the content is so different, the AI gets confused. It's like trying to teach a student to read a short story and a long novel using the same textbook, but the chapters are completely different. The AI struggles to learn both at once, leading to blurry or weird images.

The Solution: "Merging" Instead of Cutting

The authors of this paper propose a smarter way: Merging instead of cutting.

Imagine you have a group of 100 people (tokens) standing in a line.

The Old Way (Cutting): If you need fewer people, you just tell the last 50 to go home. The remaining 50 are still standing in the exact same spots.
The New Way (Merging): If you need fewer people, you ask similar-looking people to huddle together and form a single "super-person." If two people are wearing the same red shirt, they merge into one representative of that group.

Why is this better?

Consistency: Whether you have 100 people or 10 "super-people," the groups represent the same underlying structure. The "red shirt group" is still there, just condensed. This keeps the "story" consistent for the AI, no matter how many tokens you use.
Alignment: Because the groups are formed based on similarity (not just position), the AI can learn one single model that works perfectly for both short and long lists.

The Secret Sauce: "Learnable Global Merging"

There was one big hurdle. Usually, to decide who merges with whom, you look at the specific image (e.g., "These two pixels look like a cat's ear, so merge them"). But when an AI is creating a new image from scratch, it doesn't have the image yet! It can't look at the image to decide how to merge.

The authors solved this with Learnable Global Merging.
Think of this as a pre-set rulebook or a "master plan" that the AI learns during training. Instead of looking at the image to decide who merges, the AI uses a fixed, learned pattern (like a universal instruction manual) that says, "Token 1 always merges with Token 2, Token 3 merges with Token 4," regardless of what the final image looks like.

Even though this rulebook is fixed (it doesn't change based on the image), the AI trains it so well that it naturally groups similar things together. This allows the AI to know exactly how to handle the "merged" tokens even before it generates the picture.

The Results

The researchers tested this on ImageNet (a huge database of images). They found that:

Their method allows a single AI model to generate images at different speeds and qualities without needing to be retrained for each setting.
The images look better and the math works out more efficiently than previous methods that just "cut" the tokens.
They even added a tiny "fine-tuning" step (using a technique called LoRA) that makes the AI even better at specific token counts with almost no extra cost.

In short: They figured out how to make AI image generators flexible. Instead of forcing the AI to choose between a "fast" model and a "good" model, they built a single, smart translator that can shrink or expand the image data by intelligently grouping similar parts together, keeping the quality high and the cost low.

Technical Summary: Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

1. Problem Statement

Latent Diffusion Models (LDMs) rely on tokenizers to compress images into latent spaces, creating a trade-off between generation quality and computational cost. Conventional tokenizers use a fixed compression ratio, necessitating multiple model variants to accommodate different computational budgets or quality requirements.

Variable-Length Tokenizers (VLTs) have been proposed to address this by adaptively adjusting the number of tokens based on specific conditions (e.g., computational budget). The prevailing method, nested dropout, trains models to drop tail tokens from an ordered sequence, placing high-level semantics in early tokens and low-level details in later ones.

However, the authors identify a critical flaw in nested dropout when applied to Diffusion Transformers (DiTs):

Cross-Length Shift: Truncating an ordered sequence alters the pairwise similarity structure (representational alignment) of the latent space. Shorter sequences focus on core semantics, while longer sequences include residual details.
Distributional Mismatch: This shift causes the latent distribution and the corresponding score function to vary significantly across token lengths. Consequently, a single diffusion model trained jointly across varying lengths fails to generalize effectively, often performing worse than models trained on specific lengths.
Limitation of Existing Solutions: Standard representation alignment losses require matching token dimensions (N = M), which is impossible in variable-length settings where cardinality differs.

2. Methodology

The paper proposes a novel merging-based variable-length tokenizer that modulates length by combining tokens rather than truncating them. The core innovation is Learnable Global Merging (LGM), designed to ensure representational alignment across different token counts while remaining compatible with generative models.

A. Merging-Based Length Modulation
Instead of dropping tokens, the method merges tokens based on an assignment matrix $\Gamma$ .

Proportional Attention: To handle merged tokens, the decoder and diffusion transformer utilize proportional attention, where the attention score is adjusted by the effective cluster size ( $m_j$ ) of each merged token.
Full-Length Equivalent Representation: The authors demonstrate that processing modulated latents with proportional attention is mathematically equivalent to processing projected full-length latents ($Wz$) with standard attention. This equivalence allows for the direct measurement and minimization of the difference between original latents ( $z$ ) and merged latents ($Wz$), effectively minimizing representation shift by encouraging the merging of similar tokens.

B. Learnable Global Merging (LGM)
A challenge arises because standard merging (e.g., agglomerative clustering) is data-dependent, deriving merge patterns from the input image. This is incompatible with diffusion generation, where the target image is unavailable.

Data-Independence: The authors introduce Learnable Global Embeddings ( $e$ ) that are independent of the input data. These embeddings determine the merging pattern for a given target reduction.
Optimization: The merging pattern is derived via agglomerative clustering on the learnable embeddings. Since clustering involves discrete operations, a straight-through trick is applied to enable gradient-based optimization of the embeddings.
Alignment Loss: An alignment loss ( $L_{align}$ ) is introduced to ensure that the similarity structure of the learnable global embeddings ( $e$ ) matches the similarity structure of the actual latent tokens ( $z$ ). This guides the global merging pattern to group similar tokens, thereby minimizing the representation shift defined in the merging objective.

C. Diffusion Transformer Integration

Positional Embeddings: The method employs merged positional embeddings, computed in the same manner as the latents, to handle the positional information of aggregated tokens.
Post-Training: A lightweight LoRA (Low-Rank Adaptation) post-training step is introduced. After joint training across lengths, length-specific LoRA parameters are trained for a few epochs to further adapt the model to specific token counts with negligible overhead.

3. Key Contributions

Merging-Based Tokenization: A new VLT architecture that modulates length via token merging, enabling direct representational alignment across different token counts by minimizing the shift between original and merged latents.
Learnable Global Merging: A data-independent merging mechanism that resolves the incompatibility between data-dependent merging and diffusion generation, ensuring the merging pattern is accessible during inference.
Improved Quality-Compute Trade-off: The integration of this tokenizer with a Diffusion Transformer achieves superior performance in balancing generation quality (gFID) and computational cost compared to existing VLT methods.

4. Experimental Results

Experiments were conducted on ImageNet 256×256 class-conditional generation.

Representational Alignment: The proposed method achieves significantly higher Centered Kernel Nearest-Neighbor Alignment (CKNNA) scores across different token lengths compared to nested dropout. Notably, the alignment is often higher than that of independently trained fixed-length tokenizers.
Joint Training Performance:
- Nested Dropout: Shows a large performance gap (0.8–1.7 gFID) between models trained jointly across lengths and those trained separately.
- Proposed Method: Reduces this gap to less than 0.3 gFID. With LoRA post-training, the joint model achieves performance within 0.05 gFID of separately trained models.
System-Level Comparison:
- The method outperforms existing VLTs (Semanticist, FlexTok) and fixed-length baselines (SoftVQ, MAETok) in terms of the gFID–compute trade-off.
- It achieves competitive generation quality at high compression ratios (e.g., 32 tokens) without requiring a diffusion-based tokenizer decoder, unlike some baselines.
Reconstruction Quality: The tokenizer maintains strong reconstruction metrics (rFID, PSNR, SSIM) across varying token lengths, particularly at higher compression ratios.

5. Significance and Claims

The paper claims that the primary barrier to effective variable-length diffusion models is the cross-length shift in latent distribution caused by conventional truncation methods. By shifting to a merging-based approach with data-independent, learnable patterns, the authors successfully maintain representational alignment.

The significance lies in enabling a single diffusion model to operate effectively across diverse token counts, allowing for flexible quality-compute control without the need for multiple model variants. The authors position this as a step toward more efficient deployment of generative models across scenarios with varying computational constraints, noting that while a gap remains between 1D and 2D tokenizers, their approach offers a promising direction for hybrid architectures. The work is framed as an efficiency improvement rather than a new generative capability, with no claimed ethical concerns.