Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

Latent Diffusion Models (LDMs) typically rely on tokenizers with a fixed compression ratio, forcing a rigid trade-off between computational cost and image quality. While variable-length tokenizers (VLTs) have been proposed to allow for adaptive compression, existing methods—primarily based on nested dropout—truncate token sequences. This truncation disrupts the semantic consistency of the latent space across different lengths, making it difficult to train a single diffusion model that performs well across varying token counts.

The authors propose a novel variable-length tokenizer that modulates length by merging tokens rather than truncating them. By merging tokens based on a learnable, data-independent pattern, the model ensures that the merging logic is accessible during generation. This approach uses proportional attention to handle the varying number of tokens and incorporates a straight-through trick to optimize the merging assignments via gradient descent. By ensuring that the merging process is consistent and global, the model maintains a stable latent distribution, which facilitates the training of a single, flexible diffusion transformer.

The study demonstrates that merging-based length modulation significantly improves representational alignment across different token lengths compared to traditional nested dropout. By minimizing the representation shift, the proposed tokenizer allows a single diffusion transformer to achieve a superior gFID-compute trade-off. Experiments on ImageNet 256x256 show that this approach outperforms existing VLT methods, providing a more efficient and flexible framework for visual synthesis across diverse computational constraints.

This research provides a robust solution for deploying diffusion models in resource-constrained environments. By enabling a single model to adapt its computational footprint dynamically without sacrificing generation quality, this method reduces the need for maintaining multiple model variants, thereby simplifying deployment and improving efficiency in real-world applications.