Drift: Enhancing LLM Faithfulness in Rationale Generation via Dual-Reward Probabilistic Inference

A unified framework for evaluating and improving the faithfulness of large-language-model (LLM) reasoning across Natural Language Inference, Question Answering and Student Assessment.
At its core is Drift – a dual-reward probabilistic inference algorithm that couples a task-specific classifier reward with a domain-aware rationale reward to guide decoding, yielding explanations that better reflect model decisions and input evidence.

---

Why Drift?

1. Domain-sensitive rationales Drift explicitly encourages generation of tokens that align with domain-specific expert distributions, overcoming the tendency of base LLMs to produce generic explanations.
2. Dual rewards, single pass A task reward locks the first token(s) to the expert classifier’s prediction, while a look-ahead rationale reward re-weights each partial hypothesis during beam search using an expert generative model.
3. Light-weight & plug-and-play Implemented entirely at inference time; no gradient updates or model retraining required. Works with any open-source LLM that exposes logits.

Key Contributions (paper highlights)

• Introduce a probabilistic inference framework grounded in Feynman-Kac Sequential Monte-Carlo that naturally balances exploration and exploitation without costly roll-outs.
• Demonstrate significant gains on seven datasets spanning three tasks – up to +31 pp accuracy and >2 × faithfulness over backbone models.
• Show generalisability across backbones (Llama-3-8B, Mistral-7B) and with “weak” out-of-task experts, retaining benefits even when expert knowledge is only loosely related.

Supported Tasks & Datasets

Method in a Nutshell

(initial state) → TaskReward → beam expand → RationaleReward → resample … repeat

• TaskReward: mask the vocabulary so that the first token must equal the classifier’s predicted label; renormalise logits.
• RationaleReward: after each expansion, score the partial sequence with a domain expert and weight hypotheses accordingly; select the path with highest average weight at the end.
• See Algorithm 1 in the paper for full pseudocode.

Installation

# Clone repository
$ git clone https://github.com/lijiazheng99/drift.git
$ cd drift

# Install core dependencies
$ pip install -r experiments/requirements.txt

We recommend Python ≥ 3.10 and CUDA-enabled GPUs (tested on A100-40 GB and L40s with 8‑bit quantisation). There are known issues with H100 GPUs.

Quick Start

Run experiments via the helper script:

# NLI with default settings (Llama-3 backbone)
$ ./exps.sh --task nli --exp_num exp_0 --device 0

# TruthfulQA with classifier guidance (exp_3)
$ ./exps.sh --task qa --exp_num exp_3 --device 0

# ASAP prompt 2 with full Drift (exp_5)
$ ./exps.sh --task sa --dataset asap_2 --exp_num exp_5 --device 0

See experiments/faithful_evaluate_<task>.py for all command‑line flags.

Reproducing Paper Results

1. Download or prepare datasets under data/ (formats in data/README.md).
2. Obtain backbone checkpoints (e.g. meta-llama/Llama-3-8B-Instruct) and place in models/backbone/.
3. Execute the supplied scripts/* to replicate Table 1 & 2 metrics.

Tip : Faithfulness scores are normalised per dataset using the min–max ranges in Appendix A.1 of the paper.

Interpreting the Metrics

• Acc Prediction accuracy.
• Faith Semantic change (student assessment) or word‑inclusion rate (NLI/QA) under counterfactual perturbations – higher = more faithful.
• Overall Simple sum of Acc and normalised Faith.

Detailed ablations (task‑only vs rationale‑only) are illustrated in Fig. 3 of the paper.

Limitations & Future Work

• Currently evaluated on classification‑style tasks; extension to free‑form generative reasoning (e.g. maths proofs) is left to future work.
• Strong task‑reward constraints may degrade performance when the classifier is unreliable; adaptive weighting schemes are an open question.
• Exploration of a broader zoo of expert models is ongoing.

Citing Drift

If you use this framework, please cite:

@inproceedings{li2025drift,
  title     = {Drift: Enhancing LLM Faithfulness in Rationale Generation via Dual-Reward Probabilistic Inference},
  author    = {Li, Jiazheng and Yan, Hanqi and He, Yulan},
  booktitle = {Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics},
  year      = {2025}
}