Source-linked AI summary
SNCE: Geometry-Aware Supervision for Scalable Discrete Image Generation
Shufan Li, Jiuxiang Gu, Kangning Liu, Zhe Lin, Aditya Grover, Jason Kuen
TL;DR
Large VQ codebooks improve image fidelity but make discrete image generators difficult to optimize because supervision becomes sparse and one-hot targets penalize semantically similar alternatives. SNCE replaces those targets with geometry-weighted neighborhood distributions, and experiments across ImageNet-256, text-to-image generation, and image editing report faster convergence and better generation quality than standard CE.
Problem
Large codebooks improve reconstruction quality but make training difficult through sparse token supervision, while one-hot CE treats nearby semantically similar tokens as incorrect.
Method
SNCE trains discrete image generators with soft neighborhood targets whose probabilities depend on distances between code embeddings and the encoded image latent.
Results
SNCE improves convergence speed and final generation fidelity over standard CE across ImageNet-256 generation, text-to-image synthesis, and image editing.
Takeaways & Limitations
Embedding-space geometry provides a practical training signal for scaling discrete image generators to large vocabularies, positioning SNCE as a drop-in CE replacement.
Takeaways & Limitations
Generated images are not pixel-perfect and may contain small artifacts; the model also inherits hallucination and social-bias limitations from LaViDa-O.
Abstract
from arXiv · showhide
Recent advancements in discrete image generation showed that scaling the VQ codebook size significantly improves reconstruction fidelity. However, training generative models with a large VQ codebook remains challenging, typically requiring larger model size and a longer training schedule. In this work, we propose Stochastic Neighbor Cross Entropy Minimization (SNCE), a novel training objective designed to address the optimization challenges of large-codebook discrete image generators. Instead of supervising the model with a hard one-hot target, SNCE constructs a soft categorical distribution over a set of neighboring tokens. The probability assigned to each token is proportional to the proximity between its code embedding and the ground-truth image embedding, encouraging the model to capture semantically meaningful geometric structure in the quantized embedding space. We conduct extensive experiments across class-conditional ImageNet-256 generation, large-scale text-to-image synthesis, and image editing tasks. Results show that SNCE significantly improves convergence speed and overall generation quality compared to standard cross-entropy objectives.
1 Introduction
Discrete image generation offers scalable alternatives to continuous latent diffusion, but large codebooks create sparse and brittle supervision. SNCE addresses this by weighting neighboring tokens according to VQ embedding geometry, improving convergence and generation fidelity across evaluated tasks.
- Latent image generation models encode pixels into low-dimensional embeddings, which are more effective and scalable to model than raw pixels.
- Discrete generators quantize continuous image latents into tokens and commonly train token distributions with cross-entropy loss.
- Larger codebooks improve reconstruction quality but require larger models and more data, making generator training difficult.
- 1,280 average occurrences per token for a 200K codebook versus 31,250 for an 8,192-sized codebook illustrates how vocabulary growth sparsifies per-token supervision.
- One-hot cross-entropy assigns all probability to one token and treats geometrically close alternatives as equally incorrect, despite their potentially similar reconstructions.
- SNCE assigns higher target probability to tokens whose embeddings are closer to the encoded image latent, providing multiple nearby tokens with positive learning signals.
- SNCE improves convergence speed and final generation fidelity relative to standard CE across ImageNet-256, text-to-image generation, and image editing experiments.
2 Background and Related Works
Discrete image generation uses tokenized image representations with autoregressive or discrete-diffusion generation, but large codebooks remain difficult to optimize. Related work establishes the tokenizer pipeline, generation paradigms, and the underexplored role of geometry-aware soft labels.
- 2.1 Discrete Image Tokenizer: A VQ tokenizer encodes images into discrete codes using a learnable codebook, while larger codebooks can improve fidelity but may cause low utilization and latent collapse.
- 2.1 Discrete Image Tokenizer: The tokenizer maps pixels x to continuous latents z and selects each discrete token y_i by minimizing distance to codebook vectors.
- 2.2 Discrete Image Generation: Discrete image generators are categorized as autoregressive models, which generate tokens left to right, and discrete diffusion models, which progressively unmask tokens.
- 2.2 Discrete Image Generation: Most existing systems use relatively small codebooks because large-codebook optimization is challenging; Emu3.5 uses a 131,072-token codebook with a 30B-parameter model and massive data.
- 2.2 Discrete Image Generation: The shared likelihood term in autoregressive and discrete-diffusion objectives uses one-hot cross-entropy, producing weak per-token signals as codebook size grows.
- 2.3 Soft Labels: Soft labels are established in classification, segmentation, graph learning, and distillation, but remain underexplored for discrete image generation.
- 2.3 Soft Labels: SNCE is presented as the first soft-label objective explicitly targeting token sparsity caused by large codebooks in discrete image generation.
3 Method
SNCE replaces one-hot token supervision with a geometry-aware neighborhood distribution derived from distances between continuous image latents and codebook embeddings. Used as a drop-in replacement for cross-entropy in autoregressive and discrete diffusion models, it spreads learning signals across semantically similar tokens and supports smoother optimization.
- 3.1 Stochastic Neighbor Embedding: SNCE defines a categorical neighborhood distribution over codebook tokens using the distance between each continuous latent and each code embedding.The distribution uses a fixed temperature and the tokenizer’s distance metric; the setup uses τ = 0.71, negative dot product for IBQ, and L2 distance for FVQ.
- 3.1 Stochastic Neighbor Embedding: Unlike t-SNE, SNCE compares an arbitrary continuous latent with a finite codebook and uses a fixed shared temperature instead of per-sample bandwidth search.This design avoids the computational cost of binary-search bandwidth selection during training and accommodates varying codebook density.
- 3.2 Stochastic Neighbor Cross Entropy Loss: SNCE replaces the one-hot target in cross-entropy with the neighborhood distribution q(zi), weighting each token’s predicted log-probability by its embedding-based proximity.The resulting objective can be applied to both autoregressive and discrete diffusion models, with only the conditioning term differing between them.
- Interpretations: SNCE can be interpreted as stochastic quantization in expectation or as knowledge distillation from the tokenizer’s neighborhood distribution to the generator.Compared with explicit stochastic quantization, operating directly on q(zi) avoids Monte Carlo sampling variance; the tokenizer acts as a teacher encoding continuity in latent space.
- 3.3 Gradient Analysis: The soft supervision gives semantically similar tokens positive gradients proportional to proximity, whereas standard cross-entropy penalizes such alternatives despite potentially faithful reconstructions.This distributes updates across neighboring tokens and mitigates token-frequency imbalance by allowing more tokens to receive positive learning signals.
4 Experiments
Experiments show that SNCE improves optimization and generation quality for discrete image models, from a 2D toy setting to ImageNet-256, text-to-image synthesis, and image editing. The gains arise from providing positive supervision to neighboring tokens rather than only the closest code.
- Toy example: SNCE better approximates a multimodal 2D Gaussian distribution than L2 regression or CE under sparse supervision.CE fits the finite training samples but misses the underlying distribution, whereas SNCE assigns non-zero probability to neighboring tokens.
- Toy example: SNCE injects a continuity inductive bias by providing positive training signals to nearby tokens in large-vocabulary settings.This differs from one-hot CE, which supervises only the closest token.
- ImageNet-256: SNCE accelerates convergence and improves FID over CE on ImageNet-256 with large codebooks.The experiment uses 131,072- and 262,144-code tokenizers with a 342M-parameter autoregressive transformer.
- Text-to-image synthesis: SNCE outperforms CE across text-to-image evaluation metrics, including FID (−3.67) and HPSv3 (+0.12).Evaluation covers GenEval, DPG, MJHQ-30k, and HPSv3 for 1024 × 1024 synthesis.
- Image editing: SNCE improves overall image-editing quality over CE by +0.13 on ImgEdit.The adaptation uses a smaller editing dataset, which may exacerbate low token frequency for large codebooks.
- Qualitative comparison: Qualitatively, SNCE improves text alignment, spatial structure, low-level detail fidelity, and preservation of input image structure while reducing editing artifacts.These comparisons are reported for text-to-image generation and image editing.
5 Conclusion
SNCE addresses large-codebook optimization bottlenecks by assigning positive training signals to nearby tokens in embedding space. Across small- and large-scale generation and editing experiments, it accelerates convergence and improves image fidelity over standard CE.
- SNCE modifies discrete image-generator training so nearby, not only nearest, tokens receive positive signals based on embedding-space proximity.This targets optimization bottlenecks in large-codebook settings.
- Experiments span ImageNet validation, large-scale text-to-image generation, and instruction-based image editing.
- SNCE accelerates convergence and improves image fidelity across multiple benchmarks compared with standard CE.
- The authors aim for SNCE to facilitate larger codebooks in next-generation discrete image foundation models.
6 Additional Technical Details
The technical details connect SNCE to autoregressive and discrete-diffusion objectives, categorical VAE interpretation, knowledge distillation, and on-policy learning. In each case, geometric neighborhood supervision replaces or interprets standard one-hot CE supervision.
- Training objectives: SNCE replaces the standard log-likelihood term JCE with JSNCE in both autoregressive and discrete diffusion training objectives.The replacement is formulated for each model family.
- Training objectives: For both objectives, the expectation moves from quantized latents y to continuous latents z because quantization deterministically produces the relevant tokens.
- Categorical VAE interpretation: The categorical VAE interpretation treats stochastic neighborhood sampling as a categorical tokenizer distribution rather than deterministic nearest-code quantization.
- Knowledge distillation interpretation: In the distillation view, the teacher is an implicit distribution q defined by continuous latent geometry, whose samples correspond to visually similar tokenized images.
- On-policy learning perspective: The on-policy derivation identifies token probabilities as the policy and shows that the resulting gradient exactly matches the gradient of JSNCE.
7 Additional Experiment Details and Results
Additional details describe SNCE visualization, large-codebook experimental setups, label-smoothing comparisons, temperature behavior, and qualitative gains. The results support geometry-aware supervision over uniform target smoothing.
- Visualization and target construction: SNCE assigns a second-closest token 100× the probability of any other token in an illustrative soft-target construction.This avoids the infinite-versus-zero cross-entropy behavior of differing versus matching one-hot targets.
- Toy example: The toy experiment uses a two-Gaussian mixture quantized onto a 50 × 50 grid, producing 2,500 tokens, with 100 training points and 2,000 MLP steps.
- Large-scale training setup: The large-scale setup uses 50M text-to-image samples and 2M image-editing samples over 200k steps with global batch size 1,024.Image-editing data is introduced during the final 100k steps.
- Comparison with label smoothing: Label smoothing provides limited benefits because uniform probability on non-target tokens ignores embedding geometry and can still penalize semantically adjacent tokens.
- Temperature: As τ approaches 0, SNCE becomes standard CE; as τ approaches infinity, its target becomes uniform and may have low signal-to-noise.
- Temperature: τ = 0.71 achieves the best ImageNet performance, while larger and smaller temperatures degrade results.
- Qualitative results: SNCE produces higher-quality qualitative results, particularly in facial features and eyes, and preserves input structure with fewer editing artifacts.
8 Additional Discussions
The discussion contrasts SNCE with factorized-tokenizer approaches for large codebooks. SNCE instead provides a generic training solution for flat codebooks without factorization assumptions.
- Factorized VQ methods reduce the effective vocabulary through multiple smaller codebooks but impose structural assumptions on the latent representation and tokenizer.
- SNCE trains discrete image generators with a large flat codebook without assuming factorization structure.The paper highlights this setting as relevant to flat codebooks that scale well in practice.
9 Limitations
SNCE improves visual quality and text alignment, but generated images remain imperfect and may contain small artifacts. The model also inherits hallucination and social-bias limitations from its base model, LaViDa-O.
- Generated images are not pixel-perfect and may contain small artifacts despite improved visual quality and text alignment.
- The model inherits hallucination and social-bias limitations from the LaViDa-O base model.