Neuro-RIT: Neuron-Guided Instruction Tuning for Robust Retrieval-Augmented Language Model

Abstract: Retrieval-Augmented LLMs (RALMs) have demonstrated significant potential in knowledge-intensive tasks; however, they remain vulnerable to performance degradation when presented with irrelevant or noisy retrieved contexts. Existing approaches to enhance robustness typically operate via coarse-grained parameter updates at the layer or module level, often overlooking the inherent neuron-level sparsity of LLMs. To address this limitation, we propose Neuro-RIT (Neuron-guided Robust Instruction Tuning), a novel framework that shifts the paradigm from dense adaptation to precision-driven neuron alignment. Our method explicitly disentangles neurons that are responsible for processing relevant versus irrelevant contexts using attribution-based neuron mining. Subsequently, we introduce a two-stage instruction tuning strategy that enforces a dual capability for noise robustness: achieving direct noise suppression by functionally deactivating neurons exclusive to irrelevant contexts, while simultaneously optimizing targeted layers for evidence distillation. Extensive experiments across diverse QA benchmarks demonstrate that Neuro-RIT consistently outperforms strong baselines and robustness-enhancing methods.

• The paper presents a neuron-level tuning approach to robustify retrieval-augmented language models by actively suppressing irrelevant neuron activations.
• It employs integrated gradients to attribute and decouple context-specific neuron responses, enabling precise two-stage instruction tuning.
• Neuro-RIT achieves superior QA benchmarks performance while updating only about 6.6% of model parameters, enhancing efficiency and scalability.

Neuro-RIT: Neuron-Guided Instruction Tuning for Robust Retrieval-Augmented LLMs

Introduction

The vulnerability of retrieval-augmented LLMs (RALMs) to irrelevant or noisy retrieved information critically limits their effectiveness in knowledge-intensive tasks. Standard retrieval-augmented generation (RAG) systems are not explicitly robustified against distraction from non-relevant evidence, resulting in degraded answer quality, hallucination, and reduced downstream accuracy. Prior robustness-enhancement techniques have largely focused on layer- or module-level parameter adaptation or external refinement modules, neglecting the emergent neuron-level activation sparsity and functional specialization found in state-of-the-art LLMs. "Neuro-RIT: Neuron-Guided Instruction Tuning for Robust Retrieval-Augmented LLM" (2604.02194) introduces a neuron-aware paradigm for robustifying RALMs, leveraging fine-grained neuron attribution and two-stage instruction tuning targeting neuron groups most responsive to relevant and irrelevant contexts.

Neuro-RIT Framework

Neuro-RIT departs from coarse-grained adaptation by mining context-aware neurons and performing instruction tuning guided at the neuron level.

Figure 1: Overview of the Neuro-RIT framework. Neurons highly responsive to relevant and irrelevant contexts are identified and decoupled; denoising and robustification is performed via two-stage neuron-guided instruction tuning.

Context-Aware Neuron Mining

The method utilizes attribution-based neuron mining to identify neurons whose changes in activation are most diagnostic of whether a retrieved document is relevant or irrelevant to a given query. Using Integrated Gradients, the neuron contribution is quantified for each prediction under both conditions. By aggregating high-attribution neurons over many samples, three principal sets are extracted:

• $\mathcal{P}_{\text{rel}}$: neurons most specific to relevant contexts.
• $\mathcal{P}_{\text{irrel}}$: neurons specific to irrelevant (noisy/distracting) contexts.
• $\mathcal{P}_{\text{shared}}$: neurons active across both.

Intersectional analysis is used to decouple strictly context-specific versus shared neuron responses, suppressing overlap and thus maximizing the fidelity of targeted intervention.

Two-Stage Neuron-Guided Tuning

1. Denoising Stage: Neurons in $\mathcal{P}_{\text{irrel}}$ are functionally suppressed via instruction-tuning—they are trained to trigger end-of-text (EOT) emission when facing irrelevant contexts, decoupling such neurons from influencing text generation.
2. Noise Filtering Enhancement: A further tuning stage optimizes neuron-level subsets (including $\mathcal{P}_{\text{rel}}$, $\mathcal{P}_{\text{irrel}}$, $\mathcal{P}_{\text{shared}}$) and top dense layers for robust evidence distillation, ensuring relevant-context neurons drive generative behavior.

This pipeline explicitly imposes dual robustness: active suppression of spurious context processing and enhanced utilization of truly relevant evidence.

Experimental Results

Evaluation spans a range of in-domain and out-of-domain QA benchmarks, including both single-hop (NQ, TriviaQA, ASQA, SCIQ, POPQA) and multi-hop (HotpotQA, 2Wiki) datasets. The principle generator is LLaMA-3-8B-Instruct, but generalization to Mistral-7B is also demonstrated.

Figure 2: Comparative performance of four RAG-based methods on representative QA benchmarks, using an LLM-based evaluation metric.

Neuro-RIT achieves the highest or comparable accuracy to the best existing robustness-enhanced baselines (e.g., PA-RAG, InstructRAG), especially when combined with retrieval refinement modules. For example, Neuro-RIT+Reranker attains the strongest average accuracy, exceeding previous bests by up to 1.5 absolute points on certain leaderboards.

Neuronal Mechanistic Analysis

Layer-wise mapping reveals the functional partitioning of neuron groups:

Figure 3: Distribution of context-aware neuron groups in LLaMA-3-8B-Instruct on the HotpotQA benchmark. Mid layers concentrate relevant-context neurons, upper layers aggregate irrelevant and shared responsive neurons.

• $\mathcal{P}_{\text{rel}}$ clusters in mid- to upper-model layers.
• $\mathcal{P}_{\text{irrel}}$ and $\mathcal{P}_{\text{shared}}$ concentrate in the uppermost layers.
• This topological dissociation justifies selective layer-level adaptation in conjunction with neuron-wise masking.

Analogous trends are recovered for Mistral-7B, evidencing the generator-agnostic nature of the approach.

Figure 4: Neuro-RIT accuracy on Mistral-7B-Instruct-v0.2 across all benchmarks—neuron-level instruction tuning generalizes beyond LLaMA models.

Analysis and Implications

Robustness Under Adverse Contexts

Neuro-RIT exhibits the strongest resilience in settings where all retrieved documents are irrelevant, with answer accuracy consistently topping prior SOTA baselines. This demonstrates that direct neuron suppression is more effective at reducing hallucinations than dense or module-level tuning, especially for retrieval noise and context-oblivious distracting passages.

Parameter Efficiency

The approach updates ≈6.6% of model parameters (via selective neuron and layer-level adaptation), reducing the computational and memory costs by approximately 93.4% relative to full fine-tuning, while preserving or improving downstream accuracy.

Ablations

Removing either the denoising or neuron-level noise-filtering component degrades accuracy substantially, revealing both components are indispensable for robust generalization.

Practical and Theoretical Implications

The neuronal disentanglement and targeted instruction tuning manifest several advantages:

• Precision robustness control: Interventions do not propagate unintended side effects to unrelated knowledge, unlike dense or full-model fine-tuning.
• Generalization: Neuron-level robustification is dataset- and architecture-transferable, supporting generalized pipeline deployment across retrieval systems.
• Scalability: Parameter efficiency enables feasible robust adaptation even for large-scale deployment scenarios, and offers a scalable path for future robustification of even larger LLMs.

Theoretically, the results affirm the conjecture that emergent functional specialization at the neuron level in transformers can be leveraged for rapid, minimally invasive adaptation.

Future Directions

Potential avenues include:

• Extending neuron-aware fine-tuning to complex multi-task learning with explicit competence-driven neuron allocation.
• Employing neuron mining for targeted model editing and unlearning, potentially enabling continual update and error correction.
• Integrating attribution-guided neuron masking into prompt construction for real-time robust controllability.

Conclusion

Neuro-RIT establishes that context-aware neuron mining combined with two-stage neuron-level instruction tuning produces robust, efficient, and broadly generalizable retrieval-augmented LLMs. By aligning adaptation granularity to the functional units underpinning contextualization and noise processing, Neuro-RIT achieves superior robustness and accuracy with minimal compute. This provides a new methodological baseline for robust retrieval-augmented NLP and sets the stage for further research on fine-grained controllable model adaptation.