| Title: | Probabilistic Attribution Flow for Large Language Model Debugging |
| Subject: | Computer science, Software engineering, Robotics |
| Level: | Advanced |
| Description: |
BackgroundLarge language models (LLMs) are increasingly deployed in high-stakes applications, yet their internal decision processes remain largely opaque. When an LLM produces an incorrect, biased, or unexpected output, practitioners have limited tools for diagnosing the failure. Existing attribution methods such as attention weights, gradient-based saliency, and perturbation-based approaches either lack theoretical grounding, require multiple forward passes, or provide only input-level attribution without insight into which layers or components of the model contributed to the failure. Probabilistic Activation Flow (PAF)n is an attribution framework that propagates a probability distribution backwards through a neural network’s computation graph, satisfying exact mass conservation at every layer. PAF produces simultaneous layer-wise attribution distributions that enable calibrated probabilistic statements about which parts of the network drove a specific prediction. PAF has been evaluated on image classifiers and vision transformers, but has not been extended to autoregressive language models. Research ObjectiveThis thesis extends PAF to autoregressive language models and evaluates its utility as a debugging tool for understanding and diagnosing LLM failures. The primary research question is: Can PAF’s layer-wise probabilistic attribution identify the internal components of a language model responsible for incorrect or unexpected outputs, and does this information support more targeted debugging than existing attribution methods? Proposed ApproachPAF is already implemented for vision transformers. This thesis will require extending PAF to GPT-2, producing token-level attribution distributions per layer. Three debugging scenarios will be investigated: hallucination localisation (identifying which input tokens drove a factually incorrect answer), layer-wise failure localisation (identifying which layers show the largest attribution redistribution on incorrect vs correct outputs), and attention head contribution analysis (identifying which heads carry the most attribution mass for specific reasoning tasks). The implementation will be evaluated based on a principled evaluation criterion grounded in PAF’s probabilistic interpretation, using a metric measuring the fraction of attribution mass falling on human annotated evidence tokens. EvaluationThe thesis may evaluate the PAF-extension to the ERASER benchmark (human rationale annotations), TriviaQA (factual question answering with correct/incorrect answer analysis), and SST-2 (sentiment classification).
Expected Contributions1. A hook-based extension of PAF for autoregressive language models producing token-level and layer-wise attribution distributions. 2. Proposing a Mass-in-Rationale (MiR) metric for evaluating probabilistic token attribution against human rationale annotations. 3. Empirical analysis of PAF’s utility for three LLM debugging scenarios on standard benchmarks. 4. Comparison against existing attribution methods demonstrating the diagnostic value of layer-wise mass-conserving attribution. Rough TimelinePAF extension to GPT-2, hook -based implementation (Month 1), Debugging scenarios and metric design (Month 2), Experiments on ERASER, TriviaQA, SST-2 (Month 3), Thesis writing and finalisation (Month 4) ReferenceFerrando, Javier and Elena Voita. “Information Flow Routes: Automatically Interpreting Language Models at Scale.” ArXiv abs/2403.00824 (2024) |
| Start date: | |
| End date: | |
| Prerequisites: |
Knowledge on internal working mechanism on neural networks/LLMs and python programming language to work with the computation graph of neural networks/LLMs. |
| IDT supervisors: | Abu Naser Masud |
| Examiner: | |
| Comments: | |
| Company contact: |