Core Concepts: Tokens, Attention, and Redundancy

Transformers process input sequences by breaking them into discrete units called tokens, which are typically subword or word-level representations. Each token is embedded into a high-dimensional vector space, where its semantic and syntactic relationships with other tokens are modeled through self-attention mechanisms. The self-attention operation computes pairwise interactions between tokens, generating a weighted sum of values based on their relevance, as defined by:

Here, Q (queries), K (keys), and V (values) are linear transformations of the input embeddings, and d_k is the dimension of the key vectors. The softmax normalization ensures that attention weights sum to one, allowing the model to focus on the most relevant tokens dynamically.

Token Redundancy in Transformer Models

Despite their effectiveness, transformers exhibit significant computational inefficiencies due to token redundancy. Empirical studies show that many attention heads assign near-uniform weights to large subsets of tokens, particularly in deeper layers. This phenomenon arises because:

• Localized attention patterns: Many tokens attend primarily to their immediate neighbors, making distant tokens computationally irrelevant.
• Low-rank attention matrices: The rank of the attention matrix is often much smaller than the sequence length, implying that a subset of tokens can sufficiently represent the full interaction space.

Quantifying Redundancy

Redundancy can be measured using the token importance score, which quantifies how much a token contributes to the final output. For a given token t_i, its importance I_i can be approximated via gradient-based methods or attention weight aggregation:

where α_ij is the attention weight from token i to j, and ∇_xi ℒ(x_j) is the gradient of the loss with respect to the token embedding. Tokens with low I_i are candidates for merging or pruning.

Practical Implications

Redundancy-aware methods like Token Merging (ToMe) and Dynamic Token Pruning exploit these insights to reduce computational overhead without significant accuracy loss. For instance, ToMe merges similar tokens in intermediate layers using clustering, while pruning methods eliminate tokens below an adaptive threshold. These techniques achieve speedups of 1.5–2× in vision transformers and autoregressive language models.

Case Study: BERT Inference Optimization

In BERT, up to 40% of tokens in classification tasks can be pruned after the first few layers with minimal impact on accuracy. This is validated by the observation that later layers rely heavily on [CLS] token representations, making many input tokens redundant. Pruning these tokens reduces FLOPs by 30% while maintaining 98% of the original F1 score on GLUE benchmarks.

Why Merge or Prune Tokens? Efficiency vs. Performance Trade-offs

Transformer models, particularly in large-scale applications like vision or language tasks, process sequences of tokens where computational complexity scales quadratically with sequence length due to self-attention mechanisms. Token merging and pruning techniques address this inefficiency by reducing the number of tokens processed in intermediate layers, trading off minimal accuracy degradation for significant computational savings.

Computational Complexity of Self-Attention

The self-attention mechanism in Transformers has a time and memory complexity of O(n²d), where n is the sequence length and d is the embedding dimension. For high-resolution images or long documents, n can reach thousands of tokens, making attention layers prohibitively expensive. Token reduction methods mitigate this by lowering the effective n in later layers.

Token Merging: Preserving Information via Aggregation

Token merging combines semantically similar tokens, typically using cosine similarity or learned attention weights. For example, ToMe (Token Merging) merges the k most similar tokens in a layer by averaging their embeddings:

where S_i is the set of tokens merged into the i-th output token. This preserves global context while reducing redundancy, often achieving 30–50% FLOPs reduction with <1% accuracy drop in vision Transformers.

Token Pruning: Dynamic Sparsification

Pruning discards low-salience tokens entirely, often based on attention scores or learned importance metrics. For instance, a token x_i may be pruned if its L2-norm falls below a threshold τ:

DynamicViT and similar methods prune up to 40% of tokens in later layers, disproportionately benefiting inference speed since pruned tokens bypass all subsequent computations.

Trade-offs: Accuracy, Latency, and Memory

• Accuracy: Merging is more stable than pruning, as no information is fully discarded. Pruning risks losing rare but critical tokens (e.g., small objects in images).
• Latency: Pruning reduces compute more aggressively, especially for long sequences, but merging’s regularity enables better hardware utilization.
• Memory: Both methods lower peak memory usage by reducing the active token count, crucial for edge deployment.

Practical Considerations

In vision Transformers, merging is often applied after early layers where spatial redundancy is high. For NLP, pruning may be preferred in decoder-only models (e.g., GPT) to maintain autoregressive consistency. The choice hinges on the task’s tolerance for information loss versus the need for real-time throughput.

Key Metrics: FLOPs Reduction, Memory Savings, and Accuracy Impact

When evaluating token merging and pruning techniques in transformer models, three primary metrics dominate the analysis: FLOPs reduction, memory savings, and accuracy impact. These metrics provide a quantitative framework for assessing the trade-offs between computational efficiency and model performance.

FLOPs Reduction

Floating-point operations (FLOPs) serve as a proxy for computational cost. Token merging reduces FLOPs by decreasing the number of tokens processed in self-attention layers. For a transformer with N tokens, the FLOPs for self-attention scale as O(N²d), where d is the feature dimension. Merging k tokens into m merged tokens (where m < k) reduces FLOPs proportionally.

Here, r represents the merging ratio (k/m). For example, merging every 2 tokens (r=2) reduces attention FLOPs by approximately 4×. However, this assumes uniform merging; dynamic merging strategies may yield variable savings.

Memory Savings

Memory consumption in transformers primarily stems from storing attention matrices and intermediate activations. Token pruning directly reduces memory by eliminating tokens entirely, while merging compresses them. The memory footprint for attention scales as:

Pruning 30% of tokens reduces this to 4(0.7N)² = 1.96N² bytes, a 51% reduction. Merging achieves similar savings by reducing N while preserving information through weighted combinations. Memory savings extend beyond attention to layer norms and feed-forward networks, as fewer tokens pass through subsequent layers.

Accuracy Impact

The critical trade-off lies in accuracy degradation. Pruning and merging inevitably discard information, potentially harming task performance. The accuracy drop depends on:

• Token selection criteria: Importance scores based on attention weights or gradient norms preserve critical tokens better than random pruning.
• Merging strategy: Weighted averaging (e.g., ToMe) often outperforms naive concatenation.
• Task sensitivity: Image classification tolerates higher compression than fine-grained sequence tasks like named entity recognition.

Empirical studies show that careful merging can limit accuracy drops to 1-2% on ImageNet at 2× FLOPs reduction, while aggressive pruning (50% tokens) may incur 5-10% degradation. The Pareto frontier between FLOPs reduction and accuracy guides optimal strategy selection.

Practical Considerations

In real-world deployments, the choice between merging and pruning depends on hardware constraints and latency requirements. Merging introduces slight overhead from computing token similarities but maintains parallelizability. Pruning enables dynamic computation but may require sparse operations support. Profiling tools like NVIDIA Nsight help quantify actual speedups, as theoretical FLOPs reductions don't always translate linearly to wall-clock time due to memory bandwidth bottlenecks.

2.1 Similarity-Based Merging: Cosine and Euclidean Distance Approaches

Similarity-based merging in transformers relies on quantifying the resemblance between token representations to identify redundant or semantically equivalent tokens. Two dominant metrics for this purpose are cosine similarity and Euclidean distance, each offering distinct advantages depending on the nature of the token embeddings.

Cosine Similarity for Token Merging

Cosine similarity measures the angular alignment between two vectors, making it invariant to their magnitudes. Given token embeddings x and y in a d-dimensional space, their cosine similarity is computed as:

Values range from −1 (perfect anti-correlation) to 1 (identical direction). For transformer tokens, values closer to 1 indicate semantic equivalence, enabling merging of tokens with similarity exceeding a threshold τ. This approach is particularly effective in high-dimensional spaces where magnitude variations are noise-dominated.

Euclidean Distance as an Alternative Metric

Euclidean distance quantifies the straight-line separation between vectors:

Unlike cosine similarity, Euclidean distance is sensitive to both direction and magnitude. This makes it suitable for tasks where the norm of embeddings carries meaningful information, such as in energy-based models. Tokens are merged if their distance falls below a threshold ϵ, often normalized by embedding dimensionality.

Comparative Analysis

• Normalization Sensitivity: Cosine similarity inherently normalizes embeddings, while Euclidean distance requires explicit normalization (e.g., L2) for stable comparisons.
• Contextual Performance: Cosine similarity outperforms in NLP tasks where semantic alignment matters more than magnitude. Euclidean distance excels in low-dimensional or physics-inspired embeddings.
• Computational Overhead: Both metrics scale quadratically with token count, but cosine similarity avoids square root operations.

Practical Implementation

In transformer architectures, similarity-based merging is typically applied during inference or fine-tuning. For example, tokens in ViTs (Vision Transformers) with cosine similarity >0.9 are merged by averaging their embeddings, reducing sequence length without significant accuracy loss. Euclidean-based merging is preferred in graph transformers where node features have interpretable magnitudes.

Dynamic Token Merging with Learned Thresholds

Dynamic token merging extends static token reduction techniques by introducing adaptive thresholds learned during training. Unlike fixed merging rules, this approach optimizes the trade-off between computational efficiency and model performance by allowing the network to determine which tokens to merge based on their contextual importance.

Learned Threshold Mechanism

The core innovation lies in parameterizing the merging decision through a gating function G that outputs a merge probability for each token pair. For tokens x_i and x_j, the merge score is computed as:

where W_g and b_g are learnable parameters, ⊙ denotes element-wise multiplication, and τ is a temperature parameter controlling decision sharpness. The sigmoid function σ maps scores to [0,1], interpreted as merge probabilities.

Differentiable Merging via Gumbel-Softmax

To enable gradient flow through discrete merge decisions, we employ the Gumbel-Softmax reparameterization:

where g₁, g₂ are i.i.d. Gumbel noise samples. This continuous relaxation approaches an exact merge/non-merge decision as τ→0 while remaining differentiable.

Training Objective

The model jointly optimizes the primary task loss L_task and a merging regularization term:

where λ controls the computational cost trade-off, and FLOPs(l) estimates layer-wise operations reduced through merging. The expectation is taken over stochastic merge decisions during training.

Architectural Integration

Dynamic merging layers are typically inserted between transformer blocks, processing the sequence as:

1. Compute pairwise similarity scores for all tokens in local windows
2. Sample merge decisions using the gating mechanism
3. Fuse tokens via weighted averaging based on merge probabilities
4. Propagate reduced sequence to next layer

The windowed approach maintains O(n) complexity relative to sequence length n, contrasting with the O(n²) cost of global attention.

Practical Considerations

Key implementation details include:

• Gradient estimation: Straight-through estimator for merge/non-merge indicators
• Temperature annealing: Start with high τ for exploration, gradually reduce to sharpen decisions
• Memory efficiency: Sparse implementation for merge masks avoids materializing full attention matrices

Empirical results show dynamic merging achieves 2-4× speedups in vision transformers with <1% accuracy drop, outperforming static methods like ToMe by adapting to input complexity.

2.3 Hierarchical Merging Strategies for Long Sequences

Hierarchical merging addresses the quadratic complexity of self-attention in transformers by progressively reducing sequence length through multi-level token aggregation. Unlike uniform merging, which applies a fixed compression ratio globally, hierarchical approaches preserve fine-grained local information while enabling coarse-grained global reasoning.

Multi-Resolution Token Grouping

The core mechanism partitions the input sequence into nested segments using a binary tree structure. At each level l, tokens are grouped into non-overlapping clusters of size k_l, where k_l increases exponentially with depth:

where k₀ is the base cluster size (typically 2-8 tokens). Cluster representations are computed via weighted averaging using attention scores as weights:

where 𝒩_i denotes the neighborhood of tokens in cluster i at level l.

Adaptive Depth Selection

The merging depth L is dynamically determined based on sequence entropy. For an input sequence S, the stopping criterion evaluates the KL-divergence between adjacent levels:

where P^(l) represents the attention weight distribution at level l, and ε is a threshold (typically 0.1-0.3). This ensures merging stops when further aggregation would cause significant information loss.

Gradient-Aware Merging

Recent advancements incorporate gradient signals during merging to preserve salient features. The cluster weights α_ij are modified to include gradient magnitudes:

where γ is a learnable parameter and ℒ is the task loss. This approach shows particular effectiveness in vision transformers, where gradient hotspots often correspond to semantically important regions.

Computational Complexity Analysis

Hierarchical merging reduces the asymptotic complexity of self-attention from O(n²) to O(n log n) for a sequence of length n. The breakdown per level shows:

• Level 0: Full attention on n/k₀ clusters (O((n/k₀)²))
• Level l: Attention on n/(2^lk₀) clusters (O((n/(2^lk₀))²))

The total complexity becomes dominated by the finest level when k₀ ≪ n.

Importance Scoring: Attention Weights and Gradient-Based Criteria

Attention-Based Importance Scoring

The self-attention mechanism in transformers provides a natural way to assess token importance through attention weights. For a given layer l with H attention heads, the raw attention weight matrix A^l,h for head h is computed as:

where Q^l,h and K^l,h are the query and key matrices respectively, and d_k is the dimension of the key vectors. The importance score s_i for token i can be derived by aggregating attention weights across all heads and layers:

This formulation captures both the influence of token i on other tokens (via outgoing attention weights) and its receptiveness to contextual information (via incoming attention weights). In practice, some implementations use only the outgoing attention weights for computational efficiency.

Gradient-Based Importance Measures

While attention weights provide a direct measure of token interaction, gradient-based methods offer a complementary perspective by quantifying how much each token contributes to the final loss. The gradient magnitude with respect to token embeddings serves as an importance indicator:

where e_i is the embedding of token i and ℒ is the loss function. More sophisticated approaches compute the product of gradients and embeddings (similar to saliency maps in computer vision):

This method tends to be more stable than pure gradient magnitudes as it accounts for both the token representation and its influence on the loss.

Combined Scoring Approaches

State-of-the-art token pruning methods often combine multiple importance signals. A common hybrid approach linearly combines attention and gradient scores:

where α is a mixing hyperparameter typically tuned on validation data. Some recent work has proposed more sophisticated combinations using learned gating mechanisms or small neural networks to dynamically weight different importance signals.

Practical Considerations

When implementing importance scoring in practice, several factors must be considered:

• Computational overhead: Gradient-based methods require backward passes, making them more expensive than attention-only approaches
• Layer selection: Attention weights from different layers capture different types of information - early layers often focus more on local syntax while later layers capture higher-level semantics
• Normalization: Scores should be properly normalized across the sequence to avoid length-dependent biases
• Dynamic vs static pruning: Some methods compute importance scores once per input (static) while others recompute them during processing (dynamic)

Recent work has shown that the optimal importance metric varies across tasks - for example, gradient-based methods tend to work better for tasks requiring fine-grained understanding, while attention-based methods suffice for more general language modeling tasks.

Adaptive Pruning: Layer-Specific Token Elimination

Unlike static pruning methods that apply uniform compression across all layers, adaptive pruning dynamically eliminates tokens based on layer-specific importance metrics. This approach recognizes that different transformer layers exhibit varying sensitivity to token removal, with early layers often processing redundant spatial information while later layers refine semantically critical features.

Importance Scoring for Adaptive Pruning

The core mechanism relies on computing token importance scores I_l(t) at each layer l through gradient-weighted class activation:

where z_l^t represents the token embedding at position t in layer l, and ℒ is the task loss. The Hadamard product with gradients captures both activation magnitude and downstream influence.

Layer-Wise Threshold Adaptation

Pruning thresholds τ_l are computed dynamically per layer using exponential moving averages of importance scores:

with smoothing factor α typically set to 0.1-0.3. This adaptive thresholding prevents over-pruning in critical layers while aggressively compressing redundant early layers.

Practical Implementation Considerations

• Memory-efficient scoring: Importance scores can be approximated using last-attention weights in decoder layers or spatial pooling in vision transformers
• Batch-wise stability: Layer thresholds should be updated across multiple batches to avoid oscillation
• Warmup period: Most implementations disable pruning for the first 10-20% of training steps

Architectural Modifications

Effective adaptive pruning requires three key architectural changes:

1. Token regrouping: Surviving tokens are reordered contiguously to maintain efficient matrix operations
2. Cross-layer propagation: Pruning masks must propagate through skip connections
3. Gradient masking: Eliminated tokens receive zero gradients during backpropagation

Recent work in DynamicViT demonstrates that layer-specific pruning can reduce FLOPs by 40-60% in vision transformers with less than 1% accuracy drop, significantly outperforming uniform pruning baselines. The technique shows particular promise in long-sequence domains like video processing and document understanding.

3.3 Combining Pruning with Quantization for Hardware Efficiency

Pruning and quantization are complementary techniques for optimizing transformer models for hardware deployment. While pruning reduces the number of parameters by eliminating redundant connections, quantization compresses the remaining weights into lower-bit representations. When combined, these methods yield multiplicative efficiency gains in memory footprint, compute latency, and energy consumption.

Mathematical Foundations of Joint Optimization

The joint effect of pruning and quantization can be formalized through the parameter density reduction ratio D:

where s_p is the sparsity ratio from pruning, b_q is the quantized bit-width, and b_orig is the original bit-width (typically 32-bit floating point). For example, applying 70% sparsity with 8-bit quantization yields:

indicating a 13.3× reduction in parameter storage requirements. The actual hardware benefits are even greater due to:

• Reduced memory bandwidth pressure from smaller activations
• Elimination of zero-value multiplications in pruned weights
• Simplified arithmetic circuits for lower-precision operations

Hardware-Aware Co-Design Strategies

Effective joint optimization requires consideration of hardware constraints during training:

1. Block-structured pruning aligns sparsity patterns with hardware vector units (e.g., 4×4 blocks for SIMD architectures)
2. Mixed-precision quantization allocates higher bits to sensitive attention layers while aggressively quantizing feed-forward networks
3. Sparsity-aware numerics uses bitmask encoding to skip zero computations without branch penalties

Integration with Popular Transformer Architectures (BERT, GPT, ViT)

Training Strategies: End-to-End vs. Post-Training Optimization

4.3 Debugging and Profiling Merging/Pruning Pipelines

def test_pruning_equivalence(model, pruned_model, inputs):
    with torch.no_grad():
        baseline = model(inputs)
        pruned = pruned_model(inputs)
        return torch.allclose(baseline, pruned, rtol=1e-3, atol=1e-5)
  

5.1 Foundational Papers on Token Efficiency in Transformers

5.2 Open-Source Implementations and Toolkits

5.3 Emerging Research Directions and Unresolved Challenges