Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

Abstract: Latent Diffusion Models (LDMs) have become dominant in visual synthesis, but their quality-compute trade-off is largely constrained by the tokenizer's fixed compression ratio. Variable-length tokenizers (VLTs) promise adaptive compression by varying token counts, allowing diffusion models to flexibly balance quality and compute. However, conventional VLTs modulate length by truncating ordered token sequences, which makes token semantics depend on token position and breaks representational alignment across lengths. This leads to a cross-length shift in the latent distribution that hinders a single variable-length diffusion model from operating effectively. To address this, we propose a novel variable-length tokenizer that modulates length by merging tokens. We show that encouraging similar tokens to merge enables direct cross-length representation alignment when the diffusion transformer operates according to the merging pattern. Since conventional merging methods are data-dependent, making the merging pattern inaccessible during generation, we introduce learnable global merging, which is data-independent, to ensure compatibility with diffusion transformers. On ImageNet 256$\times$256 generation, our merging-based variable-length tokenizer integrated with a diffusion transformer achieves a superior gFID-compute trade-off compared to prior VLT methods. Code is available at this https URL

• The paper introduces learnable global merging, a novel method for variable-length tokenization in diffusion transformers that adapts compression by merging similar latent tokens instead of truncating them.
• Empirical results demonstrate superior generative alignment and image quality across various compression ratios, significantly outperforming methods like nested dropout.
• The approach employs position-specific learnable embeddings and a dedicated alignment loss, ensuring consistent reconstruction quality and efficient performance in image synthesis.

Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

Motivation and Problem Formulation

The paper "Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers" (2606.20076) addresses a fundamental limitation in latent diffusion pipelines for image synthesis: the inflexibility imposed by fixed-length tokenization, which limits adaptive control over quality-compute trade-offs. Conventional latent diffusion models depend on a tokenizer with a fixed compression ratio, requiring separate models for varying computational budgets and target fidelities. Variable-length tokenizers (VLTs) promise adaptive compression, but prevalent methods (notably nested dropout) introduce representational misalignment across token lengths due to positional semantic drift—early tokens encode higher-level semantics while tail truncation drops low-level detail, resulting in length-dependent shifts in latent distributions.

Learnable Global Merging: Methodological Innovations

To overcome these representational alignment issues, the authors introduce a merging-based variable-length tokenizer. Length modulation is achieved by merging similar latent tokens rather than truncating ordered sequences. The merging operation is guided via learnable global embeddings, yielding a data-independent merging pattern. This allows direct alignment of latent representations across token lengths, as each reduced-length latent can be projected to a full-length equivalent via the merging pattern, facilitating standard alignment losses over matched cardinalities.

The merging pattern is constructed using position-specific learnable embeddings, optimized via a straight-through estimator over agglomerative clustering assignments. A dedicated alignment loss encourages the embeddings to reflect latent token-wise similarity (i.e., token pairs with high cosine similarity are assigned to the same cluster). The resulting assignment matrix enables consistent proportional attention and positional embedding merges within the diffusion transformer, ensuring compatibility during both training and generation.

Figure 1: Overview of the method—length modulation via learnable global merging is performed image-agnostically but optimized for alignment and tokenization.

Empirical Evaluation and Numerical Results

Experiments are conducted on class-conditional ImageNet generation at $256\times256$ and $512\times512$ resolutions using a SoftVQ-based 1D tokenizer and LightningDiT transformer backbone. The merging-based VLT achieves substantially higher representational alignment across token lengths relative to nested dropout, as measured by Centered Kernel Nearest-Neighbor Alignment (CKNNA).

Figure 2: CKNNA between encoded tokens with different token counts reveals consistently high cross-length alignment for learnable global merging versus nested dropout.

The generation quality, evaluated using conditional Fréchet Inception Distance (gFID), demonstrates minimal deterioration under joint training for all token lengths, contrasting sharply with nested dropout where joint training introduces a 0.8–1.7 gFID gap. Using LoRA fine-tuning per length further narrows the performance gap, with gFID improving to within 0.05 of separately trained models.

Figure 3: Image generation results under different FLOPs, highlighting superior gFID-compute trade-offs for merging-based VLT.

System-level comparisons show the proposed method outperforms prior VLT baselines (Semanticist, FlexTok) and fixed-length strategies across all CRs, using fewer parameters and without necessitating complex diffusion-based decoder architectures. The method achieves rFID, PSNR, and SSIM metrics that are competitive or superior at high compression ratios.

Figure 4: Reconstruction quality (rFID, PSNR, SSIM) across token lengths—merging-based VLT excels, especially under aggressive compression.

Ablations and Analytical Observations

Ablation studies isolate the contributions of global merging, learnable embeddings, and alignment loss. Data-dependent merging without global merging—akin to ToMe—yields significantly worse gFID, underlining the need for consistency in positional and attention mechanisms within the transformer. Learnable embeddings and alignment loss robustly enhance both generation and reconstruction quality.

Token similarity analyses confirm that global merging effectively groups similar tokens, with merged pairs exhibiting substantially higher average cosine similarity than the all-pair average.

Figure 5: Comparison of reconstruction between LGM (learnable global merging) and ToMe-style merging—LGM sacrifices little reconstruction quality but attains much better generative alignment.

Qualitative Generation and Scalability

Qualitative samples illustrate high visual coherence on ImageNet at token lengths as low as 32.

Figure 6: Generated class-conditional images on ImageNet $256\times256$ with LightningDiT-XL (+LoRA) at token length 32.

Scalability experiments at $512\times512$ resolution confirm that merging-based VLT maintains superior alignment and reconstruction under increased image size and compression. gFID at 32 tokens improves by 1.66 compared to nested dropout.

Figure 7: Reconstruction quality at $512\times512$ resolution remains robust for merging-based VLT.

Figure 8: Cross-length representational alignment at $512\times512$ resolution remains high for merging-based VLT.

Practical and Theoretical Implications

The learnable global merging paradigm fundamentally improves the compatibility of variable-length tokenization with latent diffusion models. This enables a single diffusion transformer to effectively operate over a spectrum of token counts, facilitating adaptive deployment in resource-constrained and high-fidelity settings. The technique aligns well with recent advances in representation alignment [repa, reg] and transformer attention mechanisms [tome]. The approach is extensible to hybrid architectures and multi-modal generative pipelines where efficient length modulation is desirable. Theoretically, it suggests a new direction for designing latent spaces with consistent inter-sample similarity structure under variable cardinality.

Conclusion

The paper realizes a practical and theoretically sound solution for variable-length tokenization in diffusion transformers by introducing learnable global merging, which enforces representational alignment and ensures generative compatibility across token counts. Empirical results on large-scale benchmarks demonstrate improvements in both quality-compute trade-offs and representational consistency compared to established VLTs. The method's implementation is efficient, and its conceptual framework offers promising avenues for future research in scalable generative modeling and adaptive compression mechanisms in vision-language systems.