MoRFI: Monotonic Sparse Autoencoder Feature Identification

MoRFI: Monotonic Sparse Autoencoder Feature Identification

Abstract

Large language models (LLMs) acquire most of their factual knowledge during the pre-training stage, through next token prediction. Subsequent stages of post-training often introduce new facts outwith the parametric knowledge, giving rise to hallucinations. While it has been demonstrated that supervised fine-tuning (SFT) on new knowledge may exacerbate the problem, the underlying mechanisms are still poorly understood. We conduct a controlled fine-tuning experiment, focusing on closed-book QA, and find latent directions that causally contribute to hallucinations. Specifically, we fine-tune Llama 3.1 8B, Gemma 2 9B and Mistral 7B v03 on seven distinct single QA datasets, controlling for the percentage of new knowledge and number of training epochs. By measuring performance on the test set, we validate that incrementally introducing new knowledge increases hallucinations, with the effect being more pronounced with prolonged training. We leverage pre-trained sparse autoencoders (SAEs) to analyze residual stream activations across various checkpoints for each model and propose Monotonic Relationship Feature Identification (MoRFI) for capturing causally relevant latents. MoRFI filters SAE features that respond monotonically to controlled fine-tuning data mixtures of a target property. Our findings show that exposure to unknown facts disrupts the model’s ability to retrieve stored knowledge along a set of directions in the residual stream. Our pipeline reliably discovers them across distinct models, recovering knowledge through single-latent interventions.

Introduction

The paper investigates whether integrating new knowledge through fine-tuning leaves a detectable signature in the activation space. The authors conduct a controlled fine-tuning experiment on closed-book QA and study how internal activations differ across two dimensions: the proportion of training samples unknown to the pre-trained model, and the number of training epochs. The approach differs from prior work in several respects: "we select latents not by activation differences between two model states or within a single model but through bootstrapped monotonic trend detection across a controlled gradient of fine-tuning conditions, tracing how the model’s internal knowledge representations shift across fine-tuning conditions rather than contrasting two fixed states." The fine-tuning conditions vary epistemically rather than thematically, with topical content remaining diverse across relations and entities. The paper considers both increasing and decreasing latents, the latter of which prove disproportionately impactful for knowledge recovery.

The contribution is dual: studying the mechanisms of knowledge integration and developing MoRFI, a general-purpose algorithm for identifying SAE latents whose activations exhibit robust increasing or decreasing monotonic trends along a specified dimension of variation. The method combines bootstrapped resampling with statistical testing to ensure that identified latents exhibit consistent directional trends rather than transient or spurious fluctuations, providing a principled feature selection procedure validated through activation steering.

Methodology

The core of the approach is a 4D tensor A that captures SAE activations across four dimensions: (i) N input samples; (ii) P distinct dataset property configurations; (iii) F latent feature dimensions, and (iv) T timesteps corresponding to fine-tuning epochs. The property dimension P is designed to systematically source latent features whose activation magnitudes exhibit a strong, monotonic relationship with the target property. For a particular fine-tuning dataset D and a pre-trained LLM M, MD denotes a model obtained by fine-tuning M on D. MDp is defined as the model obtained by fine-tuning M on a strictly controlled version of D curated to contain exactly p% of a target property.

The MoRFI algorithm processes tensor A along a target dimension of aggregation and outputs a robust subset of features whose activation magnitudes change monotonically in response to the controlled property. The main components are:

1. Resampling and aggregation: Generates thousands of simulated test environments (replicates) by randomly sampling the original activation tensor A with replacement, then calculates the mean across these bootstrapped samples to establish a highly robust, noise-reduced baseline tensor.

2. Trend validation: Evaluates whether a feature’s activation changes consistently across the property of interest dimension using Spearman rank and Mann-Kendall tests. A feature is kept if both independent tests verify that its trajectory is significant below the αsig threshold.

3. Direction sorting and selection: Calculates the overall magnitude of change for the features identified in the previous step, classifying them into sets of increasing or decreasing features, then extracts the top-k most altered features for each individual bootstrap replicate.

4. Final ranking: Defines the final selection metric as the bootstrap frequency C/R, providing the probability that a feature ranks among the top-k most directionally altered latents across all replicates.

To identify impactful latents, the paper performs activation steering with respect to a downstream task. A uniform steering intervention using a fixed initialization magnitude (αinit = 0.4) is applied to each latent, evaluating the model’s modified forward pass on the reference dataset, isolating only latents that exceed the unsteered baseline accuracy, and retaining the top 40 for more granular tuning. Grid search on these top-40 performing latents finds the optimal strength α* for each, and the top-10 latents are returned.

Experimental Setup

The paper applies MoRFI to three major open-weight models: Llama 3.1 8B, Gemma 2 9B, and Mistral v03 7B. Pre-trained SAEs from Llama Scope, Gemma Scope, and Engels et al. are used to extract sparse representations of the residual stream from the fine-tuned models. Activations are collected only from the middle layer of each model, with decoder vectors normalized to unit length and their original magnitudes pushed to the encoder.

The dataset is based on EntityQuestions, where triplets from a diverse set of relations from Wikidata are converted to QA pairs. Train samples are annotated as Known or Unknown based on PCorrect and the SliCK categorization of knowledge. The target property is instantiated as the proportion of Unknown samples contained in the fine-tuning dataset, with seven distinct fine-tuning datasets Dp for p ∈ P = 0, 10, 25, 50, 75, 90, 100. Each Dp is equal in size, comprising p% Unknown instances.

SFT is used to train each model on their corresponding set of Dp s for a fixed number of epochs across 5 different runs for 10, 20, 30, 40 and 50 epochs, totalling 35 distinct fine-tuning runs for each model, evaluated on the test set. An additional Early Stop run uses dev set accuracy to determine the optimal stopping point. Figure 2 plots each model’s accuracy on the test set across the datasets while varying training time, demonstrating that increasing % of Unknown facts in the fine-tuning data mixture consistently leads to performance degradation, intensified with longer training.

Experiments and Results

The paper aims to identify latent directions of the residual stream that capture representational changes induced by fine-tuning on unknown facts, adversely affecting the model’s ability to access pre-trained knowledge. The bootstrap analysis identifies SAE latents whose activations exhibit robust monotonic trends along two axes: the proportion of unknown facts in the fine-tuning mixture and the number of training epochs. Activation steering then tests whether steering these candidates can recover the fine-tuned model’s accuracy on previously known facts.

Ruling out artifacts with control latents: Figure 3 shows that the top candidates discovered by the algorithm have a relatively significant effect on performance when used for steering. A control group of latents that appear unaffected by fine-tuning behaves very differently: steering with them mostly has a highly negative effect, whereas positive effects are significantly less pronounced than the top-ranked latents, indicating the effectiveness of MoRFI.

Asymmetries in latent steering dynamics: Negatively steered latents consistently yield larger gains than positively steered ones. For the ↑Unknown direction, the best positively steered latent gives +21.7 while the best negatively steered one gives +33.4 for Llama. For Mistral 7B the difference is more pronounced: +24.91 vs +39.77, and for Gemma 2 9B: +9.03 vs +18.46. This suggests that features that grow with % Unknown may become over-expressed and suppressing them recovers performance more effectively than promoting other features.

Decreasing latents are significantly more impactful relative to increasing ones: Features that decrease for models fine-tuned with more Unknown samples, when steered in either direction, tend to produce accuracy gains that match or exceed those from increasing latents. In Llama, the top decreasing latent (+21.7) when positively steered, triples the best increasing (+7.12). In Gemma, the top decreasing, positively steered latent (+19.02) outperforms every increasing positively steered one. In Mistral, the top decreasing latent reaches +39.77 relative gains.

Latent structure of Known vs Unknown knowledge: Table 2 decomposes the performance gains achieved through steering by attributing them to specific Wikidata relations. The Knowledge Recovery rate (RK) measures what fraction of the steering-induced gains correspond to MD0’s knowledge. The paper finds that 69-85% of the facts recovered by steering top-latents on MD100s are Known to MD0. The top benefitting relations: P17 (country), P36 (capital) and P495 (country of origin) are semantically related geographic predicates, and their consistent co-recovery under single-latent steering across all three models is consistent with the expectation that semantically relevant relations share representational structure in the residual stream.

Composite direction validation: To validate the direction identified by Algorithm 1, the paper steers directly with δu, the composite vector capturing the aggregate activation shift from MD0 to MD100 in SAE space. Across all three models, subtracting this direction produces accuracy gains in hallucinating models, while adding it deteriorates performance (Figure 4). However, single-latent interventions sourced with MoRFI mostly outperform δu−, indicating that the knowledge-relevant signal within the composite direction is concentrated in a sparse subset of its components, with the remainder diluting the effect or reversing it. Within-group similarity of top-performing latents is generally low (≈0.3), indicating that the top-performing latents span a distributed subspace at the middle layer rather than clustering along a shared low-rank structure.

Conclusion

The paper introduces a two-stage pipeline that identifies residual-stream directions causally linked to model behavior. Using activation snapshots, it isolates feature-level changes and their relevance. Applied to fine-tuning-induced knowledge degradation across three architectures, MoRFI shows a large gap between composite and single-latent interventions, indicating sparse knowledge signals. High alignment (69–85%) with MD0 suggests these latents control access to stored knowledge rather than performance itself. This implies forgetting reflects disrupted access, not erasure, with implications for reducing hallucinations via targeted inference-time interventions.

Based on the paper, here are the specific improvements I can implement in an AI system, and what the improved system can do:

Improvements to Implement:

1. Add a Knowledge Provenance Layer to the Model Architecture:

• Implementation: After the middle layer (e.g., layer 16 for Llama 3.1 8B), insert a lightweight, trainable knowledge gate module. This module takes the residual stream and, using a pre-trained SAE, computes activations for the specific latents identified by the paper (e.g., latents 21682, 10007, 7859 for Mistral). The gate learns a scalar weight for each of these latents, effectively controlling how much that latent's direction contributes to the final output.

• Resulting Capability: The system can now explicitly control access to its parametric knowledge. During inference, it can dynamically suppress or amplify these specific directions based on the context, directly mitigating hallucinations caused by fine-tuning on unknown facts.

2. Implement a Monotonic Drift Monitor for Post-Training:

• Implementation: Integrate the MoRFI algorithm (Algorithm 1) as a standard evaluation step during any fine-tuning or RLHF process. After each epoch, compute the activation tensor for a held-out set of known facts. Run MoRFI to identify latents that are monotonically increasing or decreasing in response to the training data. If the drift of these latents exceeds a predefined threshold, pause training and issue a warning.

• Resulting Capability: The system can now proactively detect when it is starting to forget or lose access to its pre-trained knowledge, even if overall task accuracy is still improving. This prevents the silent degradation of general knowledge during specialized fine-tuning.

3. Create a Knowledge Recovery Inference-Time Intervention:

• Implementation: Build a post-processing module that can be applied to any fine-tuned model. This module uses the paper's findings to identify the optimal steering direction (positive/negative) and strength (α) for the top-10 most impactful latents (e.g., from Table 10). When a user query is identified as a closed-book QA task, the module applies the optimal steering intervention to the residual stream at the middle layer before the final forward pass.

• Resulting Capability: The system can recover up to 85% of the knowledge that was lost during fine-tuning, without any retraining. This provides a powerful, inference-time tool to fix hallucination issues in already-deployed models.

4. Develop a Knowledge Access Audit Tool:

• Implementation: Create a diagnostic tool that, for a given model and a set of questions, uses the SAE to decompose the residual stream at the middle layer. The tool then compares the activation pattern against the known increasing and decreasing latent sets identified by MoRFI. A high activation of increasing latents (associated with unknown facts) and low activation of decreasing latents (associated with known facts) indicates a high risk of hallucination.

• Resulting Capability: The system can now provide a confidence score for its own answers based on internal mechanistic states, not just token probabilities. This allows for more reliable abstention or I don't know responses when the model's knowledge access pathways are disrupted.

5. Add a Fine-Tuning Data Purity Checker:

• Implementation: Before training on a new dataset, use the base model to compute PCorrect for all samples, classifying them as Known or Unknown. Then, use the paper's findings to predict the impact of the data mixture. If the percentage of Unknown facts exceeds a threshold (e.g., 25%), the system will flag the dataset as high-risk for inducing hallucinations and suggest either filtering out the Unknown samples or using a different training strategy.

• Resulting Capability: The system can now prevent hallucinations at the source by ensuring that fine-tuning data is aligned with the model's existing parametric knowledge, or by warning the user of the potential consequences before training begins.

What the Improved AI System Can Do:

• Self-Diagnose Hallucinations: It can identify when its internal knowledge access pathways are compromised, rather than just when its output is likely wrong.

• Recover Lost Knowledge: It can use targeted, single-latent interventions to restore access to pre-trained knowledge that was suppressed during fine-tuning, achieving significant accuracy gains.

• Provide Mechanistic Confidence Scores: It can generate confidence scores based on the state of its internal knowledge representations, offering a more robust alternative to softmax probabilities.

• Prevent Forgetting: It can monitor its own internal state during training and alert developers to the onset of knowledge degradation, allowing for early intervention.

• Audit Its Own Training Data: It can assess the risk of a new dataset before training, preventing the introduction of hallucination-inducing knowledge conflicts.

Large language models (LLMs) acquire most of their factual knowledge during the pre-training stage, through next token prediction. Subsequent stages of post-training often introduce new facts outwith the parametric knowledge, giving rise to hallucinations. While it has been demonstrated that supervised fine-tuning (SFT) on new knowledge may exacerbate the problem, the underlying mechanisms are still poorly understood. We conduct a controlled fine-tuning experiment, focusing on closed-book QA, and find latent directions that causally contribute to hallucinations. Specifically, we fine-tune Llama 3.1 8B, Gemma 2 9B and Mistral 7B v03 on seven distinct single QA datasets, controlling for the percentage of new knowledge and number of training epochs. By measuring performance on the test set, we validate that incrementally introducing new knowledge increases hallucinations, with the effect being more pronounced with prolonged training. We leverage pre-trained sparse autoencoders (SAEs) to analyze residual stream activations across various checkpoints for each model and propose Monotonic Relationship Feature Identification (MoRFI) for capturing causally relevant latents. MoRFI filters SAE features that respond monotonically to controlled fine-tuning data mixtures of a target property. Our findings show that exposure to unknown facts disrupts the model's ability to retrieve stored knowledge along a set of directions in the residual stream. Our pipeline reliably discovers them across distinct models, recovering knowledge through single-latent interventions.