Source-linked AI summary
Cubic Discrete Diffusion: Discrete Visual Generation on High-Dimensional Representation Tokens
Yuqing Wang, Chuofan Ma, Zhijie Lin, Yao Teng, Lijun Yu, Shuai Wang, Jiaming Han, Jiashi Feng, Yi Jiang, Xihui Liu
TL;DR
CubiD addresses the challenge of generating high-dimensional discrete representation tokens, whose massive structured spaces are difficult to model efficiently. It uses fine-grained masking and iterative refinement across the spatial-dimensional tensor, achieving state-of-the-art discrete generation while preserving representation capabilities for understanding and generation.
Problem
High-dimensional discrete representations produce massive structured token spaces that direct autoregressive generation would require O(hwd) steps to model.
Method
CubiD discretizes pretrained encoder features dimension-wise and performs fine-grained masking across the unified h × w × d tensor, iteratively predicting masked dimensions from partial observations.
Results
CubiD preserves the semantic quality of discretized representations and achieves state-of-the-art ImageNet 256×256 discrete generation, with a gFID of 1.88 for CubiD-XXL.
Takeaways & Limitations
The same native high-dimensional discrete tokens can support both understanding and generation without compressing or reorganizing the original representation space.
Takeaways & Limitations
Because CubiD uses a frozen pretrained encoder, reconstruction quality limits generation quality; experiments report approximately 18 dB reconstruction PSNR.
Abstract
from arXiv · showhide
Visual generation with discrete tokens has gained significant attention as it enables a unified token prediction paradigm shared with language models, promising seamless multimodal architectures. However, current discrete generation methods remain limited to low-dimensional latent tokens (typically 8-32 dims), sacrificing the semantic richness essential for understanding. While high-dimensional pretrained representations (768-1024 dims) could bridge this gap, their discrete generation poses fundamental challenges. In this paper, we present Cubic Discrete Diffusion (CubiD), the first discrete generation model for high-dimensional representations. CubiD performs fine-grained masking throughout the high-dimensional discrete representation -- any dimension at any position can be masked and predicted from partial observations. This enables the model to learn rich correlations both within and across spatial positions, with the number of generation steps fixed at $T$ regardless of feature dimensionality, where $T \ll hwd$. On ImageNet-256, CubiD achieves state-of-the-art discrete generation with strong scaling behavior from 900M to 3.7B parameters. Crucially, we validate that these discretized tokens preserve original representation capabilities, demonstrating that the same discrete tokens can effectively serve both understanding and generation tasks. We hope this work will inspire future research toward unified multimodal architectures. Code is available at: https://github.com/YuqingWang1029/CubiD.
1. Introduction
Visual generation commonly uses low-dimensional discrete tokens, but high-dimensional representations offer richer semantics while creating discretization and modeling challenges. CubiD addresses these challenges with fine-grained masking across the full spatial-dimensional tensor, enabling efficient generation and preserving representation capabilities.
- Motivation: High-dimensional pretrained features of 768–1024 dimensions could bring semantic richness to discrete visual generation, which often uses 8–32-dimensional tokens.This could support more unified visual and language tokenization.
- Challenges: Traditional vector quantization struggles at high dimensionality because sparse data requires exponentially larger codebooks, causing quantized features to drift from original semantics.The resulting drift can corrupt representation information.
- Challenges: High-dimensional discrete representations contain h×w×d tokens, making direct autoregressive generation require O(hwd) steps while standard diffusion misses intra-position dependencies.A typical 16×16×768 representation contains 196,608 tokens.
- CubiD: CubiD masks arbitrary dimensions at arbitrary spatial positions across the unified h×w×d tensor and iteratively predicts them from partial observations.This fine-grained cubic masking captures dependencies both within positions and across spatial locations.
- CubiD: CubiD uses a fixed number of refinement steps T regardless of feature dimensionality, converting sequential generation into parallel iterations while modeling spatial and dimensional correlations.Bidirectional attention operates over partially observed values.
- Results: 1.88 gFID: CubiD achieves state-of-the-art discrete generation on ImageNet 256×256 with 768-dimensional discrete tokens.Experiments also report scaling from 900M to 3.7B parameters and preservation of understanding and reconstruction capabilities.
2. Related Work
Visual tokenization typically compresses images into low-dimensional latent maps for reconstruction and generation, while newer approaches use high-dimensional pretrained vision features. Discrete generation methods improve sampling efficiency through parallel prediction but have generally operated on small token spaces.
- Visual Tokenization: Traditional VAE tokenizers encode images into low-dimensional continuous latent maps, typically with 4–32 dimensions, and jointly train encoders and decoders for reconstruction.The decoder reconstructs the image from the latent representation.
- Visual Tokenization: High-dimensional vision features from models such as DINOv2 or SigLIP can serve as latents without adaptation when paired with a specialized continuous-diffusion training schedule.This paper instead transforms such features into discrete tokens before training generative models.
- Discrete Visual Generation: Discrete autoregressive models generate tokens sequentially in O(N) steps, whereas discrete diffusion generates multiple tokens in parallel by restoring corrupted tokens.Common corruption schemes include absorbing-state masking, uniform transitions, and Gaussian-like transitions.
3. Method
CubiD discretizes frozen high-dimensional visual representations dimension-wise, then models the resulting h × w × d token tensor with fine-grained masking. Its iterative refinement predicts masked dimensions in parallel while preserving dependencies within spatial positions and across them.
- High-dimensional representation tokens: CubiD targets the joint distribution of h × w × d discrete tokens produced from high-dimensional pretrained representations.Dimension-wise quantization preserves the tensor structure, but creates a massive token space that direct autoregressive generation would require O(hwd) steps to traverse.
- High-dimensional representation tokens: Dimension-wise quantization independently maps each continuous feature value to one of L discrete levels without retraining the frozen encoder.The resulting tokens remain tractable for 768-dimensional features and preserve semantic quality on understanding tasks.
- Cubic discrete diffusion: CubiD treats the h × w × d tensor as a unified space where arbitrary dimensions at arbitrary positions can be masked and predicted from visible context.This fine-grained masking differs from methods that mask entire spatial positions and supports dependencies both within and across spatial locations.
- Cubic discrete diffusion: Training samples a masking ratio from a truncated Gaussian and independently masks positions across the tensor, replacing masked values with a learnable [MASK] token.The distribution spans [0, 1] and is biased toward aggressive masking with μ = 1.0, encouraging predictions from minimal context.
- Cubic discrete diffusion: The Transformer predicts all masked token categories from visible tokens using bidirectional attention and cross-entropy loss.Partial dimensions at each location provide context for learning how dimensions jointly encode information and constrain one another.
- Inference: At inference, CubiD starts from a fully masked tensor and progressively unmasks subsets while predicting all remaining masked tokens in parallel.A cosine schedule produces coarse-to-fine refinement, and the Transformer sequence length remains h × w regardless of feature dimensionality.
4. Experiments
Experiments show that dimension-wise quantization preserves representation quality, while CubiD benefits from fine-grained masking, iterative refinement, encoder robustness, and model scaling. On ImageNet-256, it achieves state-of-the-art discrete generation with native high-dimensional tokens.
- Discretization: Dimension-wise quantization maintains continuous-level understanding performance, whereas vector quantization significantly degrades LLaVA benchmark results.The comparison uses SigLIP2 features and dequantized inputs for discrete variants.
- Discretization: DINOv2-B matches baseline reconstruction rFID 0.57 at L = 8, while SigLIP2-B matches rFID = 0.69 at L = 16.These quantization settings are used in subsequent experiments.
- CubiD Design Studies: Per-element masking outperforms grouped alternatives: per-dim masking reaches gFID = 120.03, while fine-grained masking avoids severe artifacts and local inconsistencies.Per-element masking independently masks individual elements across the spatial-dimensional tensor.
- CubiD Design Studies: 256 inference steps improve generation over 64, with quality saturating around 512 iterations at gFID = 5.25 despite generating 196,608 discrete tokens.Autoregressive generation would require all 196,608 tokens as sequential steps.
- CubiD Design Studies: 946M-to-3.7B scaling lowers gFID from 5.25 to 4.68, demonstrating consistent gains with increased model capacity.The reported model sizes range from 946M to 3.7B parameters.
- Main Results: CubiD-XXL achieves 1.88 gFID on ImageNet-256 and improves from 2.37 to 1.88 gFID across L, XL, and XXL models with classifier-free guidance.It generates directly in the original high-dimensional representation space without reorganization or compression.
5. Conclusion
CubiD directly models native high-dimensional representation tokens through fine-grained masking across the spatial-dimensional tensor. The resulting discrete tokens preserve native representation ability for both understanding and generation while achieving state-of-the-art discrete generation without compression or reorganization.
- 5. Conclusion: CubiD turns high-dimensional discrete generation into parallel iterative refinement by masking across the entire spatial-dimensional tensor.The method avoids generating hundreds of thousands of sequential tokens.
- 5. Conclusion: Native representation ability lets the same discrete tokens support both understanding and generation without separate tokenization schemes.The conclusion states this preservation enables unified use across tasks.
Supplementary Material
The supplementary material adds implementation details, low-dimensional-token experiments, a limitations discussion, and additional image-generation results.
- Supplementary Material: The supplement provides implementation details for generation and understanding experiments.
- Supplementary Material: It includes additional CubiD experiments on low-dimensional tokens and a dedicated limitations discussion.
- Supplementary Material: It showcases additional image-generation results.
A.1. Generation Experiments
Training uses ImageNet-1K with large-batch mixed-precision optimization, EMA stabilization, and iterative unmasking for generation.
- A.1. Generation Experiments: CubiD models are trained on ImageNet-1K with a batch size of 2048 distributed across GPUs.Training uses fp16 mixed precision and an EMA of model weights with momentum 0.9999.
A.2. Understanding Experiments
The experiments evaluate whether CubiD’s discretized representations retain multimodal understanding abilities and whether the formulation generalizes across token dimensionalities. Results also examine the trade-off between compressed tokens’ generation quality and their representation quality.
- Understanding Evaluation: CubiD’s discretized tokens are evaluated with LLaVA on GQA, TextVQA, POPE, and MME to assess diverse multimodal understanding abilities.The setup compares experiments using original representations and discretized tokens.
- Low-Dimensional Tokens: CubiD is also tested on traditional low-dimensional tokens to assess whether its formulation generalizes beyond high-dimensional representations.The low-dimensional setup uses DC-AE-f32c32, producing 32-dimensional tokens.
- Low-Dimensional Tokens: 1.58 gFID and 188.7 IS on ImageNet 512×512 outperform USiT-2B’s 1.72 gFID using the same tokenizer.CubiD uses fewer parameters in this comparison.
- Compression Trade-off: Compressed 32-dimensional tokens achieve 1.55 gFID and 296.5 IS, while compression degrades the representation quality needed for understanding tasks.The tokens are obtained by reducing original 768-dimensional features with a learned projection layer optimized for reconstruction quality.
- Compression Trade-off: CubiD models the original high-dimensional tokens to preserve both generation and understanding capabilities despite the easier generation offered by lower-dimensional spaces.The paper frames this choice as a generation-understanding trade-off.
C. Limitations
CubiD remains constrained by its frozen representation encoder, its gap relative to continuous diffusion, and its inference cost. These limitations define important boundaries for generation quality and deployment efficiency.
- Representation Encoder: Approximately 18 dB reconstruction PSNR from the frozen pretrained encoder limits fine-grained detail in generated images.The encoder’s reconstruction quality sets an upper bound on generation quality.
- Continuous Generation: Discrete generation still trails continuous diffusion methods such as RAE, although the paper suggests advances in discrete generative modeling could narrow this gap.The comparison is framed as a remaining gap rather than a resolved limitation.
- Inference Efficiency: High-quality generation typically requires hundreds to a thousand inference steps, making acceleration an important direction.The paper identifies techniques developed for discrete language models as a possible avenue for improving efficiency.
D. More Visualization Results
Figures 6 and 7 present uncurated ImageNet 256×256 samples generated by CubiD-XXL under specified class conditions.
- Figure 6: Figure 6 shows uncurated ImageNet 256×256 samples from CubiD-XXL conditioned on specified classes.The caption identifies the samples and conditioning, but does not state a quantitative comparison.
- Figure 7: Figure 7 shows another set of uncurated ImageNet 256×256 samples from CubiD-XXL conditioned on specified classes.Use the figure to inspect class-conditioned visual samples rather than infer a benchmark result from the caption alone.