Mechanistic interpretability for enhancing RAG models

Marine Delvallez

Retrieval Augmented Generation - RAG

Retrieval Augmented Generation - Definition

Simple RAG Architecture (Lewis et al. 2020)

Advanced RAG Achitecture (Fan et al. 2024)

Example

RAG Architecture used here (inspired from (Tran et al. 2024))

Explainability in Artificial Intelligence

Explainability in Artificial Intelligence

Open black box models

Aims

  • Trustability
  • Understandability
  • Model Rectification

Definitions

XAI : Make model’s behavior understandable for human (Bell et al. 2022)
Understand : Predict model’s behavior (Bell et al. 2022)
Explanation : Any way to make the decision process of the model understandable for human

Interpretability
How ?

Explanability
Why ?

Explanation through creation of a model

Explanation through creation of a model

Presentation and Demonstration of MechIR

MechIR

Mechanistic interpretability

Understand the internal mechanisms of neural networks by performing causal interventions on specific model components

MechIR (Parry et al. 2025)

  • Encoder-only models

  • For Information Retrieval models

  • Identify components responsible for some behavior

  • Activation Patching Technique

Activation Patching (Chen et al. 2024)

Let \(Q \times D \subset \mathcal{Q}\times\mathcal{D}\) be a set of pairs of questions and documents
Let \(Q \times \tilde{D}\) the same set of pairs but with perturbed documents

  1. Forward pass all \(Q\times D\)
    • record \(o_{i,j}^e\) the output of each component \(n_{i,j}, \forall e \in Q\times D\)
    • record \(p_D\) the performance of the model
  2. Forward pass all \(Q\times \tilde{D}\)
    • record \(o_{i,j}^\tilde{e}\) the output of each component \(n_{i,j}, \forall \tilde{e} \in Q\times \tilde{D}\)
    • record \(p_\tilde{D}\) the performance of the model
  3. Rewrite \(D, e, \tilde{D} \text{ and } \tilde{e}\) as
    • \(\hat{D}, \hat{e}, \check{D} \text{ and } \check{e}\) if \(p_D > p_\tilde{D}\)
    • \(\check{D}, \check{e}, \hat{D} \text{ and } \hat{e}\) otherwise

Activation Patching (Chen et al. 2024)

  1. Rewrite \(D, e, \tilde{D} \text{ and } \tilde{e}\) as
    • \(\hat{D}, \hat{e}, \check{D} \text{ and } \check{e}\) if \(p_D > p_\tilde{D}\)
    • \(\check{D}, \check{e}, \hat{D} \text{ and } \hat{e}\) otherwise
  1. For each component \(n_{i,j}\) forward pass \(Q\times\check{D}\) but replace \(o_{i,j}^{\check{e}}\) by \(o_{i,j}^{\hat{e}}\) for each \(\check{e}\). Record the performance \(\bar{p}\)
  2. \(P = \frac{\bar{p} - p_\hat{D} }{p_\check{D} - p_\hat{D}}\) gives the impact of the perturbation on the model performance

Animation de l’execution de Activation patching

Step 1: Choose a perturbation

Function that applies the same modification on each document.
Example :

@perturbation
def pert1(doc:str) -> str :
  return doc.replace("solution", "answer")

@perturbation
def pert2(doc:str) -> str:
  return doc.replace("microwave", "toaster")

Perturbation creation technique

  • Identify vocabulary specific to the dataset
  • Find in the vocabulary words with several meaning \(m_D\) and \(m_D\)
  • Replace that word by a synonym of the \(m_D\) meaning

What is a good perturbation

  • Has an impact of the documents representation
  • Be useful for interpretation

Perturbation Score

Step 2 : Instantiate the model and load data

dot_model_name = "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"
dot_model = Dot(dot_model_name)
Moving model to device:  cpu
Loaded pretrained model sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco into HookedEncoder
dataset = MechIRDataset("vaswani", query_id_subset=["1"])
pert1_dot_collator = DotDataCollator(dot_model.tokenizer, pert1, q_max_length=None, d_max_length=None, perturb_type="replace")
pert1_dot_dataloader = DataLoader(dataset, batch_size=16, collate_fn=pert1_dot_collator)

pert2_dot_collator = DotDataCollator(dot_model.tokenizer, pert2, q_max_length=None, d_max_length=None, perturb_type="replace")
pert2_dot_dataloader = DataLoader(dataset, batch_size=16, collate_fn=pert2_dot_collator)

Visualisation of perturbed data

solution -> answer
Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]
Baseline Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]
Perturbed Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]
==================================================
Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]
Baseline Document: [CLS] microwave measurements of dielectric absorption in dilute solutions [SEP]
Perturbed Document: [CLS] microwave measurements of dielectric absorption in dilute answers [SEP]
==================================================
microwave -> toaster
Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]
Baseline Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]
Perturbed Document: [CLS] broadband millimetre wave paramagnetic resonance spectrometer the specimen and waveguide which can be cooled by means of a cryostat are placed between close pole pieces giving high uniform magnetic fields design details and some measurements on zero field splittings are given [SEP]
==================================================
Query: [CLS] measurement of dielectric constant of liquids by the use of microwave techniques [SEP]
Baseline Document: [CLS] microwave a measurements of dielectric absorption in dilute solutions [SEP]
Perturbed Document: [CLS] toaster measurements of dielectric absorption in dilute solutions [SEP]
==================================================

Step 4 : Measure the impact of the perturbation on the model

# Initialize lists to store baseline and perturbed performances for each dataloader
all_baseline_performance = {"pert1": [], "pert2": []}
all_perturbed_performance = {"pert1": [], "pert2": []}

# Calculate performances for each perturbation_type
calculate_performance(dot_model, pert1_dot_dataloader, all_baseline_performance["pert1"], all_perturbed_performance["pert1"])
calculate_performance(dot_model, pert2_dot_dataloader, all_baseline_performance["pert2"], all_perturbed_performance["pert2"])

plot_score_dists_mult(all_baseline_performance, all_perturbed_performance, plot_type="kde")

Step 5 : Chart the sensitivity of the model to the perturbation

Perspective : Enhance a model with MechIR

TODO

References

Bell, Andrew, Ian Solano-Kamaiko, Oded Nov, and Julia Stoyanovich. 2022. “It’s Just Not That Simple: An Empirical Study of the Accuracy-Explainability Trade-Off in Machine Learning for Public Policy.” 2022 ACM Conference on Fairness Accountability and Transparency (Seoul Republic of Korea), June, 248–66. https://doi.org/10.1145/3531146.3533090.
Chen, Catherine, Jack Merullo, and Carsten Eickhoff. 2024. “Axiomatic Causal Interventions for Reverse Engineering Relevance Computation in Neural Retrieval Models.” Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval (Washington DC USA), July, 1401–10. https://doi.org/10.1145/3626772.3657841.
Fan, Wenqi, Yujuan Ding, Liangbo Ning, et al. 2024. A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models. arXiv. https://doi.org/10.48550/arXiv.2405.06211.
Lewis, Patrick, Ethan Perez, Aleksandra Piktus, et al. 2020. “Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.” arXiv: Computation and Language.
Parry, Andrew, Catherine Chen, Carsten Eickhoff, and Sean MacAvaney. 2025. “MechIR: A Mechanistic Interpretability Framework for Information Retrieval.” Advances in Information Retrieval - 47th European Conference on Information Retrieval, ECIR 2025, Lucca, Italy, April 6-10, 2025, Proceedings, Part V, Lecture Notes in Computer Science, vol. 15576: 89–95. https://doi.org/10.1007/978-3-031-88720-8_16.
Tran, The Trung, Carlos-Emiliano González-Gallardo, and Antoine Doucet. 2024. “Retrieval Augmented Generation for Historical Newspapers.” Proceedings of the 24th ACM/IEEE Joint Conference on Digital Libraries (Hong Kong China), December, 1–5. https://doi.org/10.1145/3677389.3702542.

Demonstration (TASB & Vaswani)