[ICLR 2025] SEKA: Steering LLM Attention via Spectral Key Editing

TL;DR

Researchers from the University of Edinburgh and their collaborators have introduced Spectral Editing Key Amplification (SEKA). Unlike previous methods that "hack" the attention matrix after it's built, SEKA edits the key embeddings before the attention score is even calculated. This simple shift makes it incredibly fast, memory-efficient, and—crucially—compatible with FlashAttention.

Contextual Position: This work moves away from post-hoc attention manipulation (like PASTA) toward pre-attention representation engineering, effectively bridging the gap between "Activation Steering" and "Attention Control."

The Bottleneck: Why "Post-Hoc" Doesn't Scale

In the quest to make Large Language Models (LLMs) follow specific parts of a prompt (Prompt Highlighting), current SOTA methods like PASTA operate by scaling the attention scores. However, at a PhD level, we recognize a fundamental engineering conflict here: Modern LLM efficiency relies on FlashAttention, which fuses the attention kernel to avoid materializing the full $NimesN$ matrix in the GPU's High Bandwidth Memory (HBM).

If you want to edit the attention scores, you must materialize that matrix, destroying the $O(N)$ memory benefits and introducing massive latency. The authors found that PASTA increases inference time by over 1 second per sample—a non-starter for production.

The Insight: Relevance in the Key Subspace

The authors hypothesized that "relevance" is not just a score, but a direction in the latent space of Key ($K$) vectors. By analyzing contrastive prompts (Relevant vs. Irrelevant query for the same context), they visualized a consistent directional shift in certain attention heads.

Figure: Visualizing the structured shift in Key embeddings when token relevance changes.

The Methodology: SEKA & AdaSEKA

1. Spectral Decomposition (SVD): SEKA performs SVD on the cross-covariance matrices of key embeddings to identify the "relevance subspace."
2. Key Amplification: During inference, the key $k$ for a highlighted token is transformed as: $k^{\prime }=k+\frac{g^{+}{P}^{+}k+g^{-}{P}^{-}k}{2}$ Where $P$ is the projection matrix derived from the top singular vectors. This effectively boosts the signal of the highlighted token in the directions the model uses for retrieval.
3. AdaSEKA (Query-Adaptive): Not all prompts are the same. AdaSEKA uses a "routing" mechanism to blend different expert projections (e.g., one expert for factual recall, another for instructions) by measuring how well the current Query ($Q$) aligns with each expert's subspace.

Experimental Breakthroughs

The performance metrics are striking, particularly in tasks involving In-Context Knowledge Conflict (CounterFact) and Instruction Following (Pronoun Changing).

1. Inverting the "Lost-in-the-Middle" Curve

Standard LLMs suffer from a U-shaped performance curve, forgetting information in the middle of long contexts. By applying SEKA specifically to the middle passages, the authors managed to invert the curve, turning the performance trough into a peak.

2. Efficiency Benchmark

SEKA achieves its gains with near-zero cost. While PASTA adds 23 GB of peak memory overhead and +1.03s latency, SEKA stays within 0.03 GB and +0.03s of the original model.

Critical Perspective & Takeaway

Why it works: From a mechanistic interpretability standpoint, SEKA targets the "Routing Subspace" (Q/K interaction) while leaving the "Semantic Subspace" (Value/MLP) untouched. This surgical precision avoids the "collateral damage" common in logit-based or full-activation steering.

Limitations: The method still requires a "selection" of sensitive KV heads. While the authors provide a threshold-based heuristic (${\delta }_{min}$), the optimal heads vary between model families (e.g., Qwen vs. Gemma), suggesting that the "retrieval circuit" is architecture-dependent.

Conclusion: SEKA proves that we don't need to rebuild or even fully see the attention matrix to control it. By understanding the geometry of the Key space, we can steer LLMs at the speed of FlashAttention.