Technical Summary: FAR (Function-preserving Attention Replacement for IMC-friendly Inference)

1. Problem Statement

Transformer architectures, particularly those relying on self-attention, face significant efficiency bottlenecks when deployed on emerging In-Memory Computing (IMC) hardware, such as ReRAM crossbars. While self-attention excels on parallel digital processors (GPUs/TPUs) by enabling global token integration, its computational characteristics fundamentally conflict with the IMC execution model:

• Activation-to-Activation Multiplications: Self-attention relies heavily on operations like $QK^\top$ and softmax normalization, which require dynamic, non-local mixing of activations rather than the weight-stationary vector-matrix multiplications (VMM) that IMC hardware optimizes.
• Memory Overhead: The quadratic complexity ($O(T^2)$) of token interactions and the need to repeatedly read spatially distributed activations lead to fragmented analog operations, extensive Analog-to-Digital/Digital-to-Analog (ADC/DAC) conversions, and poor crossbar utilization.
• Current Limitations: Existing IMC accelerators often offload attention to digital units, leaving the quadratic bottleneck unresolved, or attempt to map attention directly to analog arrays without addressing the inherent data movement costs.

Concurrently, recent analyses suggest that attention layers exhibit substantial redundancy during inference, with attention maps becoming smoother and compressible across depth. This motivates the hypothesis that the functional role of attention can be approximated by sequential modules that align better with IMC locality and weight-stationarity.

2. Methodology: The FAR Framework

The authors propose FAR (Function-preserving Attention Replacement), a framework that replaces all self-attention blocks in a pretrained Vision Transformer (specifically the DeiT family) with sequential modules, specifically Multi-head Bidirectional LSTMs (BiLSTMs). The approach is designed as a "drop-in" architectural conversion that does not require training from scratch.

2.1 Replacement Strategy

• Block-wise Substitution: Every attention block $A_l$ in a pretrained transformer is replaced by a learnable substitute $A'_l$ (a Multi-head BiLSTM). All other components (MLP blocks, patch embeddings, layer norms) remain frozen.
• Functional Equivalence via Distillation: The replacement is supervised through layer-wise distillation. For each layer $l$, a similarity loss ($L_{sim}$) minimizes the difference between the output of the original attention head and the new BiLSTM head:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  This ensures the substitute replicates the representational structure of the teacher.
• Training Protocol: The process involves two stages:
  1. Distillation Phase: Only the substitute BiLSTM modules are trained using a combination of similarity loss and task loss (classification), while the rest of the model is frozen.
  2. Finetuning Phase: The entire model is jointly finetuned to recover any residual accuracy drop.

2.2 Multi-head BiLSTM Architecture

To preserve the inductive bias of multi-head attention, the substitute is structured as a Multi-head BiLSTM:

• Projection and Splitting: The input embedding is projected and split into $N$ subspaces (matching the number of teacher attention heads).
• Bidirectional Processing: Each subspace is processed by a separate BiLSTM head, capturing both forward and backward token dependencies.
• Concatenation and Projection: Outputs are concatenated and projected back to the original dimension.
• IMC Compatibility: Unlike attention, this structure relies on recurrent weight reuse and sequential dataflow, eliminating all-pairs activation interactions. This maps naturally to ReRAM crossbars where weights are stationary and data flows sequentially.

2.3 Structured Pruning for IMC Constraints

To further adapt to resource-constrained IMC arrays, the authors apply structured pruning to the hidden dimensions of the BiLSTM heads:

• Extension of DeepHoyer: They extend the DeepHoyer framework to target LSTM hidden units. Because LSTM gates (input, forget, cell, output) share hidden units, pruning a single unit requires coordinated removal across all associated weight matrices ($W_{ih}$, $W_{hh}$, and projection weights).
• Group Sparsity: A differentiable Hoyer sparsity measure is applied to groups of weights corresponding to each hidden unit, encouraging the network to learn which units are redundant.
• Layer-Adaptive Compression: Pruning is performed independently per head and direction, allowing the model to adapt its capacity to the specific redundancy of each layer, resulting in compact modules that fit IMC array sizes.

3. Key Results

The framework was evaluated on the DeiT family (Tiny, Small, Base) on ImageNet and various downstream tasks (CIFAR-10/100, Cars, Flowers, iNaturalist).

• Accuracy Preservation:
  • DeiT-Tiny: FAR-Tiny actually outperformed the baseline DeiT-Tiny by +1.2% on ImageNet (73.3% vs 72.2%), suggesting the sequential substitute acts as a beneficial inductive bias for smaller models.
  • Larger Models: For DeiT-Small and Base, FAR maintained competitive accuracy with a modest gap of approximately 1–2% compared to the teacher, demonstrating that the functional role of attention is largely reproducible via sequential recurrence.
  • Downstream Transfer: FAR models retained strong transferability to downstream classification tasks, indicating that the learned sequence-to-sequence mappings generalize effectively.
• Robustness to Pruning:
  • Applying structured pruning reduced the average hidden dimension by 40–60% with only an additional 1–2% accuracy drop.
  • Pruning patterns were non-uniform, with higher retention in middle layers and more aggressive pruning at the network ends, reflecting internal heterogeneity in representational redundancy.
• Head-wise Specialization: Visualization of token interaction maps showed that FAR heads successfully inherited the distinct semantic roles of the original attention heads, concentrating on specific object regions rather than collapsing into redundant patterns.

Hardware Efficiency Analysis

Using an analytical ReRAM IMC model, the authors estimated latency and energy:

• Latency: FAR on ReRAM IMC is estimated to be ~18x faster than DeiT on ReRAM IMC (due to avoiding quadratic attention overhead) and ~400x faster than DeiT on a GPU.
• Energy: FAR on ReRAM IMC is estimated to be ~3x more energy-efficient than DeiT on ReRAM IMC and ~150x more efficient than DeiT on a GPU.
• Mechanism: These gains stem from eliminating activation-to-activation multiplications, reducing ADC/DAC conversions, and enabling weight-stationary execution. The authors note that excluding softmax costs (which are significant bottlenecks for attention on IMC) makes these estimates conservative for FAR.

4. Significance and Claims

The paper claims that FAR provides a practical, hardware-efficient alternative for deploying pretrained transformers on IMC-based edge accelerators without the need for training models from scratch.

• Feasibility of Full Replacement: The work demonstrates that attention can be replaced end-to-end at inference time while maintaining practical accuracy, challenging the notion that all-pairs token interactions are strictly necessary for high-performance vision tasks.
• Hardware Alignment: By substituting attention with sequential, weight-stationary modules, FAR aligns the computational dataflow with the physical strengths of ReRAM crossbars, significantly reducing memory traffic and latency.
• Preservation of Pretrained Knowledge: The framework proves that the rich semantic representations learned during large-scale pretraining can be preserved and transferred to structurally different, IMC-friendly architectures through block-wise distillation.

In summary, FAR offers a pathway to restructure transformer inference around IMC-friendly sequential modules, bridging the gap between the dominance of attention-based models and the constraints of emerging memory-centric hardware.