Technical Summary: Data Attribution in Large Language Models via Bidirectional Gradient Optimization (DABGO)

Problem Statement

Large Language Models (LLMs) are increasingly deployed in critical domains such as legal drafting, medical diagnosis, and scientific research. While these models derive their capabilities from vast training corpora, a fundamental open problem remains: identifying which specific training data most influenced a model's output. Current research focuses heavily on model performance given a dataset, leaving the inverse problem—tracing model behavior back to specific data sources—underexplored. This lack of traceability hinders governance, accountability, and the ability to debug or "unlearn" specific data points. Existing Training Data Attribution (TDA) methods, such as Influence Functions, face practical limitations including the infeasibility of exact computation (requiring retraining) and fundamental issues with Hessian inversion in large-scale models. Furthermore, prior NLP attribution studies have often been restricted to fact-tracing with single-word predictions or dataset-level granularity, failing to address fully open-ended text generation.

Methodology

The authors propose DABGO (Data Attribution via Bidirectional Gradient Optimization), a framework designed to estimate backward influence functions for autoregressive LLMs. The core hypothesis relies on the "mirrored influence" concept: the influence of a training sample on a model prediction can be estimated by observing how optimizing the model on that prediction affects the model's behavior on the training sample.

Core Algorithm

1. Bidirectional Optimization: Given a pretrained model $\pi_{\theta^*}$ and a generated output sequence $\hat{x}$ (prompted by $\hat{x}_0$), the method constructs two perturbed model variants:

   • Gradient Ascent ($\theta_+$): Optimizes the model to increase the likelihood of the generated text $\hat{x}$ (simulating "unlearning" or maximizing influence).
   • Gradient Descent ($\theta_-$): Optimizes the model to decrease the likelihood of $\hat{x}$ (simulating fine-tuning or minimizing influence).
   • Both optimizations utilize a Fisher-regularized update (Elastic Weight Consolidation) to mitigate catastrophic forgetting, approximated via a diagonal Fisher Information Matrix (FIM) to ensure computational tractability.

2. Bidirectional Influence Score (BIS): For every training sample $x_i$ in the dataset $D$, the method computes the loss under both perturbed models ($\theta_+$ and $\theta_-$). The influence score is defined as the absolute change in loss:
   $$I(x_i; \hat{x}) = |\ell(x_i, \theta_-) - \ell(x_i, \theta_+)|$$
   Samples exhibiting large absolute changes in loss are ranked as the most influential for the generation of $\hat{x}$.

3. Granularity: The framework supports attribution at arbitrary levels, from the full training sample down to the token level, by measuring per-token log-likelihood differences.

Key Contributions

• Novel Framework: Introduction of DABGO, a simple yet effective attribution method for open-ended text generation that avoids expensive per-sample gradient computations or Hessian inversions.
• Bidirectional Approach: The paper demonstrates that combining both gradient ascent and gradient descent yields superior attribution accuracy compared to unidirectional approaches (using only ascent or only descent), capturing complementary signals in generative models.
• Granular Attribution: The method enables attribution at arbitrary data granularity, successfully capturing both factual content and stylistic characteristics (e.g., authorship, syntax) in attributed texts.
• Evaluation Protocol: The authors introduce and utilize the Tail-Patch Absolute (TPA) score as a computationally efficient proxy for influence, validated against the ground-truth definition of influence via full retraining experiments.

Experimental Results

The authors evaluated DABGO on GPT-2 style models trained from scratch on two distinct corpora:

1. Factual Attribution: Using a curated Wikipedia corpus (230M tokens).
2. Stylistic Attribution: Using Project Gutenberg literary texts (2M tokens) with diverse authors.

Quantitative Performance:

• Factual Setting: DABGO consistently outperformed baselines (BM25, TRACKSTAR, GECKO) across most values of $k$ (number of top attributed samples). While BM25 performed well at $k=1$ due to strong lexical overlap, DABGO excelled as $k$ increased, indicating better identification of distributed factual influence.
• Stylistic Setting: DABGO significantly outperformed word-based retrieval (BM25) and embedding-based methods (GECKO) as $k$ increased. This highlights the method's ability to capture distributed stylistic signals across multiple training samples rather than relying on single passages.
• Validation via Retraining: In counterfactual experiments where top-$k$ attributed samples were removed and models were retrained, models retrained without DABGO-attributed samples showed significantly higher loss on the test output compared to those retrained without samples attributed by other methods. This confirms DABGO's alignment with the ground-truth definition of influence.

Ablation Studies:
Experiments isolating gradient ascent and descent showed that while ascent alone provides strong performance, the combination of both directions (DABGO) achieves the highest accuracy, validating the necessity of the bidirectional approach for generative tasks.

Significance and Claims

The paper claims that DABGO enhances model interpretability, a critical requirement for accountable AI systems. By successfully tracing open-ended generations back to their most influential training data, the method:

• Closes the Feedback Loop: It connects training data directly to model behavior, enabling applications such as model debugging and data unlearning.
• Overcomes Prior Limitations: Unlike previous works restricted to curated facts or single-word predictions, DABGO operates effectively in fully open-ended generation settings.
• Provides Transparency: The ability to attribute influence at the token level offers fine-grained interpretability, distinguishing between factual contributions and stylistic nuances.

The authors acknowledge limitations, noting that DABGO remains computationally expensive (requiring two passes over the corpus) and is currently impractical for industrial-scale models with hundreds of billions of tokens. Additionally, while it identifies influential samples, it does not currently account for synergistic or adversarial interactions between subgroups of data, nor does it establish the causal role of specific highlighted tokens beyond observed loss changes. Nevertheless, the work presents a significant step toward traceable and auditable AI systems.