Spectral Attention Steering for Prompt Highlighting

🍞 Top Bread (Hook) You know how when you hand someone a worksheet and you use a highlighter to mark the most important sentence? You expect them to read that part carefully. But big language models (LLMs) don’t always notice your highlights in the same way, even if you put stars around the words.

🥬 Filling (The Actual Story)

• What the world looked like before: LLMs can read long prompts, but they often miss the exact parts people care about. If the prompt includes both helpful facts and distracting details, the model might grab the wrong thing. A classic failure is called “lost in the middle,” where models remember the beginning and end of long texts but forget the middle. People tried to fix this with a method called attention steering—nudging the model to look more at certain tokens. A popular method, PASTA, changed the attention score matrix after it was computed. It could work well, but it needed the whole attention matrix in memory, which is huge and slow.
• The problem: Modern efficient attention (like IO-aware, blockwise attention) avoids building the full attention matrix to save memory and time. But methods like PASTA need that full matrix to be edited later. So they become slow and memory-hungry, and often need expensive searches to pick which attention heads to change.
• Failed attempts: Post-hoc fixes either (1) require storing the entire giant attention table (bad for memory), or (2) adjust final outputs (logits) in rough ways that don’t truly guide focus (they can improve some cases but miss the deeper routing behavior of attention). They also often demand head-by-head searches to find where to steer—which is costly and brittle across tasks.
• The gap: We needed a way to steer attention without touching the full attention matrix, and without lots of manual tuning, while still being precise about which tokens deserve the spotlight.
• The real stakes: Think of reading a long email thread, legal document, or a medical note. If you highlight the important clause, the patient’s drug allergy, or the exact updated fact, you want the model to truly focus on it. Missing the highlight could mean wrong answers, wasted time, or even safety risks.

🍞 Bottom Bread (Anchor) Imagine you ask, “Previously, the cat was white. Now, the cat is black. What color is the cat?” If you highlight “black,” you want the model to answer “black” every time, even if it once learned the cat used to be white. This paper’s method makes that highlighting really count.

— New Concept Sandwich 1 — 🍞 Hook: You know how a teacher points to a word on the board so everyone looks right there? 🥬 The Concept: Attention Steering is a way to guide a model’s focus toward specific tokens in the prompt.

• How it works (steps):
  1. Mark the tokens you care about (the “highlights”).
  2. Adjust the model’s inner attention so queries look more strongly at those highlighted tokens.
  3. Let the model generate using this guided focus.
• Why it matters: Without attention steering, the model may treat critical and unimportant words almost the same and miss your highlight. 🍞 Anchor: When you highlight the updated fact “Kevin Garnett is a baseball player,” attention steering helps the model lock onto “baseball player,” not the old “basketball player.”

🍞 Top Bread (Hook) Imagine you wear glasses that make only the important text on a page look brighter. You still see the whole page, but your eyes are pulled to the useful parts automatically.

🥬 Filling (The Big Idea)

• The “Aha!” in one sentence: Instead of editing the big attention table after it exists, edit the key vectors before attention is computed so highlighted tokens light up more along data-driven directions.

Multiple analogies (3 ways):

1. Magnifying glass: You place the glass over special words so they look larger to the model.
2. Stage spotlight: You dim the background and brighten the actor you want the audience to watch.
3. Playlist boost: You turn up the volume only for your favorite songs without touching the rest.

Before vs. After:

• Before: Methods edited the attention matrix post-hoc, which costs memory, time, and depends on head searches.
• After: SEKA adjusts the keys first, meaning the attention naturally gives higher scores to highlighted tokens, staying fast and memory-friendly.

Why it works (intuition, no equations):

• Attention score is basically “How much does this query match that key?” If you change keys so the important tokens point more strongly in the ‘relevance’ directions, queries will match them better. That raises the attention to those tokens without needing to rewrite the entire attention table later.

Building Blocks (with Sandwiches) — New Concept Sandwich 2 — 🍞 Hook: Imagine every word in a sentence gets a tiny name tag that says what it’s about. 🥬 The Concept: Key Embeddings are the internal vectors that represent “what to look for” when other tokens decide whom to attend to.

• How it works (steps):
  1. The model turns each token into vectors (queries, keys, values).
  2. Keys act like labeled hooks; queries try to match to those hooks.
  3. Higher query–key match means more attention to that token.
• Why it matters: If you want the model to look at a specific token, shaping its key makes it easier to find. 🍞 Anchor: If the word “basketball” is highlighted, boosting its key makes the question token lock onto it more.

— New Concept Sandwich 3 — 🍞 Hook: Think of a secret storage room where you organize objects by hidden themes. 🥬 The Concept: Latent Space is the model’s hidden space where meanings and patterns live as directions.

• How it works (steps):
  1. Map tokens to vectors in a high-dimensional space.
  2. Directions in that space line up with behaviors (like “relevance to the question”).
  3. Moving along certain directions strengthens desired behavior.
• Why it matters: If “relevance” lives in a subspace, you can boost it precisely without messing up everything else. 🍞 Anchor: Sliding the “answer token” a bit more in the ‘relevance’ direction makes the model notice it.

— New Concept Sandwich 4 — 🍞 Hook: When you take apart a song into bass, drums, and vocals, you understand what’s driving the sound. 🥬 The Concept: Spectral Decomposition is a way to break data into principal directions that explain its most important variations.

• How it works (steps):
  1. Compute a matrix that captures how two sets of vectors vary together (cross-covariance).
  2. Use SVD to find top directions with the strongest shared signal.
  3. Keep top components to build a projection onto the “relevant” directions.
• Why it matters: Without finding strong directions, you’d amplify noise instead of true relevance. 🍞 Anchor: The method learns the main direction that separates “relevant” from “irrelevant” versions of the same text span.

— New Concept Sandwich 5 — 🍞 Hook: Imagine adding a small booster to a bicycle wheel so it spins faster only when you need it. 🥬 The Concept: SEKA (Spectral Editing Key Amplification) is a training-free method that edits key vectors along learned “relevance” directions to increase attention to highlighted tokens.

• How it works (steps):
  1. Offline, learn projection matrices that capture relevance directions using spectral decomposition from contrastive prompts.
  2. During inference, for highlighted tokens, add a small amplified projection of their key onto those directions.
  3. The attention mechanism naturally gives these tokens higher scores.
• Why it matters: It’s fast, memory-friendly, and makes highlighting actually work, even in long prompts. 🍞 Anchor: Highlight “They live in Berlin” and SEKA boosts the key of “Berlin,” so questions about where they live attend to it more.

— New Concept Sandwich 6 — 🍞 Hook: Think of a Swiss Army knife that picks the right tool based on the job at hand. 🥬 The Concept: AdaSEKA is an adaptive version that blends multiple expert relevance subspaces depending on the prompt’s query.

• How it works (steps):
  1. Learn several expert projections (e.g., for facts, instructions, multi-hop).
  2. Look at the prompt’s query vector and score how well it aligns with each expert’s top directions.
  3. Mix experts into one dynamic projector and apply it to highlighted keys.
• Why it matters: Different prompts need different kinds of focus; automatic routing reduces manual tuning across tasks. 🍞 Anchor: A prompt about overriding facts picks the “factual recall” expert more, while a pronoun-editing prompt picks the “instruction” expert.

Bottom Bread (Anchor) In practice, this turns your plain-text highlighting into a reliable focusing tool: when you bold or mark a phrase, the model actually pays extra attention to it during generation.

At a high level: Prompt with highlights → (A) Learn relevance directions offline → (B) Edit keys of highlighted tokens at inference → Output with boosted focus.

Step-by-step details Step 1: Build contrastive samples (offline)

• What happens: Create triplets where the same token span appears under positive (relevant question), negative (irrelevant question), and neutral contexts. Extract key embeddings for those spans across layers and heads.
• Why this exists: We need a supervision signal that cleanly separates “relevant” from “irrelevant” to discover the right directions in key space.
• Example: Context: “The portfolio manager allocates capital across equities and bonds.” Positive Q: “What does the portfolio manager allocate…?” Negative Q: “What does the climate model simulate?” The token “capital” is relevant in the positive case and irrelevant in the negative one.

Step 2: Compute cross-covariance and do spectral decomposition (offline)

• What happens: For each layer/head, compute cross-covariance matrices from neutral-with-positive and neutral-with-negative pairs. Apply SVD to get singular vectors/values. Choose top-k positive directions (strong relevance) and bottom-k negative directions (anti-relevance), forming projection matrices P+ and P−, with a threshold γ controlling how much variance to retain.
• Why this exists: SVD finds the most stable, data-driven axes that capture how keys shift when relevance changes. Without it, we’d push keys in random or noisy directions.
• Example: Suppose the top singular vector for a head aligns with the difference between positive and negative keys for many answer spans. Keeping it in P+ ensures we amplify the direction that makes “relevant” stand out.

Step 3: Select relevance-sensitive KV heads (offline analysis → runtime mask)

• What happens: Measure how much keys move between positive and negative prompts for each (layer, head). Keep only heads whose average movement (ℓ2 distance) exceeds a threshold δ_min.
• Why this exists: Not all heads do retrieval or relevance routing. Steering the wrong heads can add noise or harm performance.
• Example: In Qwen3 models, mid-to-late layers often show higher movement; early layers might not. We keep the movers.

Step 4: SEKA key editing at inference (runtime)

• What happens: For each highlighted token’s key k, compute k′ = k + g+·P+ k + g−·P− k. This adds a low-rank relevance boost before attention scores are computed.
• Why this exists: Editing keys upstream makes attention naturally prefer highlighted tokens, without building or editing the full attention matrix. Remove this step and highlighting barely changes model focus.
• Example (toy): If k projects 0.5 along the learned “relevance” direction and g+ = 0.2, then the edited key adds 0.1 along that axis, increasing the query–key match for that token.

Step 5 (AdaSEKA only): Dynamic expert routing (runtime)

• What happens: Train-free experts are learned offline. At inference, look at the last-token query per head, measure alignment with each expert’s top-K singular vectors (weighted by singular values), then mix experts into a single P_dynamic. Apply k′ = k + g·P_dynamic k for highlighted tokens.
• Why this exists: Different tasks need different relevance types. Automatic routing reduces manual hyperparameter fiddling per task/model.
• Example: If the query aligns 2× more with the “factual” expert than others, the resulting projector leans on factual directions more strongly.

Step 6: Integrate with efficient attention

• What happens: Register a lightweight hook that edits keys of only the highlighted tokens and only in selected heads, right before attention runs. This keeps compatibility with fast attention kernels.
• Why this exists: We avoid materializing or rewriting the full attention matrix, so latency and memory use stay low.
• Example: In tests with Qwen3-8B, SEKA adds about +0.03s per sample on long contexts, compared with +1.03s for post-hoc methods.

What breaks without each step

• Without contrastive data: You can’t learn clean relevance directions; performance drops and becomes unstable.
• Without SVD projections: Random projections help a bit but are clearly suboptimal; you may amplify noise.
• Without head selection: Steering every head can overwhelm the model and reduce accuracy.
• Without pre-attention editing: You lose FlashAttention compatibility and pay heavy memory/time costs.
• Without expert routing (for AdaSEKA): You must hand-tune gains per task/model, which is time-consuming and brittle.

Concrete walkthrough example Prompt: “Previously, Patrick Roy professionally plays hockey. Currently, Patrick Roy professionally plays basketball. Patrick Roy is a professional …”

• Highlight: “basketball”.
• SEKA edits the key vectors for the “basketball” tokens in relevance-sensitive heads.
• During attention, the question token’s queries match the boosted keys more strongly.
• Result: The model generates “basketball” as the profession, not “hockey,” consistently.

The secret sauce

• Targeted, low-rank key boosts along learned relevance directions give you strong control with tiny overhead.
• Pre-attention design keeps it compatible with optimized attention.
• Query-adaptive expert mixing in AdaSEKA personalizes steering to the prompt’s intent.
• Selective head steering focuses power where retrieval truly happens, avoiding collateral damage.

The test: What did they measure and why?

• They measured whether highlighting actually makes models focus and answer correctly under three scenarios:
  1. Knowledge conflicts (CounterFact): Can the model prefer the new fact in the prompt over its old memory?
  2. Occupation extraction (Bias in Bios): Can it pick the true job from the noisy biography?
  3. Instruction following (Pronoun changing): Can it follow a simple text transformation instruction while keeping content?
• They also tested the lost-in-the-middle setup to see if steering can boost recall for mid-position passages.

The competition: Compared against

• Original prompting (no steering).
• PASTA (post-hoc attention matrix editing; strong but heavy).
• SPA (logit-based steering at the output; lighter but not true attention routing).
• Ablations: SEKA with random projections; SEKA without head filtering.

The scoreboard, with context

• On CounterFact, SEKA and AdaSEKA routinely hit near-perfect Efficacy and Paraphrase Scores (e.g., ~99%) across Qwen3 sizes, outperforming the original model (often ~40–55%) and generally edging PASTA. That’s like jumping from a mid-grade to almost an A+.
• On Bias in Bios, SEKA/AdaSEKA usually land in the top two across model families. For example, with Qwen3-4B, Accuracy rises from ~80% to ~91%—a solid, reliable bump.
• On Pronoun Changing, results depend on how much the base model already responds to markdown-style marks. Qwen3 models do respond somewhat; still, AdaSEKA pushes to state-of-the-art (e.g., All-changed P. Score ~99.5%). On Gemma3-4B, which is less responsive to marks, SEKA brings especially large gains.
• Lost in the middle: Steering only the middle region flips the U-shape, turning the usual mid-context dip into a peak. Steering everything can slightly worsen the dip—showing the value of targeted, not blanket, steering.

Surprising or notable findings

• A simple “-marked” baseline can be strong on some models (like Qwen3), meaning certain models already treat formatting as a hint. Even so, AdaSEKA typically adds more gains.
• Random projections help some, but learned spectral projections plus head filtering are crucial. Removing both can tank performance (e.g., a dramatic drop on Pronoun Changing).
• Head sensitivity concentrates in mid-to-late layers—the same place mechanistic studies find retrieval heads. This alignment supports the paper’s selection strategy.

Efficiency and overhead

• SEKA adds about +0.03s per long sample and negligible extra memory—almost free.
• PASTA adds around +1.03s and large memory overhead because it needs the full attention matrix.
• AdaSEKA’s routing costs a bit more (~+0.27s) but remains far cheaper than post-hoc methods.

Takeaway

• Precise, pre-attention key editing works consistently across tasks and sizes.
• It’s not just accurate; it’s practical—fast, memory-friendly, and works with optimized attention.

Limitations

• Hyperparameter tuning: Gains g+/g− (or g), the head-selection threshold δ_min, and the variance threshold γ can influence results. Wrong settings can under-steer or over-steer.
• Model dependence: The best heads to steer and the stability of learned subspaces can vary across architectures and sizes.
• Data dependence: The quality and diversity of the contrastive samples matter. Poor samples may learn weak or noisy directions.
• Oversteering risk: Applying steering to too many heads or setting gains too high can reduce accuracy or harm generalization.

Required resources

• Storage for per-layer, per-head projection components (small compared to model size).
• A lightweight runtime hook to edit keys of highlighted tokens in selected heads.
• Optional expert banks (AdaSEKA) for different task types.

When not to use

• If your goal is to change the model’s style or long-chain reasoning semantics directly (activation steering may be better).
• If you don’t know which tokens to highlight (this method needs token indices to steer).
• If your setup forbids even tiny hooks into the attention module.

Open questions

• How universal are relevance subspaces across domains and languages? Do we need many small experts or a few big ones?
• Can routing be improved further using richer prompt signals (beyond last-token queries) while staying training-free?
• How does safety interact with attention steering—could malicious highlights bias models in harmful ways, and how do we guard against that?
• Can we jointly steer queries and keys to achieve finer control without losing efficiency?
• What is the best automatic way to decide which heads to steer for new model families with no manual analysis?

Three-sentence summary

• This paper introduces SEKA and AdaSEKA, training-free methods that steer a model’s attention by editing key vectors before attention is computed, making highlighted tokens truly stand out.
• By learning and applying spectral ‘relevance’ directions—and, for AdaSEKA, routing among multiple experts based on the prompt—these methods outperform strong baselines on multiple benchmarks.
• Crucially, they add minimal latency and remain compatible with modern fast attention, making them practical for long-context use.

Main achievement

• Turning prompt highlighting into reliable, efficient attention control via pre-attention key editing with learned spectral projections, avoiding the heavy costs of post-hoc attention matrix edits.

Future directions

• Explore broader and multilingual expert banks, smarter routing signals, and combined query–key steering.
• Study safety-aware steering and automatic guardrails for misuse.
• Integrate with retrieval-augmented systems to prioritize the most relevant passages in real time.

Why remember this

• SEKA/AdaSEKA show that a small, principled nudge at the right place (keys) can reshape attention powerfully and efficiently. It makes everyday “please focus here” prompts actually work—fast enough and well enough for real-world, long-context applications.