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.
Figures & tables
Figure 1: Limitations of Vanilla Cross Entropy Loss with One-hot Target. Given an input image, vanilla CE loss cannot distinguish between non-closet tokens in the embedding space, even though some of the tokens are close to ground truth in embedding space and can decode to semantically similar images. Addtional details of loss computation in thig figure can be found in appendix.
Figure 2: Toy examples on 2D Gaussians. (a) ground truth distribution (b) 100-sample dataset (c) discretized dataset (d) L2 regression results (e) CE loss results (f) SNCE loss results
Model
Params
Tokenizer
Tokenizer Pretraining
Codebook
Epoch
FID ↓
IBQ-B
342M
IBQ Shi et al. (2025)
ImageNet256
16,384
300
2.88
IBQ-B-Ours (CE)
577M*
Emu3.5-IBQ Cui et al. (2025)
Large-Scale T2I
131,072
100
7.53
IBQ-B-Ours (SNCE)
577M*
Emu3.5-IBQ Cui et al. (2025)
Large-Scale T2I
131,072
100
3.62
IBQ-B-Ours (CE)
577M*
Emu3.5-IBQ Cui et al. (2025)
Large-Scale T2I
131,072
300
5.44
IBQ-B-Ours (SNCE)
577M*
Emu3.5-IBQ Cui et al. (2025)
Large-Scale T2I
131,072
300
3.42
IBQ-B-Ours (CE)
846M*
FVQ Zhu et al. (2024)
ImageNet256
262,144
300
4.11
Table 1: Class-conditioned Image Synthesis on ImageNet256 datset. *Model have identical-sized transformer layer. Parameter count increased due to larger token embedding and final linear head.
Parms
Codebook
Single ↑
Two ↑
Position ↑
Counting ↑
Color ↑
Attribution ↑
GenEval ↑
SDXL Podell et al. (2023)
3B
Cont.
0.98
0.74
0.39
0.85
0.15
0.23
0.55
DALLE 3 OpenAI (2023)
-
Cont.
0.96
0.87
0.47
0.83
0.43
0.45
0.67
SD3 Esser et al. (2024)
8B
Cont.
0.99
0.94
0.72
0.89
0.33
0.60
0.74
Flux-Dev Labs (2024)
12B
Cont.
0.99
0.85
0.74
0.79
0.21
0.48
0.68
Playground v3 Li et al. (2024)
-
Cont.
0.99
0.95
0.72
0.82
0.50
0.54
0.76
BAGEL Deng et al. (2025)
14B
Cont.
0.99
0.94
0.64
0.81
0.88
0.63
0.82
Table 2: Text to Image Generation Performance on GenEval Dataset . Cont. refers to continuous latent diffusion models.
Model
Params
Codebook
DPG ↑
MJHQ-30k
FID ↓
HPSv3 ↑
SD3 Esser et al. (2024)
8B
Cont.
83.5
11.92
-
Flux-Dev Labs (2024)
12B
Cont.
-
10.15
-
Show-o Xie et al. (2024)
1B
8,192
-
15.18
-
MMaDa Yang et al. (2025)
8B
8,192
53.4
32.85
-
LaViDa-O Li et al. (2025a)
10B
8,192
81.8
6.68
8.81
Table 3: Text-to-Image Generation Performance on DPG Benchmark and MJHQ-30k Dataset.
Model
Add ↑
Adjust ↑
Extract ↑
Replace ↑
Remove ↑
Background ↑
Style ↑
Hybrid ↑
Action ↑
Overall ↑
GPT-4o OpenAI (2024)
4.61
4.33
2.90
4.35
3.66
4.57
4.93
3.96
4.89
4.20
Qwen2.5VL+Flux Wang et al. (2025)
4.07
3.79
2.04
4.13
3.89
3.90
4.84
3.04
4.52
3.80
FluxKontext dev Labs et al. (2025)
3.76
3.45
2.15
3.98
2.94
3.78
4.38
2.96
4.26
3.52
OmniGen2 Wu et al. (2025a)
3.57
3.06
1.77
3.74
3.20
3.57
4.81
2.52
4.68
3.44
UniWorld-V1 Lin et al. (2025)
3.82
3.64
2.27
3.47
3.24
2.99
4.21
2.96
2.74
3.26
BAGEL Deng et al. (2025)
3.56
3.31
1.70
3.30
2.62
3.24
4.49
2.38
4.17
3.20
Table 4: Image Editing Performance on ImgEdit benchmark.
Table 5: Training configurations. We report the relevant hyperparameters for training, including the learning rate, number of training steps, optimizer setup, image resolution fo generation tasks.
Table 6: Additional Experiment Results on ImageNet256.