PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Static Interpretability: Anticipating Post-SFT Mechanisms from Pre-SFT Parameters for Better Tuning

Hang Chen, Jiaying Zhu, Wenya Wang

arXiv:2608.24482v1cs.LGcs.AIcs.CL

TL;DR

Pre-SFT mechanistic localization can misidentify mechanisms needed after tuning, particularly for novel tasks. This paper predicts post-SFT mechanisms from current parameters and the target dataset by modeling SFT as a continuous trajectory and extrapolating Taylor-based sensitivity estimates. Experiments report superior SFT guidance, robust localization, and scalability across model sizes, while the authors identify limitations for multi-token generation and conflicting mechanisms.

  • Problem

    Pre-SFT localization may identify parameters that do not reflect mechanisms emerging during SFT, especially when the target task is novel.

  • Method

    The framework estimates ideal post-SFT mechanisms from current parameters and the target dataset by modeling SFT as continuous updates and extrapolating from an initial direction.

  • Results

    The method achieves state-of-the-art localization performance, superior SFT guidance, and robust scalability across diverse LLM families and increasing model sizes.

  • Takeaways & Limitations

    The work extends mechanistic interpretability from retrospective analysis toward predictive utility for targeted post-training optimization.

  • Takeaways & Limitations

    Current methods focus mainly on next-token mechanisms, limiting the locating-then-tuning paradigm for long-sequence instruction-following tasks.

Abstract

from arXiv · show

Mechanistic Localization bridges mechanistic interpretability and post-training optimization by isolating critical parameters via interpretative approaches and then guiding parameter-efficient Supervised Fine-Tuning (SFT) in a ``locating-then-tuning'' paradigm. However, due to the retrospective nature of mechanistic interpretability, directly interpreting pre-SFT models introduces misleading conclusions. Specifically for novel tasks, initially identified neurons differ drastically from those governing the final model, introducing biases that actively disrupt SFT. To address this, we propose a forward-looking localization framework that accurately estimates the post-SFT interpretability state using only pre-SFT parameters and the target dataset. Theoretically, we model SFT as a continuous parameter evolution, leveraging Taylor expansion to rigorously bridge the post-tuning mechanistic objective with the pre-SFT model's dynamic gradients. Practically, we design dual-granularity (neuron- and component-level) localization pipelines. Extensive experiments demonstrate that our approach not only provides superior SFT guidance but also exhibits robust performance and temporal scalability across increasing model sizes. This work transcends the fundamental limitation of traditional interpretability-its inability to identify task-critical mechanisms before they are trained-pioneering a predictive frontier that unites mechanistic interpretability with targeted optimization.

1 Introduction

The paper argues that static pre-SFT localization can misidentify task-critical parameters, especially for novel tasks, and proposes predicting post-SFT mechanisms from current parameters. It develops a lightweight estimation pipeline and reports strong localization, scalability, and SFT-guidance results.

  • Motivation: Pre-SFT localization can misidentify mechanisms that emerge during SFT, potentially freezing parameters needed for novel-task learning.This motivates replacing retrospective localization with a forward-looking estimate of post-SFT task-critical parameters.
  • Objective: The framework predicts neurons or components responsible for a target task in an ideal post-SFT model using only pre-SFT parameters and the target dataset.
  • Approach: SFT is decomposed into estimating the initial update direction and extrapolating the update distance toward the post-SFT state.
  • Efficiency: The estimation pipeline requires 1% of the training dataset for one epoch, with two forward passes and one backward pass for attribution data.
  • Results: Experiments on Mistral-7B, Llama-2-13B, and Qwen3-30B report superior SFT guidance, robustness to task difficulty, and scalability.
  • Results: The paper reports state-of-the-art localization performance across diverse LLM families with robust scalability and stable computational efficiency as model parameters increase.

2 Preliminaries

The preliminaries define the current and ideal post-tuning models, review causal localization methods, and explain why static pre-SFT analysis becomes biased during SFT.

  • Task Definition: The current model performs poorly on the target dataset, while the ideal post-tuning model achieves optimal target performance without degrading pre-existing capabilities.
  • Task Definition: Existing locating-then-tuning methods identify task-relevant parameter subsets and update only those parameters.
  • Task Definition: The paper reformulates localization to predict the critical parameter subset of the ideal post-SFT model rather than the current model.
  • Causal Localization: Causal localization measures how activating a parameter, neuron, or component affects an output-dependent metric on the target dataset.
  • Causal Localization: A larger ∆E indicates a stronger causal relationship with the task output, using clean and corrupted activations as the comparison states.
  • Attribution Patching: Attribution Patching approximates causal influence through a first-order Taylor expansion, replacing exhaustive counterfactual ablation with two forward passes and one backward pass.
  • Dynamic Interpretability: Mechanistic shifts in circuit strength, topology, and key nodes during SFT make localization based solely on pre-SFT parameters biased.

3 Method

The method models SFT as a continuous trajectory, estimates the initial gradient direction with a Taylor approximation, and extrapolates along that direction to predict ideal-state causal effects. Sampling step scalars accommodates the non-unique set of possible post-SFT configurations.

  • Problem: The central challenge is estimating future component importance from current parameters and the target dataset without executing full SFT.
  • Framework: SFT is modeled as a trajectory that first follows an initial gradient direction and then travels a distance toward a local optimum.
  • Single-Step Estimation: The intermediate state θ′ is obtained by an infinitesimal update of θ directed by the SFT objective on DT.
  • Single-Step Estimation: A first-order Taylor expansion around θ estimates θ′ causal effects using the original model parameters rather than evaluating the hypothetical θ′ model directly.
  • Multiple-Step Extrapolation: The transition from θ to θI is represented as an infinite sequence of continuous additive updates whose incremental effects accumulate along the parameter trajectory.
  • Multiple-Step Extrapolation: An exponential translation argument shows that sensitivity at the extrapolated state can be computed as S(θ + K∆θ), avoiding evaluation of intermediate states.
  • Multiple-Step Extrapolation: K is sampled across a range because the ideal post-SFT state is non-unique, and the mean causal effect across projected configurations estimates future mechanism importance.
  • Assumptions: The approximation assumes ∆θ captures the target task’s general optimization direction while remaining small enough for K to explore plausible parameter configurations.

4 Implemetation

The method uses a small probing SFT to estimate post-tuning mechanisms, then applies importance-guided tuning at neuron and component granularities. It compares direct localization on the probing model with the proposed extrapolated localization and reports increasing advantages as task difficulty rises.

  • Probing SFT: 1% of the training data and one epoch produce the probing model θ′ for estimating post-SFT parameter importance.The probing SFT uses 1% of the original learning rate.
  • Localization comparison: Direct attribution patching on θ′ serves as the Localization from θ′ baseline, while the proposed method estimates the ideal post-SFT state.The paper notes that θ′ remains close to θ and distant from the ideal model, requiring unrealistic trajectory assumptions for equivalence.
  • Localization comparison: As target SFT difficulty increases, localization from the estimated ideal state increasingly outperforms Localization from θ′.The comparison evaluates both localization paradigms empirically.
  • Neuron-level pipeline: The tuning pipeline supports neuron-level updates that retain only the top 20% most critical neurons and freeze the remaining 80%.For Mistral-7B, this updates 1.4B of 7B neurons using next-token cross-entropy loss.
  • Component-level pipeline: Component-level tuning assigns LoRA ranks dynamically from importance rankings, using ranks from 1 to 32 with an average rank of 8.The components are attention and MLP parameter matrices, including Wq, Wk, Wv, Wo, Wup, Wgate, and Wdown.

5 Experiments

Experiments evaluate neuron- and component-level localization across language, reasoning, arithmetic, and model-scale settings. The proposed method generally improves target-task performance and preserves general capabilities, while remaining robust as task difficulty and model size increase.

  • Evaluation Setup: Experiments cover GLUE-based NLU, BOOL logical reasoning, and Arithmetic subtasks ranging from 2 to 7 digits.Mistral-7B is the primary model, with LLaMA-2-13B and Qwen3-30B used for scalability validation.
  • Evaluation Setup: The evaluation compares neuron-level and component-level pipelines against gradient-guided and causal-effect localization baselines using TTA and PTA.TTA measures target-task accuracy, while PTA tracks retention of general capabilities.
  • Main Results: Our methods establish a substantial margin over baselines on both target-task accuracy and general-capability retention across LR, NLU, and MR domains.Arithmetic results in Table 1 average the 2- to 5-digit subtasks.
  • Main Results: Component-level localization significantly outperforms neuron-level localization, indicating that finer mechanistic granularity does not necessarily improve performance.CircuitLoRA is the second-best method, and the authors suggest LoRA’s adaptability to single, formalized tasks may contribute to this result.
  • Scalability: As model size increases, fine-tuning performance remains optimal while localization time grows prominently and overall computational complexity stays nearly lowest among evaluated baselines.The authors note that the gradual fine-tuning-time growth may partly reflect the simplicity of the target tasks, while localization complexity exceeds O(n).
  • Ablation Study: Table 2 shows that anticipating future parameter states provides greater foresight than relying on current parameters, while localization supplies interpretability-driven guidance.The ablation compares Full-Param, Random, Static, and Probing strategies.
  • Ablation Study: Across increasingly difficult arithmetic tasks, Static and Probing performance degrades sharply, whereas the proposed method preserves substantial robustness.The comparison spans 2- to 7-digit subtasks and 1 to 6 reasoning steps; Probing or Static may remain viable for rudimentary tasks under severe compute constraints.
  • Mechanistic Validation: Mechanistic analysis finds the estimated parameters produce circuits more closely aligned with the fully fine-tuned model and identify intermediate nodes encoding task subskills.Circuit graphs compare the initial, probing, and estimated ideal post-SFT states.

6 Conclusion and Limitations

The paper presents predictive pre-SFT localization for estimating post-SFT mechanisms and guiding targeted tuning, while identifying boundaries involving multi-token generation and conflicting mechanisms.

  • Conclusion: The framework predicts post-SFT mechanistic importance before SFT using base parameters and a probing model trained on 1% of the dataset.It supports targeted, fine-grained allocation of tuning intensity.
  • Limitations: Current interpretability methods focus mainly on next-token prediction, restricting locating-then-tuning to relatively narrow downstream applications rather than long-sequence instruction tasks.Extending localization across multiple forward passes remains an open challenge.
  • Limitations: Neuron polysemanticity means task-specific fine-tuning can degrade other capabilities, while existing localization identifies conflicting neurons without providing effective mitigation strategies.This limits the benefits of multi-objective joint optimization.

A The Derivation of Equation 2

The derivation approximates the causal effect of an infinitesimal parameter update using quantities computed under the current parameters, applying first-order Taylor expansions to activations and gradients.

  • Activation approximation: The derivation treats activations as differentiable functions of parameters and Taylor-expands clean and corrupted activations under the updated parameters.This produces a linear approximation of their activation difference.
  • Gradient approximation: The gradient under current parameters is Taylor-expanded for the updated parameters and rewritten using mixed second-order derivatives.The gradient term is expressed through parameter sensitivity of the activation-gradient relationship.
  • Final formulation: Substituting the activation and gradient approximations into Attribution Patching and omitting O((∆θ)^2) terms yields the final infinitesimal-update formulation.The resulting sensitivity captures how the target neuron’s attribution changes with parameter updates.

B Derivation of the Translation Operator for Sensitivity S(θ)

The appendix models sensitivity along the SFT update direction with a differential operator and folds its higher-order Taylor series into a translation operator, while showing why direct probing-model substitution is generally unsound.

  • Ideal-state estimation: The ideal-state sensitivity can therefore be estimated by scaling the update distance, yielding S(θI) = S(θ + K∆θ).This identity provides the theoretical bridge from pre-SFT parameters to the anticipated ideal state.
  • Operator construction: A differential operator D = ∆θ · ∇θ represents directional derivatives along the parameter-update direction, and repeated application generates higher-order terms.These terms correspond to Hessians, higher-order tensors, and beyond.
  • Translation identity: Folding the infinite Taylor series into an exponential translation operator shows that sensitivity at step k equals sensitivity evaluated after shifting parameters by k∆θ.The operator avoids explicitly computing the higher-order derivatives represented by the series.
  • Assumptions and setting: The derivation uses current and probing states’ activation differences and gradients, with an HVP assumption introduced to approximate gradient variation.The localization setting also constrains dataset labels to a single token and partitions arithmetic tasks by digit and step complexity.
  • Probing-model limitation: Directly substituting a probing model’s attribution requires activation differences to vary perfectly linearly across a macroscopic update, an assumption generally violated by nonlinear networks.The derivation therefore characterizes naive substitution as theoretically unsound.

E Experiment of Multi-Task Fine Tuning

The experiments evaluate multi-task fine-tuning, K selection, and probing-based prediction across varied task combinations and adaptation conditions. The proposed approach achieves high and stable performance while probing reduces the need for repeated K sampling.

  • Multi-task setup: The study constructs joint task sets containing 2 to 6 randomly combined tasks and evaluates mean and variance of Target Task Accuracy.At least three distinct combinations are used for each task-set size.
  • Multi-task results: Our method achieves the highest average multi-task accuracy and nearly the lowest variance across joint fine-tuning combinations.The result indicates limited severe inter-task conflict under the evaluated setup.
  • Conflict analysis: Our approach significantly minimizes conflicting circuit nodes during 0 to 400 fine-tuning iterations, surpassed only by unconstrained full-parameter fine-tuning.Baseline methods perform comparably to random localization and fail to reduce conflicting nodes.
  • K sampling: Tasks with near-zero initial TTA maintain high expected overlap across K ∈[10, 80], whereas higher-initial-TTA tasks have concentrated optimal K values.The values in the figure legend indicate initial Target Task Accuracy before fine-tuning.
  • Automated K prediction: A linear probe predicts task-specific K before tuning and matches or surpasses 10-sample averaging in circuit overlap and downstream TTA.Middle or deep residual-stream probes outperform first-layer probes for this prediction.

F.3 Robustness of Sampling K

The K-sampling analysis finds stable rankings for the most important components despite variation across sampled extrapolation values. This stability is explained by the dominance of gradient-dependent causal effects and the concentration of attribution in the top components.

  • Rank stability: Top-ranked components exhibit minimal rank variance across ten independent K samplings, while lower-ranked components vary more widely.Average sampled-component rankings range from the top 10 to beyond 2000.
  • Mechanistic explanation: The K-dependent extrapolation term is at least one order of magnitude larger than the baseline causal-effect term, so robust gradient direction governs rank stability.This supports identifying critical nodes across varying K values.
  • Attribution concentration: The top 10 components account for nearly 90% of the total causal effect, making lower-ranked rank fluctuations mathematically negligible.The concentration preserves overall localization efficacy despite sampling variation.

G Robustness and Sensitivity Ablation across Stochastic Factors

The framework is tested against stochastic variation from fine-tuning, probing, data subsampling, K sampling, and their joint combination. It maintains stable downstream performance and highly correlated critical-component rankings.

  • Perturbation factors: The robustness analysis isolates five stochastic factors: SFT seeds, probing samples, probing seeds, K sampling, and joint perturbation.These factors affect trainable initialization, probing states, extrapolation estimates, or combinations thereof.
  • Evaluation measures: Robustness is measured through Top@50 component-rank Spearman correlation ρ and downstream TTA coefficient of variation CV.Higher ρ indicates more invariant computational subgraphs, while lower CV indicates greater performance insensitivity.
  • Robustness results: CV ≤0.0018 and ρ ≥0.83, reaching 1.00 under varying SFT seeds, across isolated and joint perturbation settings.The reported results indicate near-zero downstream performance variance and consistently high Top@50 ranking correlation.

H Analysis from Mechanistic Interpretability

Mechanistic analyses compare predicted circuits with those from the fully fine-tuned surrogate across component and neuron levels. The proposed method most closely matches the ideal circuit, preserves structural fidelity under approximation, and concentrates effectiveness in top-ranked components.

  • Circuit alignment: Tables 5 and 6 compare component- and neuron-level circuit overlap and KL divergence against the full-parameter fine-tuned model.The proposed method yields circuits most closely aligned with the ideal model.
  • Approximation ablation: The mixed second-order approximation maintains superior structural fidelity to the ideal circuit compared with pure second-order computation across expanding edge capacities.Approximation introduces mild algebraic perturbations but preserves topology more effectively.
  • Causal-effect concentration: Top@50 components account for nearly 90% of total causal effect, so accurate ranking of prominent components chiefly governs framework efficacy.This concentration makes the method’s performance less dependent on lower-ranked components.
  • Circuit structure: The predicted IOI circuit contains more intermediate nodes and more intricate input-to-output pathways than circuits derived from θ or θ′.The structure is interpreted as anticipating richer task-specific mechanisms.
  • Probing hyperparameters: Reducing probing learning rate and epochs is evaluated through Top@50 overlap with the full-parameter fine-tuned model across probing SFT setups.Figure 12 reports the corresponding ablation comparison.
Loading 2608.24482v1…