Steering Sparse Autoencoder Latents to Control Dynamic Head Pruning in Vision Transformers (Student Abstract)

1. Problem Statement

Vision Transformers (ViTs) utilize multi-head self-attention mechanisms, but many attention heads are redundant, leading to unnecessary computational costs without proportional performance gains. While Dynamic Head Pruning frameworks (e.g., AdaViT) exist to adaptively remove redundant heads based on input data, they suffer from two critical limitations:

• Opacity: The pruning decisions rely on dense residual embeddings, making the decision process a "black box" that is difficult to interpret.
• Lack of Control: Because the policies operate on dense, polysemantic representations, it is difficult to explicitly steer or control which heads are pruned for specific semantic concepts or classes.

2. Methodology

The authors propose a novel framework that integrates Sparse Autoencoders (SAEs) with dynamic head pruning to make the pruning process interpretable and controllable at the latent level. The methodology consists of three main stages:

A. Baseline Dynamic Pruning (AdaViT)

• Architecture: A ViT is trained with lightweight decision networks attached to each layer.
• Input: Each decision network receives the CLS token from the residual stream (encoding global context).
• Output: The network outputs logits ($a_\ell$) representing head importance.
• Masking: Binary masks ($M_{\ell,i}$) are generated via Gumbel-Sigmoid sampling to select which heads to keep, enforcing a target sparsity ratio while maintaining accuracy.

B. Sparse Autoencoder (SAE) Training

• Objective: To disentangle the dense, polysemantic CLS token embeddings into sparse, monosemantic features.
• Input: The CLS token ($x$) from the final layer of the trained ViT.
• Architecture:
  • Encoder: Projects the 384-d input into a 3072-d latent space ($8\times$ expansion).
  • Sparsity Enforcement: Uses a Top-K activation ($k=64$) to retain only the $k$ largest latent dimensions.
  • Decoder: Reconstructs the input ($\hat{x}$) from the sparse latent vector ($z$).
• Loss: Minimized Mean Squared Error (MSE) between the original input and reconstruction ($L_{rec} = \|x - \hat{x}\|^2$).

C. Latent Steering and Pruning Control

Once the SAE is trained, the framework modifies the pruning behavior by "steering" the latent space:

1. Amplification: Selected latent dimensions ($S$) are amplified by adding a strength factor $\alpha$:
   $$z'_i = \begin{cases} z_i + \alpha & \text{if } i \in S \\ z_i & \text{if } i \notin S \end{cases}$$
2. Re-sparsification: A Top-K operation is applied to the amplified vector to maintain sparsity ($\tilde{z}$).
3. Reconstruction: The steered latent vector is decoded back to the embedding space ($\tilde{x} = W_{dec}\tilde{z} + b_{dec}$).
4. Decision: The steered embedding $\tilde{x}$ is fed into the original decision network to generate a new pruning mask.

Steering Strategies: The authors tested three strategies for selecting the index set $S$:

• Per-class frequent: Top-$k$ latents most frequently activated for a specific class.
• Global frequent: Top-$k$ latents activated across all classes.
• Random: Randomly selected latents.

3. Key Contributions

• Interpretability via Disentanglement: The work demonstrates that SAEs can effectively disentangle dense ViT embeddings into sparse features that correspond to class-specific concepts.
• Controllable Pruning: It introduces a mechanism to explicitly control dynamic pruning decisions by amplifying specific latent dimensions, moving beyond opaque, data-driven heuristics.
• Class-Specific Head Discovery: The framework reveals that different classes rely on distinct, compact subsets of attention heads, and these subsets can be identified and optimized via latent steering.

4. Experimental Results

The framework was evaluated on CIFAR-100 using a ViT-Small (12 layers, 6 heads, 384-d) pre-trained on ImageNet.

• Baseline Performance: The standard dynamic pruning model achieved 89.79% accuracy while pruning 30% of heads (down from the baseline 91.27% on the full model).
• SAE Fidelity: Replacing the original embedding with the SAE reconstruction caused negligible performance loss (−0.12% accuracy), confirming the SAE preserves essential pruning information.
• Steering Efficacy:
  • Per-class Steering: As the amplification strength ($\alpha$) increased, this strategy significantly reduced head usage while maintaining or improving accuracy.
  • Comparison: Global and Random steering strategies resulted in significant accuracy drops as $\alpha$ increased, highlighting that class-specific features are distinct from global features (low overlap of 0.1641 between global and per-class top-$k$ latents).
• Case Studies (Final Layer):
  • "Bowl" Class: Accuracy improved from 76% to 82%, while head usage dropped from 0.72 to 0.33. The model relied heavily on heads h2 and h5.
  • "Pine Tree" Class: Accuracy improved from 79% to 84%, with head usage dropping from 0.93 to 0.35, relying on heads h2 and h3.
  • Semantic Relationships: Semantically related classes (e.g., "bowl" and "plate") shared similar head subsets (h2, h5), suggesting the SAE captures meaningful semantic relationships among attention heads.

5. Significance and Future Work

• Bridging Efficiency and Interpretability: This work successfully bridges the gap between computational efficiency (pruning) and mechanistic interpretability. It proves that pruning decisions are not just statistical artifacts but can be controlled via specific, interpretable latent features.
• Implications: The ability to steer pruning based on class-specific signals suggests a path toward more robust, explainable, and efficient ViT architectures.
• Future Directions: The authors plan to extend this framework to earlier layers of the network and apply it to foundation models (larger-scale models) beyond the current small-dataset (CIFAR-100) and final-layer focus.