Adaptive Token Reduction Image Representation

---
name: adaptive-token-reduction-image-representation
title: "When Less is Enough: Adaptive Token Reduction for Efficient Image Representation"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2503.16660"
keywords: [Token Reduction, Vision Transformers, Feature Selection, Image Compression, Efficiency]
description: "Adaptively prune visual tokens from vision encoders by reconstructing discarded features from retained ones, reducing computational cost by 50% while maintaining task performance on OCR and image understanding tasks."
---

## 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:

```python
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:

```python
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:

```python
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