PaperlayerView exact revision on arXiv

Source-linked AI summary

Spectral Attention Steering for Prompt Highlighting

Weixian Waylon Li, Yuchen Niu, Yongxin Yang, Keshuang Li, Tiejun Ma, Shay B. Cohen

arXiv:2603.01281v1cs.CLcs.AI

TL;DR

Existing attention-steering methods require storing the full attention matrix, limiting compatibility with memory-efficient attention. SEKA and AdaSEKA instead edit key embeddings before attention, achieving strong benchmark performance with negligible overhead and optimized-attention compatibility.

  • Problem

    Existing attention-steering methods require explicit full-attention-matrix storage, limiting compatibility with memory-efficient implementations such as FlashAttention.

  • Method

    SEKA uses spectral relevance subspaces to edit key embeddings before attention, while AdaSEKA query-adaptively blends task-specific expert projections through training-free routing.

  • Results

    Both methods generally outperform baselines and existing methods across tasks and model scales, with SEKA adding ≈0.03s latency per sample versus PASTA’s +1.03s.

  • Takeaways & Limitations

    SEKA and AdaSEKA provide training-free, efficient attention steering compatible with optimized implementations for more controllable and efficient long-context applications.

  • Takeaways & Limitations

    Standard SEKA requires hyperparameter tuning across tasks and model families because its projections are static, motivating AdaSEKA’s adaptive selection.

Abstract

from arXiv · show

Attention steering is an important technique for controlling model focus, enabling capabilities such as prompt highlighting, where the model prioritises user-specified text. However, existing attention steering methods require explicit storage of the full attention matrix, making them incompatible with memory-efficient implementations like FlashAttention. We introduce Spectral Editing Key Amplification (SEKA), a training-free steering method that tackles this by directly editing key embeddings before attention computation. SEKA uses spectral decomposition to steer key embeddings towards latent directions that amplify attention scores for certain tokens. We extend this to Adaptive SEKA (AdaSEKA), a query-adaptive variant that uses a training-free routing mechanism to dynamically combine multiple expert subspaces based on the prompt's semantic intent. Our experiments show both methods significantly outperform strong baselines on standard steering benchmarks while adding much lower latency and memory overhead, in compatibility with optimised attention.

1 INTRODUCTION

The paper introduces SEKA, a training-free method that steers attention by editing key vectors before score computation, and AdaSEKA, which dynamically blends task-specific experts according to semantic intent. Both methods improve standard steering benchmarks while remaining compatible with FlashAttention and adding negligible overhead.

  • SEKA: SEKA directly modifies key vectors before attention scores are calculated, replacing attention-output intervention with input-side steering.It learns a universal relevance subspace through spectral decomposition of key embeddings from contrastive prompts.
  • AdaSEKA: AdaSEKA blends task-specific expert projections through a cheap, training-free routing mechanism conditioned on the prompt’s semantic intent.Example experts target factual recall and instruction following.
  • Implementation: SEKA operates directly on key embeddings and is fully compatible with FlashAttention with negligible computational overhead.The reported implementation avoids editing the full attention output.
  • Experiments: SEKA and AdaSEKA achieve superior results on standard benchmarks covering knowledge conflicts, occupation extraction, and instruction following.AdaSEKA’s query-adaptive routing dynamically tailors steering to each prompt’s semantic intent.

2 PROBLEM DEFINITION AND MOTIVATIONS

This section formulates prompt highlighting as increasing attention to user-specified tokens and motivates pre-attention key editing as a memory-efficient alternative to post-computation steering. Contrastive analyses indicate that contextual relevance produces structured directional shifts in key representations, supporting projection-based steering.

  • Problem Definition: Prompt highlighting steers queries to assign increased focus to a designated subset H of tokens within a T-token prompt.The standard unnormalised attention score is Attn(i, j) = q_i^⊤k_j/√d_k.
  • Objective: The objective is to add a controllable term Δ_ij to attention scores for highlighted-token indices, yielding A′_ij = A_ij + Δ_ij.The additive term is designed to selectively boost attention toward user-specified highlighted tokens.
  • Motivation: Existing post-computation methods such as PASTA and positional calibration require explicit full attention-matrix storage, limiting compatibility with FlashAttention.PASTA may also require costly head search to determine which attention heads to steer.
  • Empirical Motivation: Contrastive positive-versus-negative prompts reveal robust, consistent directional shifts in key representations for certain attention heads.The analysis uses synthetic neutral, positive, and negative prompt triplets, extracts corresponding key embeddings, and visualises paired shifts with PCA.
  • Pre-Attention Steering: These structured relevance shifts motivate editing keys before attention via k′_j = k_j + gP k_j, where P defines a per-head relevance subspace and g controls scaling.The intervention preserves compatibility with efficient attention implementations and provides a geometrically interpretable steering mechanism.

3 SPECTRAL ATTENTION STEERING FOR PROMPT HIGHLIGHTING

SEKA steers prompt highlighting by editing key embeddings along spectrally learned relevance directions before attention computation, while AdaSEKA routes queries across expert projections. The framework also selects relevance-sensitive KV heads to target steering where prompt-dependent representation shifts are strongest.

  • 3 SPECTRAL ATTENTION STEERING FOR PROMPT HIGHLIGHTING: SEKA and AdaSEKA achieve prompt highlighting by directly editing key embeddings before attention computation, adapting spectral activation steering to token-wise attention control.SEKA is the base method, while AdaSEKA is its query-adaptive variant.
  • 3.1 SPECTRAL LEARNING OF RELEVANCE-ALIGNED PROJECTIONS (OFFLINE): Learned projections are stored per layer and head, enabling fine-grained steering at inference time.The projections are learned from token-level key embeddings obtained from synthetic contrastive prompts.
  • 3.1 SPECTRAL LEARNING OF RELEVANCE-ALIGNED PROJECTIONS (OFFLINE): Positive projections use top singular vectors associated with the largest singular values, whereas negative projections use least-significant vectors targeting directions least associated with relevance.The threshold γ controls retained variance when constructing the projection matrices.
  • 3.2 SPECTRAL EDITING FOR HIGHLIGHTED TOKENS (INFERENCE): During inference, SEKA injects learned projections into key embeddings before attention scores are computed, with independently adjustable positive and negative steering gains.The edited keys modify attention logits through a key-dependent relevance term.
  • 3.2 SPECTRAL EDITING FOR HIGHLIGHTED TOKENS (INFERENCE): SEKA is equivalent to adding a low-rank relevance bias to attention scores while avoiding direct attention-matrix manipulation by editing key vectors instead.This provides a structured and interpretable mechanism operating entirely on key representations before attention computation.
  • 3.3 VARIANT: QUERY-DRIVEN ADAPTIVE SEKA: AdaSEKA automatically selects and combines expert projections using query-specific relevance signals, reducing manual configuration across tasks and model families.Each expert learns projections from a distinct domain-specific dataset, and four expert projections are derived in the described implementation.
  • 3.3 VARIANT: QUERY-DRIVEN ADAPTIVE SEKA: AdaSEKA routes using the last prompt token’s query, whose alignment with expert projection directions determines dynamic coefficients and the weighted combination of expert projections.The last token serves as the prompt’s global aggregator; routing preserves the sign of alignment after normalization.
  • 3.4 SELECTING RELEVANCE-SENSITIVE KEY-VALUE HEADS: SEKA is most effective when applied selectively to KV heads sensitive to prompt relevance, with larger relevant–irrelevant shifts consistently emerging in mid-to-late layers.Early layers and some heads show minimal shifts, suggesting retrieval behavior is less likely there.

4 EXPERIMENTAL SETUP

The experiments evaluate SEKA and AdaSEKA on standard PASTA attention-steering benchmarks and an additional lost-in-the-middle positional-recall setting. Comparisons use direct prompting, highlighted-context prompting, PASTA, and SPA, with exact match under a 60-token generation limit.

  • Standard Benchmarks: Standard benchmarks cover Pronouns changing, Bias in Bios, and CounterFact, following PASTA’s benchmark selection and highlighted-token conventions.Bias in Bios and Pronouns changing derive from BIASBIOS, while CounterFact uses the COUNTERFACT dataset.
  • Baselines and Methods: The benchmark compares direct prompting, marked highlighted-context prompting, SEKA, AdaSEKA, PASTA, and Selective Prompt Anchoring (SPA).Markers identify steering token indices for attention-steering methods and are removed from the input IDs.
  • Lost-in-the-Middle Experiment: The lost-in-the-middle experiment tests whether highlighting middle passages can invert the usual U-shaped recall curve into a peak.The setting targets degraded recall when relevant information appears in the middle of long contexts, unlike the beginning and end.
  • Metrics: Evaluation uses exact match (EM), counting predictions containing the ground-truth short answer span, with answers limited to a maximum of 60 tokens.The metric follows Liu et al. (2024) and limits verbose or off-topic completions.
  • Lost-in-the-Middle Experiment: SEKA is tested on either passages 4 through 25 or all context passages in the positional-recall experiment.The comparison includes direct prompting of the base LLM and PASTA; Hsieh et al. (2024) is excluded because its code is unavailable.

5 RESULTS

SEKA and AdaSEKA generally outperform baselines and existing steering methods across models and tasks, while attention-level steering is especially effective for factual recall and less emphasis-responsive models. Ablations and lost-in-the-middle experiments further validate the spectral and head-selection components and show that selective steering can reshape positional performance.

  • Overall results: SEKA and AdaSEKA generally rank among the top two and outperform baseline and existing methods across tasks and model scales.The experiments use Qwen3 models from 4B to 14B and Gemma3 models from 4B to 12B.
  • Overall results: 99.02 and 97.16 are the Qwen3-4B CounterFact scores for SEKA and PASTA, respectively, while SPA lags considerably.These near-perfect scores support attention-level interventions for factual recall and knowledge conflicts.
  • Overall results: AdaSEKA reaches an A. P. Score of 99.52 on Qwen3-8B Pronoun Changing and is especially valuable for smaller models less responsive to markdown emphasis.The “**-marked” baseline is notably strong for Qwen3, whereas Gemma3-4B is less responsive to markdown emphasis.
  • Overall results: Removing learned projections and head filtering drops the Qwen3-4B Pronoun task A. P. Score from 90.52 to 36.95.Random projections with head filtering remain beneficial but clearly suboptimal, highlighting the value of spectral learning.
  • Lost-in-the-middle analysis: Applying SEKA to middle passages substantially increases central exact-match scores and can eliminate the lost-in-the-middle performance trough.Applying SEKA uniformly across all passages is also evaluated, while selective middle-passage steering is identified as highly effective for inverting the U-shaped profile.
  • Lost-in-the-middle analysis: Decreasing δmin increases the number of steered KV heads, with thresholds 0.16, 0.165, 0.17, and 0.18 selecting 58, 48, 41, and 31 heads for Qwen3-8B-Base.The control experiment fixes g+ at 0.2 and g− at 0.1 while varying δmin.

6 OVERHEAD ANALYSIS

The section evaluates inference overhead for the pre-computation approach, emphasizing compatibility with optimized attention mechanisms such as FlashAttention. Measurements use 100 Section 5.2 samples with a Qwen3-8B-Base model on one NVIDIA-GH200-120GB GPU.

  • Compatibility: Pre-computation is compatible with optimized attention mechanisms such as FlashAttention.The passage identifies this compatibility as a key advantage of the approach.
  • Evaluation setup: Overhead is measured on 100 samples averaging 4362 tokens from Section 5.2.The evaluation uses inference samples drawn from the paper’s Section 5.2 benchmark setting.
  • Evaluation setup: The evaluation uses Qwen3-8B-Base on a single NVIDIA-GH200-120GB GPU.Table 3 reports time per sample and average peak memory usage for this model and hardware configuration.

7 RELATED WORK

Related work divides LLM steering into activation steering, which controls high-level semantic outputs through MLP interventions, and attention steering, which directs focus to specific prompt tokens. Prompt-highlighting methods in the latter category intervene on attention scores or final logits but can add latency and conflict with FlashAttention.

  • Steering paradigms: LLM steering comprises activation steering for high-level semantic outputs and attention steering for directing focus to specific input tokens.Activation steering intervenes in MLP layers, whereas attention steering targets token-level focus within the prompt.
  • Activation Steering: Activation steering, also called representation engineering, adds steering vectors to MLP activations to control semantic attributes.Reported applications include honesty, safety, style, reasoning, and knowledge editing.
  • Attention Steering: Prompt-highlighting methods address failures to attend to key information by editing attention scores or final logits after computation.These interventions can introduce significant latency, and full attention-matrix editing is incompatible with optimisations such as FlashAttention.

8 CONCLUSION

The paper introduces SEKA and AdaSEKA as training-free attention-steering methods that modify key embeddings before attention computation. This pre-attention design addresses prior efficiency limitations while remaining compatible with optimised implementations and achieving state-of-the-art results on standard benchmarks.

  • 8 CONCLUSION: SEKA and AdaSEKA form a new class of training-free attention-steering methods that modify key embeddings before attention computation.AdaSEKA is described as the adaptive variant of SEKA.
  • 8 CONCLUSION: The pre-attention approach overcomes prior efficiency limitations and is fully compatible with optimised attention implementations.Key embeddings are edited before the attention computation rather than requiring explicit storage of the full attention matrix.
  • 8 CONCLUSION: Both methods achieve state-of-the-art results across a range of standard steering benchmarks.The conclusion reports experimental confirmation of these results but supplies no numerical values in the provided passage.

REPRODUCIBILITY STATEMENT … E.3 PRONOUNS CHANGING

The paper documents its reproducibility materials, constructs synthetic and task-specific data for relevance supervision and AdaSEKA, and explains SEKA’s targeted geometric transformation and algorithms. It evaluates the methods on adapted benchmarks, including COUNTERFACT, BIASBIOS, and Pronouns Changing, with an enhanced pronoun-sensitive metric.

  • REPRODUCIBILITY STATEMENT: All code, experiment instructions, projection matrices, and pre-processed datasets are publicly available, while original evaluation datasets, metrics, and hyperparameters are documented in cited sections and appendices.The repository is available at https://github.com/waylonli/SEKA; BIASBIOS, COUNTERFACT, and “Lost in the Middle” are distributed under the MIT License.
  • A SYNTHETIC DATASET FOR TOKEN-LEVEL RELEVANCE SUPERVISION: The synthetic supervision dataset contrasts two context–question–answer settings so identical token spans can receive relevance labels from aligned, mismatched, and neutral prompts.Each answer span yields positive h+, negative h−, and neutral h key embeddings; samples are automatically generated with GPT-4o without manual annotation.
  • B MULTI-EXPERT PROJECTION LEARNING SAMPLES FOR AdaSEKA: Four datasets provide neutral–positive prompt pairs for multi-expert projection learning in AdaSEKA, including training-derived pairs from Counterfact and BiasBios.The supplied passage states that three task-specific datasets were added to the synthetic dataset, making four total.
  • C GEOMETRIC INTUITION OF THE SEKA TRANSFORMATION: SEKA decomposes each key into relevance-subspace and orthogonal components, scaling the aligned component by 1 + g while preserving orthogonal directions.With g = 1, the relevance-aligned component is doubled; the operation is a targeted linear modification rather than a projection or orthogonal transformation.
  • C GEOMETRIC INTUITION OF THE SEKA TRANSFORMATION: SEKA modifies the attention-routing relevance subspace rather than semantic representations, consistent with the separation of Q/K routing from V/MLP semantic writes.The supplied passages characterize Q/K vectors as implementing token matching and algorithmic routing, with semantic information primarily in values and MLP activations.
  • D SEKA AND AdaSEKA ALGORITHMS: Standard SEKA learns fixed positive and negative projection matrices offline using SVD and applies them to highlighted-token keys through an inference hook, while AdaSEKA dynamically combines expert projections.The algorithms describe static projections for SEKA and query-driven dynamic projections for AdaSEKA before returning edited keys to attention computation.
  • E DETAILS OF STANDARD BENCHMARKS: The evaluation uses three established PASTA-adapted benchmarks with case-insensitive scoring, omitting JSON Formatting because modern models achieve near-perfect performance.COUNTERFACT tests prioritization of new contextual facts over pre-trained associations, BIASBIOS predicts one of 28 occupations, and Pronouns Changing tests linguistic transformation.

F TECHNICAL SETUP … I.3 PRONOUNS CHANGING EXAMPLES

The appendix specifies SEKA and AdaSEKA’s tuning procedures, examines SEKA’s sensitivity and attention effects, compares PASTA configurations, and provides qualitative examples across benchmark tasks. Together, these sections describe implementation choices, mechanistic behavior, runtime tradeoffs, and task-specific steering outcomes.

  • F TECHNICAL SETUP: SEKA tunes four hyperparameters, whereas AdaSEKA tunes only the KV-head selection threshold and one steering gain coefficient.CounterFact and Bias in Bios use a 500-sample validation grid search, while Pronoun Changing uses a separate development set.
  • F TECHNICAL SETUP: δmin and g+ are most influential because they control which heads are steered and amplification strength, with extremes degrading performance.Performance declines when too few or too many heads are selected, or when positive gain is too small or too large.
  • F TECHNICAL SETUP: SEKA’s sensitivity patterns are similar within model families but vary across tasks and models, with Gemma 3 showing greater variance with respect to γ.Qwen3-4B and Qwen3-8B both favor low δmin and remain stable across γ on CounterFact.
  • G MECHANISTIC INSIGHT VIA ATTENTION VISUALISATION: SEKA shifts attention toward manipulated prompt content in the CounterFact visualization, increasing focus on information relevant to the edited fact.The example compares layer-wise mean attention before and after steering in Qwen3-4B-Base.
  • H COMPLETE RESULTS OF PASTA WITH DIFFERENT CONFIGURATIONS: PASTA’s configurations differ substantially in head-selection runtime, from 1–2 minutes for SEKA-based methods to ≈2 hours for the original search routine.The compared approaches include SEKA’s KV-head configuration, transformed attention-head configuration, and PASTA’s original head search.
  • I.1 COUNTERFACT EXAMPLES: SEKA steers models toward newly provided factual information, overriding pretrained knowledge in CounterFact examples.Tables 9 and 10 illustrate examples involving an overridden location and job title.
  • I.2 BIAS IN BIOS EXAMPLES: SEKA focuses attention on the correct highlighted sentence in biographies while ignoring distracting information in Bias in Bios examples.The examples demonstrate both correcting a prediction and increasing confidence in the correct answer.
  • I.3 PRONOUNS CHANGING EXAMPLES: SEKA improves instruction following in some Pronoun Changing cases, although original models generally exceed 0.9 P. Score.Tables 13 and 14 show examples of specific failures that SEKA can address.

J PROJECTION SAMPLE EFFICIENCY ANALYSIS

SEKA’s projection quality is generally data-efficient, with performance typically stabilising after roughly 50 synthetic samples, although thresholds vary by task, architecture, and model size. Additional samples mainly improve stability by reducing variance, while model families and sizes exhibit distinct sample-efficiency patterns.

  • J PROJECTION SAMPLE EFFICIENCY ANALYSIS: Performance typically stabilises once roughly 50 synthetic samples are used, with the threshold varying by task, architecture, and model size.This analysis evaluates end-to-end performance across the three standard benchmark tasks using SEKA projections extracted from different numbers of synthetic samples.
  • J PROJECTION SAMPLE EFFICIENCY ANALYSIS: More samples do not always improve peak performance, but larger sample sizes consistently reduce variance and produce more stable behaviour.With only a few samples, projections may overfit synthetic pairs and introduce unpredictable variance even when accuracy plateaus.
  • J PROJECTION SAMPLE EFFICIENCY ANALYSIS: Within model families, Qwen3 CounterFact stabilises relatively early, whereas Gemma3 models—especially Gemma3-12B—require more samples for the same task.BiasBios and Pronouns Changing generally stabilise faster across most settings.
  • J PROJECTION SAMPLE EFFICIENCY ANALYSIS: Model size still creates noticeable sample-efficiency differences despite family-level similarities, with Qwen3-8B identified as the clearest example.The passage reports both shared behaviour patterns within families and meaningful differences associated with model size.

K COMPLETE RESULTS FOR δmin THRESHOLD ON LOST-IN-THE-MIDDLE

The optimal KV-head selection threshold δmin varies with model size. Figure 10 examines how different δmin thresholds affect SEKA performance on lost-in-the-middle for Qwen3-4B and Qwen3-14B.

  • The optimal KV-head selection threshold δmin can vary with model size.This variation is noted as relevant to threshold selection.
  • The threshold comparison covers both Qwen3-4B and Qwen3-14B models.The figure illustrates the effect of varying δmin for these two model sizes.
  • Figure 10 varies δmin thresholds when applying SEKA to the middle region of the lost-in-the-middle task.The figure reports exact match scores under these threshold settings.
Loading 2603.01281v1…