💻 computer science

Latent Attribution Regularization: Cross-Domain Attribution-Consistency Training for Stable NLP Explanations

This paper introduces Latent Attribution Regularization (LAR), a training-time method that enforces attribution consistency under semantically preserving synonym substitutions, demonstrating that it significantly improves the stability of gradient-based explanations in fine-tuned Transformers without compromising model accuracy, while also establishing that proper token alignment is essential for accurately measuring such stability.

Original authors: Junaid Hassan

Published 2026-09-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Junaid Hassan

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, a common way to understand how a computer makes a decision is to look at which parts of the input it paid the most attention to. Imagine a machine reading a sentence and deciding if it is happy or sad; researchers can generate a map showing which words the model considered most important for that judgment. These maps, known as feature attributions, are increasingly used to explain AI decisions to people. However, a troubling pattern has emerged: if you slightly change the sentence by swapping a word for a synonym that means the same thing, the map often changes completely. The computer might suddenly point to a different word as the reason for its decision, even though the meaning of the sentence and the final answer have not changed. This makes the explanation feel untrustworthy, as if the AI is unreliable or confused, when in reality, the change might be an illusion created by how the computer reads the text.

A new study by independent researcher Junaid Hassan investigates this problem, revealing that much of the apparent instability is actually a measurement error caused by how computers break down words. Modern language models do not read whole words like humans do; instead, they chop text into smaller pieces called subword tokens. When a single word is replaced with a synonym, the number of these small pieces can change, shifting the position of every single token that follows. If researchers compare the importance scores of words based on their position in the list, they are often comparing the wrong things entirely, like matching the first page of one book to the second page of another. The researcher developed a method to fix this by matching words based on their actual identity rather than their position, and found that the stability of these explanations was far better than previously thought.

Building on this corrected way of measuring, the study introduces a new training technique called Latent Attribution Regularization. This method teaches the AI model to keep its explanation maps consistent, even when the input text is slightly altered with synonyms. During the training process, the model is shown a sentence and a slightly changed version of that sentence, and it is rewarded if the importance scores for the unchanged words remain similar in both versions. This is done alongside the standard task of learning to classify the sentiment of the text correctly. The goal is not just to get the right answer, but to ensure the reasoning behind that answer stays steady when the wording shifts.

The researcher tested this approach using a model called DistilBERT on a dataset of movie reviews. The experiments were run at two different scales: a smaller test with 1,500 examples and a larger, more realistic test with 20,000 examples. In both cases, the model was trained with and without the new consistency technique. The results showed that the models trained with the consistency technique produced significantly more stable explanations. When the text was altered with light, medium, or heavy amounts of synonym substitutions, the explanations remained consistent far more often than those from models trained without the technique. For instance, in the larger test, the stability improvement grew as the text changes became more difficult, showing a clear benefit under harder conditions.

Crucially, this improvement in stability did not come at the cost of the model's ability to understand the text. The models trained with the new technique achieved the same high accuracy in classifying movie reviews as the models trained without it. In the larger experiment, both groups reached an accuracy of around 82 percent, proving that the model could learn to be more consistent without forgetting how to do its main job. The study also confirmed that the initial finding of extreme instability was indeed an artifact of the measurement method; once the alignment was corrected, the baseline stability was already high, but the new training method pushed it even higher.

The findings suggest that for developers building AI systems that explain their decisions to users, there is a practical way to make those explanations more reliable. By adding this consistency check during the training phase, the AI can learn to ignore the noise caused by minor wording changes and focus on the true meaning. The researcher notes that while this work was done on a specific type of model and a single task, the method offers a low-risk addition to standard training that could make AI explanations more trustworthy for real-world applications. The study concludes that the instability often blamed on the AI models themselves may have been a trick of the measurement, and that with the right tools and training, we can build systems that explain themselves with greater steadiness.

Technical Summary: Latent Attribution Regularization (LAR)

Problem Statement

Gradient-based feature attributions in fine-tuned Transformer models often exhibit high instability when inputs undergo small, semantically preserving perturbations (e.g., synonym substitutions). If an explanation changes drastically while the model's prediction and the input's meaning remain constant, the explanation is untrustworthy.

The paper identifies a critical, previously overlooked measurement artifact in this domain: tokenization drift. Standard subword tokenization (WordPiece) means that substituting a single word can alter the number of subword tokens, shifting the position of all subsequent tokens. Evaluating attribution stability by comparing raw positional indices (as done in prior work like Yin et al., 2022) causes a severe misalignment. The author demonstrates that this straightforward approach collapses stability correlations from a plausible $\approx 0.90$ to as low as $\approx 0.03$ , creating a false impression of total instability independent of actual model behavior.

Methodology

1. Corrected Evaluation Protocol

To address the measurement artifact, the author introduces an exact token-alignment procedure:

Alignment: Instead of comparing tokens by position, the algorithm matches tokens in the original and perturbed sequences by their exact string identity, preserving order.
Scope: The metric computes Spearman rank correlation only over "background" tokens (those unchanged by the substitution). Substituted words are excluded from the correlation calculation by design.
Requirement: A sample contributes to the statistic only if at least three tokens are successfully matched.

2. Latent Attribution Regularization (LAR)

Building on the corrected evaluation, the author adapts a training-time regularization strategy for NLP:

Objective: The training loss combines standard cross-entropy ( $L_{CE}$ ) with a cosine-similarity penalty between the feature attributions of the original input ( $x$ ) and a perturbed input ( $\tilde{x}$ ).
$L_{LAR} = L_{CE}(f(x), y) + \lambda \left(1 - \cos(a(x), a(\tilde{x}))\right)$
Attribution Method: The study uses the $\ell_2$ norm of the input gradient ( $\nabla_{e_i} \log_y(f(e))$ ) rather than Integrated Gradients (IG), chosen for computational simplicity within the double-backward training loop.
Perturbation: Inputs are perturbed by substituting adjectives and adverbs with WordNet synonyms. Unlike some prior contrastive generation methods, this process does not filter for model confidence preservation.
Training vs. Evaluation Asymmetry: While evaluation uses exact token identity matching, the training loss compares attributions positionally (truncating to the shorter sequence). This is necessary because identity-based matching is a discrete operation incompatible with the differentiable double-backward computation required for gradient-of-gradient calculations.

3. Experimental Setup

Model: DistilBERT (66M parameters), initialized from a non-task-tuned checkpoint (randomly initialized classification head) to ensure genuine learning rather than re-memorization.
Task: Binary sentiment classification on SST-2.
Scales: Experiments were run at two scales:
- Small: 1,500 examples, 1 epoch.
- Large: 20,000 examples, 2 epochs (reaching $\approx 82\%$ validation accuracy).
Control: A matched cross-entropy-only control ( $\lambda=0$ ) using the same seeds, data, and optimizer settings.
Perturbation Strengths: Evaluated at Light (15%), Medium (35%), and Heavy (60%) substitution rates.

Key Results

1. The Token-Alignment Artifact

In preliminary experiments without regularization, straightforward positional alignment yielded a stability correlation of $\bar{\rho} \approx 0.035$ . Applying the exact token-alignment correction to the same data and model raised the baseline stability to $\bar{\rho} \approx 0.90$ . This confirms that the apparent instability in prior literature was largely a measurement artifact.

2. Stability Improvements with LAR

Across both training scales and 10 independent seeds, LAR ( $\lambda=0.5$ ) produced statistically significant improvements in attribution stability compared to the control:

Small Scale: Significant improvements at all perturbation strengths (paired t-test $p \le 0.012$ ).
Large Scale: Significant improvements at all perturbation strengths (paired t-test $p = 0.002\text{--}0.004$ ; Wilcoxon $p = 0.004\text{--}0.006$ ).
Magnitude: The improvement in stability grew with perturbation severity (e.g., $+0.026$ for light, $+0.036$ for heavy perturbations in the large-scale setting).
Consistency: 9 out of 10 seeds agreed on the direction of improvement at every perturbation strength.

3. Accuracy Trade-offs

Crucially, the stability gains came without measurable accuracy degradation:

Small Scale: Control (73.2%) vs. LAR (73.3%).
Large Scale: Control (81.6%) vs. LAR (82.3%, $p=0.42$ ).

Significance and Claims

The paper positions its contributions at the intersection of attribution robustness (previously studied in vision under adversarial attacks) and NLP stability evaluation.

Correction of Measurement Artifacts: The primary claim is that previous evaluations of attribution stability in NLP were fundamentally flawed due to tokenization drift. The author provides a necessary correction (exact token alignment) to enable trustworthy comparisons.
First Application of Consistency Regularization to NLP: The author bridges the gap between vision-based attribution regularization (e.g., Wang and Kong, 2022) and NLP. They adapt the cosine-similarity objective to Transformer fine-tuning under semantically preserving perturbations, rather than adversarial noise.
Direct Training Objective: Unlike prior work that used stability metrics solely for post-hoc evaluation, this paper uses the stability metric (specifically the synonym-substitution stability introduced by Yin et al., 2022) as a direct training objective.
Practical Implication: For practitioners deploying fine-tuned Transformers where explanations are visible to end-users, LAR offers a low-risk addition to standard fine-tuning. It provides statistically robust stability improvements, particularly under significant semantic perturbations, without sacrificing classification performance.

Limitations and Scope

The author explicitly notes the following constraints:

Single Task/Architecture: Results are specific to SST-2 and DistilBERT.
Attribution Method: The study uses input-gradient norms, not Integrated Gradients (IG).
Training/Evaluation Asymmetry: The training loss uses positional alignment while evaluation uses identity alignment; closing this gap is a direction for future work.
Perturbation Generation: The synonym substitution does not filter for model confidence preservation, differing from the original protocol of the stability metric.
Convergence: The large-scale experiments reached $\approx 82\%$ accuracy, which is below a fully converged SST-2 fine-tune ( $\approx 90\%+$ ), leaving open the possibility of different behavior at full convergence.