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

Sep 30, 2026

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

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.

Xu Huang, Ye Huang, Zijun Liao +6

Peking University · Singapore University of Technology and Design

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

Ask PDF

Oct 9, 2026

APEX: Asynchronous Prefix Denoising for Extreme Compression

Fewer latent tokens promise cheaper diffusion. Yet aggressive latent compression creates a reconstruction-generation dilemma: adding channels improves reconstruction but makes generation harder. We introduce APEX, short for Asynchronous Prefix denoising for EXtreme compression, which pairs a hierarchical latent space with asynchronous denoising over channel groups. Our autoencoder orders channel prefixes from structure to detail, while one shared denoiser generates them at staggered local times without adding tokens. On ImageNet-512, APEX outperforms existing methods at every evaluated channel count under 64x and 128x compression. At matched training and guided sampling compute, it also outperforms a 256-token model at every evaluated budget, with larger gains at lower compute. Code is available at https://daviddinkevich.github.io/APEX.

David Dinkevich, Nisan Chiprut, Yoav HaCohen +1

Lightricks · Hebrew University of Jerusalem

Latent Diffusion Models Diffusion Model Compression Diffusion-Based Image Generation Generative Image Compression Diffusion-Based Denoising Diffusion Model Acceleration

Ask PDF

Aug 29, 2026

Elastic Token Compression for Pixel-Space Diffusion Transformers

Natural images concentrate their detail in a small fraction of the frame, yet diffusion models spend a full token on every patch, in every layer and at every timestep. The waste is largest in pixel-space models, with no autoencoder to absorb low-level redundancy first. Probing a pretrained pixel text-to-image transformer, we find its middle-block tokens redundant wherever the image is flat. The redundancy occupies connected, content-shaped regions, and exploiting it requires tokens with the same geometry. Cutting a Hilbert ordering of the patches provides them. Consecutive positions are always image neighbours, so any contiguous run is a connected region whose size and shape follow the content, and grouping in two dimensions becomes a cut in one. Existing reductions each lose part of this. Similarity merging scatters its groups, latent bottlenecks discard position, and skipping deletes what it should summarize. We cut where the model's features change most and pool each run into one region token. Our Region Token Interface (\method{}) adapts a diffusion model to these tokens, with the region count drawn at random during fine-tuning so one checkpoint serves every budget. \method{} leads prior reduction methods at matched budgets, matches dense quality at

2.0\times

the speed, and stays close at

2.6\times

. The code and models are open-sourced at https://eduardzamfir.github.io/rti

Eduard Zamfir, Christian Reisswig, Zongwei Wu +2

University of Würzburg · Google

Visual Token Compression Diffusion Model Compression Diffusion Model Acceleration Diffusion Models Diffusion Transformer Pixel-Space Diffusion Models

Ask PDF

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