Source-linked AI summary
Key Path Identification for Resolving Knowledge Conflicts via SAE-based Steering
Wenbo Zhang, Zhongxiang Sun, Zhiguang Han, Jun Xu
TL;DR
Knowledge conflicts expose weaknesses in mass SAE steering because correlation-based selection and ignored feature interactions introduce redundant features. KPI identifies causally critical features and constructs key paths for more selective steering. In RAG tasks with knowledge conflicts, it improves accuracy by 18% on average over the best mass-steering baseline.
Problem
Mass SAE steering can select redundant features because correlation-based identification ignores feature interactions, weakening precise steering for knowledge conflicts.
Method
KPI captures feature interactions, identifies causally critical key features, and constructs key paths for steering with fewer feature modifications.
Results
18% average accuracy improvement over the best mass-steering baseline was achieved in RAG tasks with knowledge conflicts.
Takeaways & Limitations
KPI provides a more precise and interpretable model-editing perspective by filtering redundant features and alleviating side effects.
Takeaways & Limitations
KPI requires continuous-layer SAEs, but satisfying SAE resources are unavailable for larger 70B+ models and training them from scratch is too costly.
Abstract
from arXiv · showhide
Sparse autoencoder (SAE)-based steering has been widely used to address knowledge conflicts by guiding LLMs to be more faithful to the contextual knowledge. Existing methods usually perform mass steering, which modifies a large batch of SAE features identified via correlation-based methods. However, due to the inaccurate correlation and the neglected feature interactions, mass steering methods fail to precisely identify the features that play the key roles in steering and introduce a large number of redundant ones, which add noise and weaken the steering effects. Our empirical studies reveal that steering only a small subset of the identified features can achieve comparable or even better performance. Motivated by this finding, we propose Key Path Identification (KPI), a novel method that identifies key steering features characterized by strong causal dependencies with both upstream and downstream features. From these features, KPI constructs key paths and steers through less feature modifications. In this way, KPI advances SAE-based steering from quantity-driven to quality-focused, offering a perspective for more precise and interpretable model editing. Experiments in RAG tasks with knowledge conflicts show that our method improves the accuracy by 18% on average compared to the best baseline of mass steering, effectively filtering redundant features, alleviating side effects and demonstrating the core role of key paths in steering.
1 Introduction
Knowledge conflicts can make RAG models favor parametric knowledge over contextual evidence, while correlation-based SAE steering selects redundant features and weakens control. KPI addresses this by identifying causally critical feature paths, achieving higher accuracy with fewer redundant modifications.
- Knowledge conflicts can cause models to produce outputs unfaithful to retrieved context by preferring inherent parametric knowledge.
- Correlation-based SAE steering treats features independently, selecting redundant features whose interactions dilute the impact of critical features.
- KPI identifies key features and causal paths through feature interactions to improve the precision and effectiveness of steering.
- 18% average accuracy improvement over the best mass-steering baseline was reported in RAG tasks with knowledge conflicts.
- The reported improvements are associated with precise SAE-feature localization and alleviated side effects.
2 Preliminary
SAEs encode model representations into sparse, higher-dimensional features and reconstruct them for feature-level steering. Correlation-based selection ranks features by mutual information with target behavior, but can retain many redundant features because it ignores feature interactions.
- SAEs encode model representations into sparse features z and reconstruct them as h_sae using learned encoder and decoder parameters.
- During steering, a selected feature’s decoder vector serves as its steering vector for modifying the model representation.
- Correlation-based identification: Mutual information measures correlation between an SAE feature’s activation and behavior selection, while positive features activate more during target than nontarget behavior.
- Correlation-based identification: Correlation-based methods rank features by mutual information and select the most correlated features using K.
- Correlation-based identification: More than 50 features can remain in one layer, and multi-layer steering retains redundant features with minor effects that weaken overall performance.
3 Redundancy Comes from Inaccurate Correlation and Neglected Interactions
Single-feature tests reveal that a small number of features can steer more effectively than mass steering, exposing inaccurate correlation rankings and harmful redundancy from ignored interactions.
- Some individual features outperform mass steering over all positive features, while most selected features produce relatively weak steering effects.
- In Gemma-2-9B layer 30, weak features can rank above the most useful feature, making correlation-based selection inaccurate.
- Increasing K to recover low-ranked useful features also includes more weak features, whose neglected interactions introduce redundancy and weaken steering.
4 Method
KPI identifies causally important SAE features from interaction patterns, ranks graph bottlenecks, and constructs cross-layer key paths to steer with fewer redundant modifications.
- KPI captures feature interactions, finds graph-based key features, and builds key paths across functional layers.
- Feature Interaction Pattern Capture: The method records activation increases in downstream features after singly steering top positive features, using a small development set to capture interaction patterns.
- Key Feature Finding: KPI represents features as nodes and causal interactions as directed edges, then uses indegree to identify convergence points for multiple paths.
- Key Path Building: Top-ranked key features form cross-layer interaction paths, but their individual steering effects vary and peak in particular layers.
- Key Path Building: KPI steers top-ranked features in a key layer and preceding layers because steering affects the current and subsequent layers.
5 Experiment Setting
Experiments evaluate SAE steering across three language models, two knowledge-conflict datasets, and multiple baselines, with additional analyses of selected-feature effects and KPI ablations.
- Models and SAEs: Experiments use Gemma-2-9B, Llama-3.1-8B, and Llama-3-8B with 131K-wide SAEs trained on residual streams.
- Datasets: NQ-Swap and Macnoise evaluate whether steering makes model outputs more faithful to substituted contextual knowledge.
- Baselines: STA and SPARE provide mass-steering SAE baselines, while ICL and CAD represent instruction-based and contrastive-decoding methods.
- Analyses: Figure 4 compares single-steering effects of selected key features for Llama-3.1-8B across two datasets.
- Analyses: Figure 5 reports a Gemma-2-9B KPI ablation in which KPI surpasses every single steering.
6 Experimental Results
KPI achieves the strongest reported performance while using fewer, more causally relevant features than mass-steering baselines. Analyses link its gains to improved feature selection, attention changes across layers, and reduced redundant-feature noise.
- 6.1 Main Results: KPI achieves the best performance across both datasets for every evaluated model.
- 6.1 Main Results: KPI still outperforms single steering despite the decreasing tendency of SAE-based methods on Llama-3.1-8B’s NQ-Swap results.The results indicate that SAE-based steering depends on SAE feature quality; SPARE also suffers from inaccurate layer selection on Gemma-2-9B Macnoise.
- 6.2 Feature Selection Analysis: KPI selects only 8 features, far fewer than mass-steering baselines, while achieving a higher average single-steering effect.
- 6.3 Side Effect Checks: KPI achieves comparable or better performance on faithful examples, indicating relatively small side effects from filtering redundant features.The side-effect evaluation tests examples without knowledge conflicts.
- 6.4 Attention Score Influence Result: Key features and other strong features consistently enhance attention to golden answer tokens in most heads of the next key layer, unlike a high-mutual-information feature with weak steering.The weak feature produces lower enhancements and more declines, illustrating the inaccuracy of correlation-based identification.
- 6.4 Attention Score Influence Result: Some strong features activate useful features rather than immediately increasing answer-token attention, with effects emerging in subsequent functional layers.These features can show more negative effects in the adjacent next layer but later enhance attention through model inferences.
- 6.5 Cosine Similarity Result: Cosine similarities and steering effects generally diminish with distance from the key layer, with the analyzed behavior-selection pattern forming mainly across layers 26-29.The analysis uses decoding-vector similarities to characterize feature interactions across layers.
7 Conclusion
The paper presents KPI as a precise steering method for resolving knowledge conflicts in RAG through causally critical features and key paths. Its experiments and mechanistic analyses report improved performance, fewer redundant modifications, reduced side effects, and layered formation of knowledge-selection behavior.
- Conclusion: KPI addresses redundancy from inaccurate correlation and neglected feature interactions by identifying causally critical features and key paths.
- Conclusion: KPI improves feature collaboration, filters redundant features, and alleviates side effects while steering with fewer feature modifications.
- Conclusion: Attention and feature-similarity analyses indicate that steering enhances attention to golden answer tokens and that knowledge-selection behavior forms gradually across multiple layers.
Limitations
The paper identifies several boundaries: KPI does not exclude backup circuits, requires cross-layer SAE resources that are scarce for 70B+ models, and assumes retrieved evidence should be followed in RAG.
- Basic Methods: The work does not deeply investigate denoising through data filtering.
- Multiple Circuits: KPI finds a simple circuit but does not exclude backup circuits with similar functions.Activating multiple circuits may instead provide no benefit because of saturation and mutual interference.
- Models and SAEs: The method requires continuous-layer SAEs, but satisfying resources are unavailable for many models larger than 70B.Training such SAEs from scratch is described as too costly.
- RAG Setting: The experiments assume retrieved evidence should be followed and focus on contextual faithfulness rather than broader reasoning ability.
D Results in More Complex Faithful Scenario
KPI is evaluated on FaithEval-counterfactual-v1.0, a more complex counterfactual setting, using steering settings derived from two simpler datasets. The reported results show consistent and significant improvement, supporting generalizability to out-of-distribution knowledge conflicts.
- Evaluation Setting: The evaluation tests whether settings derived from two simple datasets transfer to a scenario without factual context and corresponding answers.
- Results: KPI shows consistent and significant improvement on FaithEval-counterfactual-v1.0.The dataset uses completely designed contexts for counterfactual answers and is more complex than entity replacement.
E Side Effect Results on Other Model Capabilities
The paper evaluates side effects on MMLU and GSM8K, finding that mass-steering baselines can collapse under broader token steering while KPI maintains capability with small side effects.
- MMLU: Mass-steering baseline performance collapses when steering is applied to all tokens on MMLU.The comparison uses the all-token steering setting reported in Table 8.
- Side Effect Evaluation: KPI shows small side effects and maintains a high level of capability on additional model-capability evaluations.The evaluations use MMLU and GSM8K.
- Feature Selection: Selected-feature overlap remains stable across different layer top-k settings.
- Feature Selection: Top-k set to 10 gives the best reported performances while excluding many noise features and retaining potentially useful ones.This comparison uses Gemma-2-9B on Macnoise and Llama-3-8B on NQ-Swap.
F.3 Stability of Feature Interaction Pattern Capture with Different Numbers of Sampled Instances
The interaction patterns captured by KPI’s single-steering method remain stable as the number of sampled instances changes, supporting the use of small development datasets in the main experiments.
- Evaluation Setting: The overlap analysis varies the sampled-instance count while evaluating Gemma-2-9B on Macnoise and Llama-3-8B on NQ-Swap.
- Stability Analysis: Interaction patterns captured by single steering are quite stable across different numbers of sampled instances.
- Stability Analysis: The observed stability supports the reasonableness of the small development datasets used in the main experiments.
G Efficiency Analysis
KPI constructs an interaction graph by measuring how singly steered features increase downstream activations, then ranks and selects key paths for validation-based steering. Its graph construction costs 3.89 GPU hours under the reported Gemma-2-9B setting, using forward passes to improve efficiency.
- Runtime: KPI uses forward passes with single-feature steering to improve efficiency relative to backward-propagation approaches that require substantial computing and storage.The reported efficiency comes from using mutual-information rankings to identify important features before graph construction.
- Runtime: 3.89 GPU hours is the total graph-construction cost for Gemma-2-9B under the reported default setting.The measurement uses one NVIDIA A100-SXM4-40GB GPU, a development set of size 10, layers 18 to 36, and a maximum memory use of 36.45 GiB.
- Key feature ranking: Features are ranked jointly by interaction-graph indegree, producing a global ranking used to select a limited number of key features.The method defines an unweighted directed graph over candidate features and ranks features by decreasing indegree.
- Interaction graph: KPI estimates directed feature interactions by measuring mean activation increases from singly steered source features to later-layer target features.Candidate targets are restricted to the source layer or later layers, and the five largest positive increases are retained per target layer.
- Key-path selection: KPI selects a key layer and final key path using validation performance, then tunes the steering strength over a candidate coefficient set.Layer selection evaluates singly steered features, while coefficient tuning evaluates steering all features in the selected path.