BiGain: Unified Token Compression for Joint Generation and Classification

Jiacheng Liu; Shengkun Tang; Jiacheng Cui; Dongkuan Xu; Zhiqiang Shen

cs.CVMar 12, 2026

BiGain: Unified Token Compression for Joint Generation and Classification

Authors: Jiacheng Liu, Shengkun Tang, Jiacheng Cui, Dongkuan Xu, Zhiqiang Shen

Organizations: VILA Lab, MBZUAI · North Carolina State University

Abstract

Acceleration methods for diffusion models (e.g., token merging or downsampling) typically optimize synthesis quality under reduced compute, yet often ignore discriminative capacity. We revisit token compression with a joint objective and present BiGain, a training-free, plug-and-play framework that preserves generation quality while improving classification in accelerated diffusion models. Our key insight is frequency separation: mapping feature-space signals into a frequency-aware representation disentangles fine detail from global semantics, enabling compression that respects both generative fidelity and discriminative utility. BiGain reflects this principle with two frequency-aware operators: (1) Laplacian-gated token merging, which encourages merges among spectrally smooth tokens while discouraging merges of high-contrast tokens, thereby retaining edges and textures; and (2) Interpolate-Extrapolate KV Downsampling, which downsamples keys/values via a controllable interextrapolation between nearest and average pooling while keeping queries intact, thereby conserving attention precision. Across DiT- and U-Net-based backbones and ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017, our operators consistently improve the speed-accuracy trade-off for diffusion-based classification, while maintaining or enhancing generation quality under comparable acceleration. For instance, on ImageNet-1K, with 70% token merging on Stable Diffusion 2.0, BiGain increases classification accuracy by 7.15% while improving FID by 0.34 (1.85%). Our analyses indicate that balanced spectral retention, preserving high-frequency detail and low/mid-frequency semantics, is a reliable design rule for token compression in diffusion models. To our knowledge, BiGain is the first framework to jointly study and advance both generation and classification under accelerated diffusion, supporting lower-cost deployment.

Figures & tables

Figure 1: Framework of our BiGain TM method. A Laplacian filter is applied to hidden-state tokens to compute local frequency scores. In each spatial stride, the lowest-scoring token is selected as a destination token, while the others form the source set. Destination and source tokens are gathered globally, and a bipartite matching selects top source-destination pairs.

Figure 2: Impact of token compression on diffusion models as our motivation on COCO2017 and ImageNet-100. Left: ToMe [ 3 ] (baseline) vs. Laplacian-Gated Merge (ours) as the merge ratio increases. Right: ToDo [ 26 ] (baseline) vs. Interpolate-Extrapolate KV-Downsampling (ours) as the downsample factor grows. Curves report percent change relative to the uncompressed model (

\uparrow

better; for FID we plot

-\Delta

FID%). Blue: classification accuracy. Orange: generation quality (FID).

Method	Acceleration Type	FLOPs Reduction $\uparrow$	Acc@1 $\uparrow$ (%)	$\Delta$ vs. Baseline $\downarrow$
Baseline (No Accel.)	None	–	81.03	–
ToMe [ 3 ]	Token Merging/Pruning	$10$ %	72.96	$\downarrow$ 8.07
DiP-GO [ 32 ]	Model Pruning	$10$ %	76.53	$\downarrow$ 4.50
SiTo [ 31 ]	Token Merging/Pruning	$7$ %	68.84	$\downarrow$ 12.19
MosaicDiff [ 10 ]	Model Pruning	$10$ %	77.38	$\downarrow$ 3.65
BiGain ${}_{\texttt{TM}}$ (Ours)	Token Merging/Pruning	$10\%$	78.38	$\downarrow$ 2.65

Table 1: Classification accuracy (Acc@1) on Pets dataset under similar FLOPs reduction.

Pets
Method	No Accel.	Classification $\uparrow$ (TD $\times$ )							No Accel.	Generation $\downarrow$ (TD $\times$ )
Method	No Accel.	2 $\times$	3 $\times$	4 $\times$	5 $\times$	6 $\times$	7 $\times$	8 $\times$	No Accel.	2 $\times$	3 $\times$	4 $\times$
Avg-pooling (baseline)		77.02	73.45	71.26	69.00	67.56	66.66	65.13	35.01	38.50	39.42	39.74
ToDo [ 26 ]	81.03	81.30	79.15	77.46	72.74	66.74	62.87	56.16		33.52	32.38	31.48
BiGain ${}_{\texttt{TD}}$ (Ours)		81.52	79.91	78.03	74.92	70.86	69.33	66.03		32.19	30.44	29.21
$\Delta\uparrow$		$\uparrow$ 0.22	$\uparrow$ 0.76	$\uparrow$ 0.57	$\uparrow$ 2.18	$\uparrow$ 4.12	$\uparrow$ 6.46	$\uparrow$ 9.87		$\downarrow$ 1.33	$\downarrow$ 1.94	$\downarrow$ 2.27

Table 2: SD-2.0 Token Downsampling : Classification (Acc@1 on Pets, ImageNet-100; Acc@1 and mAP on COCO-2017) and generation fidelity (FID

\downarrow

) vs. downsampling factor. For classification, we fix the interextrapolation factor at 0.9 across all timesteps to ensure stability. For generation, we linearly vary the factor from 0.8 (early steps) to 1.2 (later steps), shifting emphasis from low- to high-frequency information. Gray color indicates the same generation results as the above group.

Pets
Method	No Accel.	Classification $\uparrow$ (Token Merging Ratio)							No Accel.	Generation $\downarrow$ (Token Merging Ratio)
Method	No Accel.	10%	20%	30%	40%	50%	60%	70%	No Accel.	10%	20%	30%	40%	50%	60%	70%
ToMe	81.03	80.10	79.88	78.44	76.42	72.96	69.93	65.76		35.05	35.30	35.71	36.26	37.00	37.63	38.35
BiGain ${}_{\texttt{TM}}$ (Ours)		81.16	81.16	80.40	80.07	78.38	76.04	74.63	35.01	35.00	35.12	35.01	35.99	36.52	36.99	37.73
$\Delta\uparrow$		$\uparrow$ 1.06	$\uparrow$ 1.28	$\uparrow$ 1.96	$\uparrow$ 3.65	$\uparrow$ 5.42	$\uparrow$ 6.11	$\uparrow$ 8.87		$\downarrow$ 0.05	$\downarrow$ 0.18	$\downarrow$ 0.70	$\downarrow$ 0.27	$\downarrow$ 0.48	$\downarrow$ 0.64	$\downarrow$ 0.62
ImageNet-100

Table 4: SD-2.0 Token Merging : Classification (Acc@1 on Pets, ImageNet-100/1K; Acc@1 and mAP on COCO-2017) and generation fidelity (FID

\downarrow

) vs. Token Merging Ratio.

Figure 3: Qualitative comparison of BiGain ${}_{\texttt{TM}}$ and ToMe on SD-2.0 backbone. From left to right: BiGain ${}_{\texttt{TM}}$ with 70%, 50%, and 30% merge ratios, no acceleration, then ToMe with 30%, 50%, and 70% merge ratios.

Figure 4: Qualitative comparison of BiGain ${}_{\texttt{TD}}$ and ToDo on SD-2.0 backbone. From left to right: BiGain ${}_{\texttt{TD}}$ with downsampling factors $4\times$ , $3\times$ , and $2\times$ , no acceleration, then ToDo with factors $2\times$ , $3\times$ , and $4\times$ .

	SA only			SA+CA			SA+MLP			SA+CA+MLP
Method	0.7	0.5	0.3	0.7	0.5	0.3	0.7	0.5	0.3	0.7	0.5	0.3
ToMe [ 3 ]	65.76	72.96	78.44	61.68	68.41	74.46	51.43	58.71	66.35	50.86	59.53	66.20
BiGain ${}_{\texttt{TM}}$ (Ours)	74.63	78.38	80.40	73.89	78.03	79.56	68.27	74.93	77.98	68.25	74.84	77.95

Table 6: Ablation of token-merging locations in Stable Diffusion 2.0 on Pets. Self-Attention (SA) is always merged; Cross-Attention (CA) and MLP are toggled. Results reported at merge ratios

r\!\in\!\{0.7,0.5,0.3\}

. The underlined results indicate the best performance across all configurations.

Dataset	Method	GFLOPs	10%	20%	30%	40%	50%	60%	70%
Pets	Laplacian Gated Merge	704.99	81.16	81.16	80.4	80.07	78.38	76.04	74.63
	Cached Assignment Merge	698.88	80.29	79.97	79.89	79.01	78.11	75.91	74.49
	Adaptive Block Merge	695.08	80.40	80.16	79.99	79.18	77.84	75.96	74.13
ImageNet-100	Laplacian Gated Merge	704.99	71.94	71.62	70.64	68.38	67.24	65.20	61.28
	Cached Assignment Merge	698.88	71.76	71.16	70.44	69.38	67.78	64.56	61.28
	Adaptive Block Merge	695.08	72.58	71.94	70.58	70.52	68.04	65.36	60.98

Table 7: Further speedup on SD-2.0 Token Merging : classification performance vs. merge ratio. Acc@1 for single-label datasets; Acc@1 and mAP for multi-label COCO-2017. GFLOPs are measured at merge ratio r = 0.7.

Scoring method	0.7	0.5	0.3
Global mean deviation	72.96	77.84	79.91
$\ell_{1}$ -norm	73.02	77.11	79.86
$\ell_{2}$ -norm	72.72	77.95	79.61
Channel variance	73.04	77.95	79.83
Laplacian Filter $\ell_{1}$	74.63	78.38	80.40
Laplacian Filter $\ell_{2}$	74.24	77.81	79.80

Table 8: Ablation over token scoring heuristics for Stable Diffusion 2.0. Top-1 acc. (%) on Pets dataset across merge ratios. Local Laplacian signals outperform global or spectral metrics.

Figure 5: Visualization of our Laplacian-based frequency heuristic on hidden representations from Stable Diffusion-2.0. We probe U-Net at the highest-resolution upsampling stage. The visualization is computed from a noised image without a text prompt, showing the model’s intrinsic frequency-aware reconstruction dynamics. To reduce variance, we randomly sample 100 independent noise realizations and visualize the averaged token salience map.

Figure 6: Comparison of token merging schemes. Left: ToMe [ 3 ] ; Right: Our BiGain ${}_{\texttt{TM}}$ . Merging is applied with a merge ratio 90% at the highest-resolution latent layer of the U-Net transformer in Stable Diffusion 2.0 at denoising step $t=200$ . Grayscale indicates merged tokens.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Classes	Split	# Images (Cls.)	# Images (Gen.)
ImageNet-100 [ 30 ]	100	Val.	5,000	5,000
ImageNet-1K [ 22 ]	1,000	Val.	2,000	50,000
Oxford-IIIT Pets [ 17 ]	37	Test	3,669	3,669
COCO-2017 [ 14 ]	80	Val.	5,000	5,000

Appendix

Table 9: Dataset statistics with official splits used in our experiments.

Dataset	$N_{\text{stages}}$	TrialList	KeepList
ImageNet-100	2	[5, 20]	[5, 1]
COCO-2017	2	[5, 20]	[5, 1]
Oxford-IIIT Pets	2	[5, 20]	[5, 1]
ImageNet-1K	3	[5, 20, 100]	[50, 10, 1]

Appendix

Table 10: Adaptive diffusion-classifier parameters per dataset.

N_{\text{stages}}

is the number of pruning stages; TrialList is the cumulative number of Monte Carlo trials per candidate by stage; KeepList is the number of candidates retained after each stage.

Method	Time $\downarrow$ (s / batch)	Acceleration $\uparrow$ ( % )	FLOPs $\downarrow$ (G)
Baseline (No Accel.)	11.98	–	804.26
SiTo [ 31 ]	8.71	27.30	748.49
ToMe [ 4 ]	7.37	38.48	704.87
Laplacian Gated Merge (Ours)	7.37	38.48	704.99
Cached Assignment Merge (Ours)	7.29	39.15	698.88
Adaptive Block Merging (Ours)	7.27	39.32	695.08

Appendix

Table 11: Stable Diffusion 2.0 efficiency (batch size 4). Wall-clock sampling time per batch (seconds) excluding VAE encode/decode. All rows use merge ratio

r=0.7

Downsampling strategy	Acc@1 $\uparrow$	KV token sparsity
Todo (Nearest-Neighbor) [ 26 ]	72.30	75%
Low-frequency tokens (lowest-laplacian)	45.58	75%
High-frequency tokens (Highest-laplacian)	26.56	75%

Appendix

Table 12: Classification results on frequency-based KV selection on ImageNet-100. We compare the standard ToDo strategy with frequency-aware variants that select tokens with the highest or lowest Laplacian scores globally. Retaining only high- or low-frequency tokens severely degrades classification performance, highlighting the need to preserve a balanced spectrum.

Metric Name	Formula
Global mean deviation	$F_{i,j}=\\|\bm{x}_{i,j}-\bm{\mu}\\|$
$\ell_{1}$ norm	$F_{i,j}=\\|\bm{x}_{i,j}\\|_{1}$
$\ell_{2}$ norm	$F_{i,j}=\\|\bm{x}_{i,j}\\|$
Channel variance	$F_{i,j}=\frac{1}{C}\sum_{c=1}^{C}\left((\bm{x}_{i,j})_{c}-\frac{1}{C}\sum_{c^{\prime}=1}^{C}(\bm{x}_{i,j})_{c^{\prime}}\right)^{2}$
Laplacian ( $\ell_{1}$ )	$F_{i,j}=\frac{1}{C}\sum_{c=1}^{C}\left\|({\bm{X}}*{\bm{L}})_{i,j,c}\right\|$ , ${\bm{L}}=\begin{bmatrix}0&1&0\\ 1&-4&1\\ 0&1&0\end{bmatrix}$
Laplacian ( $\ell_{2}$ )	$F_{i,j}=\sqrt{\frac{1}{C}\sum_{c=1}^{C}\left(({\bm{X}}*{\bm{L}})_{i,j,c}\right)^{2}}$

Appendix

Table 13: Formulas of different metrics.

Method	No accel.	ToMe (70%)	BiGain ${}_{\text{TM}}$ (70%)	ToDo (2 $\times$ )	BiGain ${}_{\text{TD}}$ (2 $\times$ )
Acc. ( $n{=}2000$ ) (%) $\uparrow$	$57.05\pm 2.17$	$37.35\pm 2.13$	$44.50\pm 2.19$	$55.75\pm 2.17$	$56.30\pm 2.17$
Acc. ( $n{=}10000$ ) (%) $\uparrow$	$57.92\pm 0.97$	$39.85\pm 0.96$	$45.25\pm 0.98$	$56.50\pm 0.97$	$57.59\pm 0.97$

Appendix

Table 14: ImageNet-1K subset robustness. Diffusion-classifier accuracy on 2K and 10K subsets with 95% Wilson confidence intervals.

$\alpha$	0.0	0.5	0.8	0.9 ( BiGain ${}_{\texttt{TD}}$ )	1.0	1.2
Acc. ( $s{=}2$ ) (%) $\uparrow$	77.02	78.74	79.97	81.52	81.30	64.27
Acc. ( $s{=}4$ ) (%) $\uparrow$	71.26	75.63	77.48	78.03	77.46	45.46

Appendix

Table 15: IE-KVD

\alpha

sensitivity. SD-2.0 / Oxford-IIIT Pets diffusion classification accuracy (Acc., %).

Linear schedule $(\alpha_{\text{start}}\!\rightarrow\!\alpha_{\text{end}})$	FID $\downarrow$
ToDo	33.52
IE-KVD: $0.2\!\rightarrow\!0.8$	35.56
IE-KVD: $0.5\!\rightarrow\!1.0$	33.95
IE-KVD: $0.7\!\rightarrow\!1.0$	33.89
IE-KVD: $0.8\!\rightarrow\!1.2$	32.19
IE-KVD: $0.0\!\rightarrow\!1.2$	32.54

Appendix

Table 16: IE-KVD schedule ablation. SD-2.0 / Oxford-IIIT Pets generation quality (FID

\downarrow

) under different linear schedules.

Method	Acc@1 $\uparrow$ (%)	FID $\downarrow$
No accel.	81.03	35.01
ToMe	65.76	38.35
ToDo	81.30	33.52
BiGain ${}_{\text{TM}}$	74.63	37.73
BiGain ${}_{\text{TD}}$	81.52	32.19
L-GTM(enc) + IE-KVD(dec)	79.53	34.84

Appendix

Table 17: Joint use of L-GTM and IE-KVD. SD-2.0 / Oxford-IIIT Pets with

r{=}0.7

s{=}2

\alpha{=}0.9

Method	Time $\downarrow$ (ms/step)	Speedup $\uparrow$ ( $\times$ )	FLOPs $\downarrow$ (G)
No accel.	235.65	1.00	804.26
ToMe	144.31	1.63	704.87
ToDo	145.50	1.62	717.44
BiGain ${}_{\text{TM}}$	142.88	1.65	704.99
BiGain ${}_{\text{TD}}$	150.30	1.57	717.44
L-GTM(enc) + IE-KVD(dec)	147.43	1.60	716.23

Appendix

Table 18: Wall-clock UNet step time. SD-2.0 on RTX 4090 (FP16, batch=4, 50 steps; averaged over 50 runs; VAE excluded).

Token Compression

See all 26 papers

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence
Sep 30, 2026Xu Huang, Ye Huang, Zijun Liao +6Autoencoders Diffusion-Based Image Generation
Peking University · Singapore University of Technology and Design
APEX: Asynchronous Prefix Denoising for Extreme Compression
Oct 9, 2026David Dinkevich, Nisan Chiprut, Yoav HaCohen +1Latent Diffusion Models Diffusion Model Compression
Lightricks · Hebrew University of Jerusalem
Elastic Token Compression for Pixel-Space Diffusion Transformers
Aug 29, 2026Eduard Zamfir, Christian Reisswig, Zongwei Wu +2Visual Token Compression Diffusion Model Compression
University of Würzburg · Google

Sep 30, 2026J/K move · Enter open · S save

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence

Xu Huang, Ye Huang, Zijun Liao, Yuwei Niu+5

Peking University · Singapore University of Technology and Design

Autoencoders Diffusion-Based Image Generation Diffusion Model Acceleration T2I Diffusion Models Visual Token Compression

Open paper AskPDF arXiv

High-compression tokenizers are essential for scaling latent image generative models. However, aggressive compression creates a fundamental tradeoff between reconstruction fidelity and generation efficiency: high compression image encoder always increases the learning difficulty of diffusion training, resulting in slow model convergence. Recent representation autoencoders speed up the diffusion training by improving the latent feature's expressive capability by replacing VAE encoders with pretrained semantic encoders, yet they are typically limited to moderate compression and lose pixel-level details necessary for faithful reconstruction. To achieve both high compression and fast diffusion training, we propose DC-SAE, a Decoupled Compact Semantic Autoencoder designed for high-compression image generation with accelerated diffusion model convergence. DC-SAE consists of two key components: (1) a macro-level architecture design that leverages semantic encoders to enable higher compression ratios, and (2) a pixel-level encoder that preserves low-level details, ensuring high-fidelity image reconstruction. We empirically demonstrate that DC-SAE performs strongly on image generation tasks, achieving both compact latent representations and efficient training dynamics. Specifically, on the ImageNet dataset with

512 \times 512

resolution, DC-SAE achieves

32\times

spatial compression, with 29.79 PSNR and 3.37 gFID, substantially outperforming the previous state-of-the-art high-compression tokenizer baselines DC-AE by 13.5% and 54.9% on PSNR and gFID, respectively, maintaining comparable throughput and faster diffusion model training convergence. Beyond class-conditional generation, a

1.6

B-parameter DiT using DC-SAE achieves 0.84 on GenEval and 86.007 on DPG-Bench for text-to-image generation at

1024\times1024

resolution.

Abstract

Token Compression

Explore similar work

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence

APEX: Asynchronous Prefix Denoising for Extreme Compression

Elastic Token Compression for Pixel-Space Diffusion Transformers