Attribution-Guided Masking for Robust Cross-Domain Sentiment Classification

Technical Summary: Attribution-Guided Masking for Robust Cross-Domain Sentiment Classification

Problem Statement

Pre-trained Transformer models (e.g., BERT, RoBERTa) achieve state-of-the-art performance on in-domain sentiment classification but frequently suffer severe performance degradation when transferred to out-of-domain data. The authors hypothesize that this generalization gap ($\Delta$) is driven by the models' over-reliance on domain-specific spurious tokens (e.g., @mentions, hashtags, slang) rather than domain-invariant sentiment markers.

The paper initially investigates whether post-hoc token-level attribution drift can predict this failure. They introduce the Attribution Drift Score (ADS) to quantify shifts in token-level importance between domains. However, comprehensive testing across multiple formulations revealed no meaningful correlation between ADS and the generalization gap. This negative finding indicates that post-hoc attribution drift is an insufficient signal for predicting transfer difficulty, necessitating a shift from prediction to training-time intervention.

Methodology: Attribution-Guided Masking (AGM)

The authors propose Attribution-Guided Masking (AGM), a framework designed to dynamically detect and penalize reliance on spurious tokens during the fine-tuning phase. The method operates without requiring target-domain labels or human annotation.

Core Components:

1. Attribution Computation: The model utilizes a Gradient$\times$Input approximation to compute token-level attribution scores ($a$) based on the cross-entropy loss ($L_{CE}$).
   $$a = \nabla_h L_{CE} \odot h$$

2. Spurious Token Detection: Tokens with absolute attribution magnitudes exceeding the 75th percentile ($\tau_{high} = 0.75$) within a sequence are flagged as potentially spurious.

3. Counterfactual Generation: Flagged tokens are replaced with [MASK], and a shared-weight Masked Language Model (MLM) predicts contextually appropriate replacements. The counterfactual sequence is retained only if the model's predicted sentiment label remains unchanged, ensuring the core sentiment is preserved.

4. Training Objective: The AGM loss function combines the standard classification loss with two auxiliary terms:

   • Attribution Masking Loss ($L_{mask}$): Directly penalizes high attribution scores on flagged tokens.
     $$L_{mask} = \frac{1}{|M|} \sum_{i \in M} (a_i)^2$$
   • Counterfactual Contrastive Loss ($L_{CCL}$): (Optional) Aligns the pooled representations of the original and counterfactual inputs using Mean Squared Error to enforce domain-invariant representations.

   The total objective is:
   $$L_{AGM} = L_{CE} + \lambda_1 L_{mask} + \lambda_2 L_{CCL}$$

Experimental Setup and Results

The method was evaluated using a strict zero-shot, leave-one-out transfer protocol across four diverse domains: IMDb, Amazon, TripAdvisor, and Sentiment140. The primary metric was the Generalization Gap ($\Delta = |F1_{source} - F1_{target}|$), where a lower value indicates better robustness.

Key Findings:

• Performance on Hard Transfers: On the most difficult transfer task (Sentiment140), AGM achieved a generalization gap of $\Delta = 0.244$. This outperformed Domain-Adversarial Training (DANN, $\Delta=0.264$), Distributionally Robust Optimization (DRO, $\Delta=0.248$), and Fish ($\Delta=0.247$). It performed comparably to Invariant Risk Minimization (IRM, $\Delta=0.238$) but with substantially lower variance ($\pm 0.015$ for AGM vs. $\pm 0.036$ for IRM).
• Stability and Interpretability: Unlike IRM, which required extensive hyperparameter tuning and showed instability, AGM provided consistent results across eight random seeds. Uniquely, AGM offers token-level interpretability, revealing exactly which features drive the generalization gap.
• Qualitative Analysis: Heatmaps confirmed that AGM successfully suppresses attribution on domain-specific artifacts (e.g., @mentions, numbers, slang) and shifts reliance toward domain-invariant sentiment markers.
• Ablation Study: The study confirmed that attribution-guided masking is the critical component. Removing $L_{mask}$ or replacing the attribution-guided selection with random token selection consistently degraded performance on difficult transfers (increasing $\Delta$ to $\approx 0.260$). The counterfactual contrastive loss ($L_{CCL}$) was found to be optional; the "mask-only" variant achieved equivalent performance with less computational overhead.

Significance and Contributions

The paper claims three primary contributions:

1. Diagnostic Insight: It demonstrates that post-hoc attribution drift is an unreliable predictor of generalization failure, explaining that shared vocabulary tokens carry stable attributions while the problematic domain-specific tokens are excluded from cross-domain comparisons.
2. Novel Intervention: It introduces AGM, a training-time intervention that automatically identifies and penalizes spurious token reliance via gradient-based attribution, requiring no external supervision or target labels.
3. Robust Generalization: Through rigorous evaluation, it shows that AGM achieves competitive generalization on the hardest transfer tasks, outperforming global representation alignment methods (like DANN) which fail when stylistic distances are extreme.

Limitations:
The authors note that the study is limited to binary English sentiment classification. The approach introduces a computational cost of approximately 2–3$\times$ slower training due to the double backward pass and counterfactual generation. Additionally, while AGM performs competitively with IRM on Sentiment140, the differences fall within overlapping confidence intervals, suggesting the results should be interpreted with caution regarding statistical significance.

Conclusion:
The work provides a pathway for developing sentiment classifiers that rely on intrinsic semantic meaning rather than exploitable domain-specific artifacts, offering a scalable solution for zero-shot transfer that is both robust and interpretable.