🤖 machine learning

BiGain: Unified Token Compression for Joint Generation and Classification

BiGain is a training-free, plug-and-play framework that accelerates diffusion models by employing frequency-aware token compression operators to simultaneously enhance classification accuracy and maintain or improve generation quality.

Original authors: Jiacheng Liu, Shengkun Tang, Jiacheng Cui, Dongkuan Xu, Zhiqiang Shen

Published 2026-03-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jiacheng Liu, Shengkun Tang, Jiacheng Cui, Dongkuan Xu, Zhiqiang Shen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart artist (a Diffusion Model) who can paint beautiful pictures from scratch. This artist is also a brilliant detective who can look at a picture and tell you exactly what it is (e.g., "That's a golden retriever!").

However, this artist is slow. To paint a picture or solve a mystery, they have to look at every single tiny pixel in the image, one by one. It's like trying to read a whole encyclopedia to find one word.

The Problem: The "Fast" Artists Are Bad Detectives

Scientists have tried to make this artist faster by telling them to ignore some pixels. They say, "Hey, just look at the big shapes and skip the tiny details!"

The Result: The artist becomes fast and still paints pretty pictures.
The Catch: When the artist tries to act as a detective, they fail miserably. Why? Because the "tiny details" (like the edge of a cat's ear or the texture of a dog's fur) are exactly what the detective needs to identify the animal. By skipping them to save time, the artist becomes blind to the clues.

The Solution: BiGain (The "Smart Filter")

The authors of this paper, BiGain, came up with a clever way to speed up the artist without making them blind. They realized that images have two types of information:

The Big Picture (Low Frequency): The general shape, colors, and layout. (Good for painting).
The Fine Details (High Frequency): Edges, textures, and sharp lines. (Crucial for identifying).

Most speed-up methods act like a blender: they mix everything together, losing the sharp edges. BiGain acts like a smart sieve.

How BiGain Works: Two Magic Tools

1. The "Laplacian Gated Merging" (The Smart Glue)

Imagine you have a pile of LEGO bricks representing the image.

Old Way: You grab two random bricks and glue them together to save space. If you glue a brick from a smooth blue sky to a brick from a jagged tree branch, you get a muddy mess. The detective can't tell the tree from the sky anymore.
BiGain Way: Before gluing, BiGain uses a special scanner (a Laplacian filter) to check how "bumpy" or "smooth" each brick is.
- If a brick is from a smooth area (like a blue sky), BiGain says, "Glue this one to its neighbor! It's safe."
- If a brick is from a bumpy area (like a cat's whiskers), BiGain says, "Stop! Don't touch this one. It's a vital clue!"
- Result: You save space by gluing the boring parts, but you keep all the important, sharp edges intact.

2. The "Interpolate-Extrapolate KV Downsampling" (The Selective Zoom)

In the artist's brain, there are three types of notes: Queries (What am I looking for?), Keys (What is here?), and Values (What does it look like?).

Old Way: When speeding up, they shrink all the notes. The artist forgets exactly where to look.
BiGain Way: They shrink the Keys and Values (the "what is here" and "what it looks like" notes) but leave the Queries (the "what am I looking for" notes) at full size.
- Analogy: Imagine you are searching a library. You keep your search list (Query) perfectly detailed so you know exactly what book you want. But you shrink the bookshelves (Keys/Values) to fit more books in a smaller room. You can still find the right book because your search list is sharp, even if the shelves are packed tighter.

Why This Matters

BiGain is like a dual-purpose tool that doesn't force you to choose between speed and accuracy.

For the Artist: It keeps the painting looking beautiful because the smooth parts are still there.
For the Detective: It keeps the classification accurate because the sharp edges and textures are preserved.

The Bottom Line

Before this, speeding up these AI models meant sacrificing their ability to "see" details. BiGain introduces a frequency-aware approach: it knows the difference between a boring, smooth background and a critical, detailed edge. By treating them differently, it makes the AI faster, smarter, and more accurate all at once, without needing to retrain the model from scratch. It's the first time we've successfully taught an AI to run fast and keep its eyes wide open.

1. Problem Statement

Diffusion models have become the backbone of modern generative systems, but their high computational cost during sampling has spurred the development of acceleration techniques, primarily token merging (e.g., ToMe) and token downsampling (e.g., ToDo).

The Gap: Existing acceleration methods are optimized almost exclusively for synthesis quality (e.g., maintaining low FID scores). They often ignore the model's discriminative capacity.
The Consequence: While these methods may preserve visual "look," they frequently degrade performance in downstream recognition tasks (classification). Naive compression tends to remove high-frequency details (edges, textures, fine structures) that are critical for classification, even if the global semantic content remains intact.
The Challenge: There is a need for a token compression strategy that simultaneously preserves generative fidelity (requiring low/mid-frequency semantic coherence) and discriminative utility (requiring high-frequency detail preservation), without requiring model retraining.

2. Core Insight: Frequency Separation

The authors propose that the solution lies in frequency separation.

Low/Mid Frequencies: Represent global shapes, layouts, and semantics (crucial for coherent generation).
High Frequencies: Represent edges, fine textures, and small objects (crucial for classification).
Design Principle: A successful compression method must perform balanced spectral retention. It should merge smooth, redundant regions (low-frequency) while preserving high-contrast, detail-rich regions (high-frequency).

3. Methodology: BiGain Framework

BiGain is a training-free, plug-and-play framework composed of two frequency-aware operators designed to be inserted at inference time.

A. Laplacian-Gated Token Merging (L-GTM)

Goal: Merge spectrally smooth tokens while protecting high-contrast tokens.
Mechanism:
1. Frequency Scoring: A 2D Laplacian filter is applied to the hidden state tokens to compute a local frequency magnitude score for each token. High scores indicate edges/textures; low scores indicate smooth regions.
2. Destination Selection: In each spatial grid, the token with the lowest frequency score is selected as a "destination" (anchor).
3. Source Selection: All other tokens in the grid form the "source" set.
4. Merging: Top similar source-destination pairs are identified via bipartite matching (based on cosine similarity) and merged via averaging.
Effect: This acts as a selective low-pass filter, collapsing flat regions to save compute while retaining edges and micro-structures essential for classification.

B. Interpolate-Extrapolate KV-Downsampling (IE-KVD)

Goal: Reduce attention computation costs by downsampling Keys (K) and Values (V) while keeping Queries (Q) intact.
Mechanism:
- Queries (Q): Kept at full resolution to preserve the model's ability to localize and attend precisely (critical for both generation alignment and classification scoring).
- Keys/Values (K/V): Downsampled using a controllable operator $D_{\alpha,s}$ that interpolates between nearest-neighbor pooling (preserves detail) and average pooling (smooths).
- Control Parameter ( $\alpha$ ): Allows tuning the trade-off between detail preservation and smoothing.
Effect: Reduces memory and FLOPs ( $O(N \tilde{N}d)$ ) without destroying the discriminative cues embedded in the attention mechanism.

4. Key Contributions

Reframing Token Compression: The first work to treat token compression in diffusion models as a bi-objective optimization problem (Generation + Classification) rather than a single-objective (Synthesis) problem.
Frequency-Aware Operators: Introduction of L-GTM and IE-KVD, which explicitly use spectral analysis (Laplacian filtering) to guide compression decisions, ensuring high-frequency details are preserved.
Training-Free & Plug-and-Play: The framework requires no retraining, fine-tuning, or architectural changes. It works with both U-Net (Stable Diffusion) and DiT (Diffusion Transformer) backbones.
Theoretical Justification: The paper provides a spectral analysis showing that balanced spectral retention tightens the decision bound for diffusion classifiers by minimizing margin loss while controlling variance.

5. Experimental Results

The authors evaluated BiGain on Stable Diffusion 2.0 (U-Net) and DiT-XL/2 across multiple datasets: ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017.

Classification Performance:
- BiGain significantly outperforms baselines (ToMe, ToDo, DiP-GO) in classification accuracy under acceleration.
- Example (ImageNet-1K, SD 2.0): With a 70% token merging ratio, BiGain improves classification accuracy by 7.15% compared to the baseline, while maintaining or slightly improving generation quality (FID improved by 0.34).
- In downsampling scenarios, BiGain consistently reduces the accuracy drop seen in baselines as the downsampling factor increases.
Generation Performance:
- BiGain maintains FID scores comparable to or better than unaccelerated models and other acceleration methods.
- Qualitative analysis shows that BiGain preserves structural details (e.g., cat edges in the Pets dataset) better than ToMe, which tends to blur fine structures.
Efficiency:
- Achieves significant FLOP reductions (up to ~40% in some configurations) and wall-clock speedups (approx. 1.6x) on RTX 4090 GPUs.

6. Significance and Impact

Dual-Purpose Deployment: BiGain enables the deployment of diffusion models for dual-purpose systems (e.g., medical imaging for both diagnosis and reconstruction, or industrial inspection for defect detection and generation) without sacrificing performance in either task.
Design Rule: It establishes "balanced spectral retention" as a robust design rule for token compression in diffusion models, moving beyond simple similarity-based merging.
Accessibility: By being training-free, it lowers the barrier for deploying efficient, high-performance generative models in resource-constrained environments.

In summary, BiGain bridges the gap between "what looks good" and "what classifies well" in accelerated diffusion models by intelligently preserving the high-frequency information necessary for recognition while compressing redundant low-frequency data.