HyperAIHyperAI

Command Palette

Search for a command to run...

Best Paper AwardACL 2024

Causal Estimation of Memorisation Profiles

Pietro Lesci Clara Meister Thomas Hofmann Andreas Vlachos Tiago Pimentel

Abstract

Understanding memorisation in language models has practical and societal implications, e.g., studying models' training dynamics or preventing copyright infringements. Prior work defines memorisation as the causal effect of training with an instance on the model's ability to predict that instance. This definition relies on a counterfactual: the ability to observe what would have happened had the model not seen that instance. Existing methods struggle to provide computationally efficient and accurate estimates of this counterfactual. Further, they often estimate memorisation for a model architecture rather than for a specific model instance. This paper fills an important gap in the literature, proposing a new, principled, and efficient method to estimate memorisation based on the difference-in-differences design from econometrics. Using this method, we characterise a model's memorisation profile—its memorisation trends across training—by only observing its behaviour on a small set of instances throughout training. In experiments with the Pythia model suite, we find that memorisation (i) is stronger and more persistent in larger models, (ii) is determined by data order and learning rate, and (iii) has stable trends across model sizes, thus making memorisation in larger models predictable from smaller ones.

One-sentence Summary

University of Cambridge and ETH Zürich propose a difference-in-differences method from econometrics that efficiently estimates per-instance memorisation profiles in language models by observing a small set of instances throughout training, with experiments on the Pythia suite revealing that memorisation is stronger in larger models, governed by data order and learning rate, and predictable across model sizes.

Key Contributions

  • An unbiased and computationally efficient estimator for instance-level counterfactual memorisation is developed using a difference-in-differences design from econometrics, estimating how including a training instance changes a specific model's ability to predict that instance.

  • A method for characterising a model's memorisation profile across training is proposed by observing a small set of instances throughout training, separating checkpoint timing, data order, and learning rate effects from architecture-level averages.

  • Experiments with the Pythia model suite show that memorisation is stronger and more persistent in larger models, determined by data order and learning rate, and stable across model sizes, making memorisation in larger models predictable from smaller ones.

Introduction

Large language models are trained on web-scale corpora, yet they can reproduce training sequences verbatim, which matters for copyright, privacy, factual understanding, and training dynamics. Prior work defines memorisation causally, but measuring it is difficult because the counterfactual outcome, a model’s performance without training on an instance, is unobservable. Earlier estimates either require multiple training runs and therefore describe model architectures rather than specific models, or assume the counterfactual is negligible and may conflate memorisation with ordinary predictability. The authors leverage potential outcomes and econometrics to formalise counterfactual memorisation and propose an observational estimator based on training checkpoints and loglikelihood measurements. This yields memorisation profiles across training and shows that larger models memorise more persistently, learning rate and data position matter, and profiles are stable enough to predict memorisation in larger models from smaller ones.

Method

The authors leverage a causal framework to characterize when a language model memorizes training examples. Each example x\boldsymbol{x}x is treated at a specific timestep g=G(x)g = G(\boldsymbol{x})g=G(x) when it is first used for training. Let Yc(x;g)Y_c(\boldsymbol{x};g)Yc​(x;g) denote the model performance on x\boldsymbol{x}x at checkpoint ccc under the condition that x\boldsymbol{x}x was included in training at timestep ggg. The counterfactual outcome Yc(x;∞)Y_c(\boldsymbol{x};\infty)Yc​(x;∞) represents the hypothetical performance on the same example if it had never been used for training. Under this framing, counterfactual memorization is defined as the causal effect of having trained on an example at timestep ggg on its later evaluation at timestep ccc:

τx,c=Yc(x;g)−Yc(x;∞).\tau_{\boldsymbol{x},c} = Y_c(\boldsymbol{x};g) - Y_c(\boldsymbol{x};\infty).τx,c​=Yc​(x;g)−Yc​(x;∞).

Because individual counterfactual outcomes cannot be observed, the method focuses on an average treatment effect over all examples treated at timestep ggg:

τg,c=Ex[Yc(x;g)−Yc(x;∞)∣G(x)=g].\tau_{g,c} = \underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};g) - Y_c(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = g \right].τg,c​=xE​[Yc​(x;g)−Yc​(x;∞)∣G(x)=g].

The collection of values τg,c\tau_{g,c}τg,c​ across treatment timesteps ggg and checkpoint timesteps ccc forms a memorization profile. Each row of this profile represents the memorization path for a batch of examples introduced at timestep ggg. There is no memorization for checkpoints before exposure, so τg,c=0\tau_{g,c}=0τg,c​=0 when c<gc<gc<g. The effect is called instantaneous memorization when c=gc=gc=g, persistent memorization when c>gc>gc>g, and residual memorization when c=Tc=Tc=T.

For estimation, the authors first observe the factual outcome for examples in the training batch BgB_gBg​. They define the batch-average observed performance at checkpoint ccc as:

Y‾c(g)=1∣Bg∣∑x∈BgYc(x).\overline{Y}_c(g) = \frac{1}{|\mathcal{B}_g|} \sum_{\boldsymbol{x} \in \mathcal{B}_g} Y_c(\boldsymbol{x}).Yc​(g)=∣Bg​∣1​x∈Bg​∑​Yc​(x).

The first term in the causal estimand can be directly estimated with this quantity. The second term,

Ex[Yc(x;∞)∣G(x)=g],\underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = g \right],xE​[Yc​(x;∞)∣G(x)=g],

is counterfactual because examples in BgB_gBg​ were in fact used for training. To address this issue, the authors derive two estimators that replace the unobserved counterfactual with observable proxy outcomes.

The first is the difference estimator. Under the assumption that training and validation examples are independently and identically distributed and that examples are randomly assigned to batches or validation sets, the expected performance on untreated validation examples approximates the expected counterfactual performance of treated examples. The corresponding estimator is:

τ^g,cdiff=Y‾c(g)−Y‾c(∞),\widehat{\tau}_{g,c}^{\mathrm{diff}} = \overline{Y}_c(g) - \overline{Y}_c(\infty),τg,cdiff​=Yc​(g)−Yc​(∞),

where B∞B_\inftyB∞​ denotes a held-out validation set. This estimator uses only observed outcomes at checkpoint ccc. It compares the average performance of a trained batch with the average performance of untrained validation examples. The method requires strong distributional matching between training and validation sets, and its variance can remain high because the batch mean Y‾c(g)\overline{Y}_c(g)Yc​(g) is restricted to examples in BgB_gBg​.

The main estimation procedure uses a difference-in-differences design. This approach uses the time dimension to identify the treatment effect by comparing how performance changes over checkpoints for treated examples and untreated examples. It relies on parallel trends and no anticipation. Parallel trends means that, in the absence of training, the expected change in performance across checkpoints would be the same for treated and untreated examples:

Ex[Yc(x;∞)−Yc′(x;∞)∣G(x)=g]=Ex[Yc(x;∞)−Yc′(x;∞)∣G(x)=∞].\underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};\infty) - Y_{c'}(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = g \right] = \underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};\infty) - Y_{c'}(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = \infty \right].xE​[Yc​(x;∞)−Yc′​(x;∞)∣G(x)=g]=xE​[Yc​(x;∞)−Yc′​(x;∞)∣G(x)=∞].

No anticipation means that training at timestep ggg cannot affect performance at any checkpoint c<gc<gc<g:

Ex[Yc(x;g)∣G(x)=g]=Ex[Yc(x;∞)∣G(x)=g].\underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};g) \mid G(\boldsymbol{x}) = g \right] = \underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = g \right].xE​[Yc​(x;g)∣G(x)=g]=xE​[Yc​(x;∞)∣G(x)=g].

Under these conditions, the target causal estimand can be rewritten using observable panel outcomes. The difference-in-differences estimand is:

τg,cdid=Ex[Yc(x;g)−Yg−1(x;g)∣G(x)=g]−Ex[Yc(x;∞)−Yg−1(x;∞)∣G(x)=∞].\tau_{g,c}^{\mathrm{did}} = \underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};g) - Y_{g-1}(\boldsymbol{x};g) \mid G(\boldsymbol{x}) = g \right] - \underset{\boldsymbol{x}}{\mathbb{E}}\left[ Y_c(\boldsymbol{x};\infty) - Y_{g-1}(\boldsymbol{x};\infty) \mid G(\boldsymbol{x}) = \infty \right].τg,cdid​=xE​[Yc​(x;g)−Yg−1​(x;g)∣G(x)=g]−xE​[Yc​(x;∞)−Yg−1​(x;∞)∣G(x)=∞].

In practice, the authors evaluate model performance for a subset of training batches and validation examples at the checkpoint immediately before treatment, g−1g-1g−1, and at a later checkpoint ccc. This creates a panel of observed outcomes, analogous to repeated measurements in causal inference. The difference-in-differences estimator is computed as:

τ^g,cdid=(Y‾c(g)−Y‾g−1(g))⏟change for trained batch−(Y‾c(∞)−Y‾g−1(∞))⏟change for validation set.\widehat{\tau}_{g,c}^{\mathrm{did}} = \underbrace{\left( \overline{Y}_c(g) - \overline{Y}_{g-1}(g) \right)}_{\text{change for trained batch}} - \underbrace{\left( \overline{Y}_c(\infty) - \overline{Y}_{g-1}(\infty) \right)}_{\text{change for validation set}}.τg,cdid​=change for trained batch(Yc​(g)−Yg−1​(g))​​−change for validation set(Yc​(∞)−Yg−1​(∞))​​.

This workflow uses only observed outcomes from a model checkpoint sequence. For each batch BgB_gBg​, the method measures how much the model’s performance on that batch changes after the batch becomes available for training, while subtracting the background change observed on examples that remain untrained. This difference controls for general improvements that happen over time across the model, rather than improvements caused by exposure to specific training batches.

The authors also describe influence functions as an alternative instance-level approximation for estimating counterfactual memorization at the final checkpoint. Influence functions estimate how parameters would change if an example were removed from the training set, without retraining:

θ−x,T≈θT+1NHθ−1∇θL(θT,x),\boldsymbol{\theta}_{-\boldsymbol{x},T} \approx \boldsymbol{\theta}_{T} + \frac{1}{N} \mathrm{H}_{\boldsymbol{\theta}}^{-1} \nabla_{\boldsymbol{\theta}} \mathcal{L}(\boldsymbol{\theta}_{T}, \boldsymbol{x}),θ−x,T​≈θT​+N1​Hθ−1​∇θ​L(θT​,x),

where Hθ\mathrm{H}_{\boldsymbol{\theta}}Hθ​ is the Hessian of the training loss evaluated at θT\boldsymbol{\theta}_{T}θT​. The estimated counterfactual performance is obtained by evaluating the example using the adjusted parameters:

Y−x,T(x)=γ(θ−x,T;x).Y_{-\boldsymbol{x},T}(\boldsymbol{x}) = \gamma(\boldsymbol{\theta}_{-\boldsymbol{x},T};\boldsymbol{x}).Y−x,T​(x)=γ(θ−x,T​;x).

The corresponding influence-function memorization estimate is:

τ^x,Tinfl=YT(x)−Y−x,T(x).\widehat{\tau}_{\boldsymbol{x},T}^{\mathrm{infl}} = Y_T(\boldsymbol{x}) - Y_{-\boldsymbol{x},T}(\boldsymbol{x}).τx,Tinfl​=YT​(x)−Y−x,T​(x).

This approach avoids repeated retraining and provides an instance-level estimate, but it depends on convexity, positive definiteness of the Hessian, and convergence assumptions that are generally weak for large language models. It also applies only at the final checkpoint, so it cannot characterize memorization dynamics across intermediate checkpoints.

Experiment

The experiments evaluate memorisation in Pythia language models ranging from 70M to 12B parameters using a subsampled panel of training instances measured across checkpoints with sequence-level log-likelihood. They find that instantaneous memorisation is stronger early in training and increases with model size, while persistent memorisation is weak or absent in small models and tends to plateau after a fixed number of steps. Residual memorisation is often statistically insignificant, though a recency effect appears near the end of training, and overall memorisation patterns across model sizes are broadly predictable.