ADu2021/adaptive-token-reduction-image-representation

Core Concept

This skill implements adaptive token reduction for vision encoders. The key insight is that many visual tokens are redundant—their information can be reconstructed from a smaller subset of more informative tokens. Rather than random pruning, this approach learns which tokens are essential by training a selector network to identify valuable tokens and a reconstructor to verify that discarded tokens can be faithfully recovered.

Architecture Overview

The system has three main components:

• Feature Selector (S): Three Transformer layers with a Gumbel-Softmax head that generates binary masks, choosing which tokens to keep or discard
• Feature Reconstructor (R): Three Transformer layers that reconstruct discarded tokens from retained ones plus a shared learnable masked embedding
• Optimization Objective: Balances reconstruction fidelity against pruning efficiency using modified regularization

The training uses an autoencoder-like framework where the selector learns to identify redundant tokens, and the reconstructor validates that removed tokens are recoverable.

Implementation

The feature selection mechanism uses Gumbel-Softmax for differentiable discrete choices. The following code shows the core selector and reconstructor modules:

import torch
import torch.nn as nn
from torch.nn.functional import gumbel_softmax

class TokenSelector(nn.Module):
    """Selects which tokens to retain using Gumbel-Softmax."""
    def __init__(self, hidden_dim, num_layers=3, num_heads=8):
        super().__init__()
        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                dim_feedforward=hidden_dim * 4,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        self.selection_head = nn.Linear(hidden_dim, 1)

    def forward(self, features, temperature=1.0, training=True):
        """
        Args:
            features: (B, N, D) token embeddings
            temperature: Gumbel-Softmax temperature
            training: if True, use stochastic selection
        Returns:
            masks: (B, N) binary masks
            logits: (B, N) selection logits
        """
        x = features
        for layer in self.transformer_layers:
            x = layer(x)

        logits = self.selection_head(x).squeeze(-1)

        if training:
            # Gumbel-Softmax for differentiable discrete selection
            masks = gumbel_softmax(logits.unsqueeze(-1), tau=temperature, hard=True)
            masks = masks.squeeze(-1)
        else:
            masks = (logits > 0).float()

        return masks, logits

The reconstructor mirrors this structure but focuses on recovering discarded tokens:

class TokenReconstructor(nn.Module):
    """Reconstructs removed tokens from retained ones."""
    def __init__(self, hidden_dim, num_layers=3, num_heads=8):
        super().__init__()
        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                dim_feedforward=hidden_dim * 4,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        self.masked_embedding = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.reconstruction_head = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, features, masks):
        """
        Args:
            features: (B, N, D) original token embeddings
            masks: (B, N) binary masks from selector
        Returns:
            reconstructed: (B, N, D) reconstructed features
        """
        # Create input with masked tokens replaced
        masked_input = features * masks.unsqueeze(-1)
        masked_input = masked_input + self.masked_embedding * (1 - masks.unsqueeze(-1))

        x = masked_input
        for layer in self.transformer_layers:
            x = layer(x)

        reconstructed = self.reconstruction_head(x)
        return reconstructed

Training uses L2 reconstruction loss with a pruning efficiency regularizer:

def compute_loss(original_features, reconstructed_features, masks, pruning_weight=0.1):
    """
    Args:
        original_features: (B, N, D) original tokens
        reconstructed_features: (B, N, D) reconstructed tokens
        masks: (B, N) selection masks
        pruning_weight: weight for pruning efficiency term
    """
    batch_size = masks.shape[0]

    # Reconstruction loss for discarded tokens only
    discard_mask = 1 - masks
    reconstruction_loss = torch.sum(
        ((original_features - reconstructed_features) ** 2) * discard_mask.unsqueeze(-1)
    ) / (torch.sum(discard_mask) + 1e-6)

    # Pruning efficiency: encourage removal of tokens
    pruning_ratio = torch.mean(1 - masks)

    # Modified regularization: max(L_pr, p) prevents trivial solutions
    pruning_loss = torch.max(
        pruning_weight * reconstruction_loss,
        torch.tensor(pruning_ratio)
    )

    total_loss = reconstruction_loss + pruning_loss
    return total_loss, reconstruction_loss, pruning_ratio

Practical Guidance

When to Use:

• Reducing inference latency in vision-language models (LLaVA, etc.)
• Processing high-resolution images where token count becomes a bottleneck
• OCR and image understanding tasks that benefit from aggressive compression
• Deployment scenarios where model size or memory is constrained

When NOT to Use:

• Complex reasoning tasks that require fine-grained spatial details
• Tasks where every pixel matters (e.g., small object detection)
• Low-resolution inputs where 50% pruning would be too aggressive
• When inference speed is not a critical constraint

Key Hyperparameters:

• num_layers (3-4): Depth of selector/reconstructor networks; deeper = more capacity
• temperature (0.5-2.0): Gumbel-Softmax temperature; lower = sharper, higher = softer
• pruning_weight (0.01-0.5): Balance between reconstruction and efficiency; higher = more aggressive pruning
• Training dataset size: Paper uses 100,000 COCO images

Common Pitfalls:

• Setting pruning_weight too high leads to trivial solutions that remove nearly all tokens
• Using token selection during training but hard masking during inference causes distribution mismatch
• Not evaluating on task-specific benchmarks; pruning effectiveness varies by task
• Forgetting to freeze the vision encoder while training selector/reconstructor

Performance Notes

• OCR tasks: Up to 50% token reduction with negligible degradation
• Reasoning-heavy tasks: Minimal benefit from aggressive pruning (suggest 20-30%)
• Training time: Approximately 24 hours on 100K COCO images with standard GPUs
• Inference overhead: Selection adds ~5-10% latency but saves much more downstream

References

• Gumbel-Softmax paper for differentiable discrete sampling
• Vision Transformer (ViT) and DeiT architectures
• LLaVA and LLaVA-NeXT multimodal models
• COCO dataset for training selector networks