Elucidating the Design Space of LM Data Attribution

Technical Breakdown§

The paper introduces the Controlled Attribution Benchmark (CATT) to isolate the effects of dataset properties on data attribution for language models. By fine-tuning LLMs on entirely fictitious knowledge, CATT provides unambiguous ground truth for fact tracing, enabling controlled experiments by varying one dataset configuration at a time (e.g., document redundancy, lexical distractors).

Core Methodology§

Data attribution methods typically compute per-token gradients and aggregate them into document-level scores. The authors identify two underexplored design choices:

1. Loss function for gradient computation: Standard cross-entropy (CE) loss gradients are compared to negative log-odds (NLO) gradients.
2. Aggregation of per-token scores: Averaging across all tokens vs. selecting top token scores.

Let $\mathbf{g}_t = \nabla_\theta \ell_t$ be the gradient of the loss at token $t$. For a training document $d$ with tokens $t \in d$, the attribution score for a test query $q$ is often computed as:

$$\text{Attr}(d, q) = \sum_{t \in d} \langle \mathbf{g}_t, \mathbf{g}_q \rangle$$

where $\mathbf{g}_q$ is the gradient of the query loss. Standard methods average these inner products: $\frac{1}{|d|} \sum_{t \in d} \langle \mathbf{g}_t, \mathbf{g}_q \rangle$.

The authors find that true causal evidence is highly localized within a document. Averaging dilutes this signal. Instead, they propose using NLO gradients, which are defined as:

$$\ell_{\text{NLO}} = -\log \frac{p(y_t|x_{<t})}{1 - p(y_t|x_{<t})} = -\log \frac{p(y_t|x_{<t})}{\sum_{y' \neq y_t} p(y'|x_{<t})}$$

This loss emphasizes tokens where the model is uncertain between the correct token and alternatives, surfacing critical tokens. Then, they aggregate only the top-$k$ token scores (e.g., maximum) to compute the document score:

$$\text{Attr}_{\text{top-k}}(d, q) = \text{mean of top-}k \; \langle \mathbf{g}_t^{\text{NLO}}, \mathbf{g}_q^{\text{NLO}} \rangle$$

Implementation Details§

The benchmark fine-tunes LLMs (e.g., Llama, GPT) on synthetic documents containing fictitious facts. Attribution methods include gradient-based approaches like TracIn, Influence Functions, and Representer Point Selection. For each test query, ground truth is the document containing the fact. The authors evaluate recall (fraction of queries where the true document is ranked in top-$k$) and counterfactual retraining performance (FTRACE-TREX).

Key findings:

• Redundancy and distractors consistently degrade attribution performance.
• NLO gradients + top-token aggregation improves recall by 35–72 percentage points over baselines.
• On FTRACE-TREX, counterfactual retraining performance improves for every method tested, by up to 4.9x.

PyTorch Illustration§

import torch
from torch import nn

def compute_attribution_scores(model, train_docs, query, top_k=5):
    """Compute attribution scores for each training doc using NLO gradients and top-k token aggregation."""
    model.eval()
    # Compute query gradient with NLO loss
    query_loss = nlo_loss(model, query)
    query_grad = torch.autograd.grad(query_loss, model.parameters(), retain_graph=True)
    query_grad_flat = torch.cat([g.view(-1) for g in query_grad])
    
    scores = []
    for doc in train_docs:
        token_scores = []
        for token_idx in range(len(doc)):
            # Forward pass for each token prediction
            logits = model(doc[:token_idx])
            loss = nlo_loss_from_logits(logits, doc[token_idx])
            grad = torch.autograd.grad(loss, model.parameters(), retain_graph=True)
            grad_flat = torch.cat([g.view(-1) for g in grad])
            token_scores.append(torch.dot(grad_flat, query_grad_flat).item())
        # Aggregate top-k token scores
        top_scores = sorted(token_scores, reverse=True)[:top_k]
        scores.append(sum(top_scores) / len(top_scores))
    return scores

def nlo_loss(model, input_ids):
    logits = model(input_ids).logits
    # Shift for next-token prediction
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()
    log_probs = torch.log_softmax(shift_logits, dim=-1)
    # NLO = -log(p/(1-p)) = -log p + log(1-p)
    log_p = log_probs.gather(-1, shift_labels.unsqueeze(-1)).squeeze(-1)
    log_1mp = torch.log1p(-torch.exp(log_p))
    return (-log_p + log_1mp).mean()

This approach highlights the importance of localized causal evidence and provides a diagnostic framework for future attribution research.