Measuring memorization in RLHF for code completion
Aneesh Pappu, Billy Porter, Ilia Shumailov, Jamie Hayes
cs.LG, cs.CL, cs.SE
2024-06-18
Google DeepMind finds RLHF memorizes 0.9% of reward-model code examples versus 17.6% from fine-tuning, while IPO on Gemma 2B/7B memorizes 18.2%/19.5%.
Memorization during supervised fine-tuning is a mapped risk: feed a prefix, watch the model emit the rest of a training example. Production code-completion models almost always add a second stage, RLHF. Fine-tune, train a reward model on human preferences, then reinforce the policy against that reward. Whether that second stage copies preference data into the deployed model has barely been measured stage by stage.
Reward-model data is the sensitive slice. It often comes from live users and can contain file paths, keys, and internal names. If the aligned model later emits those strings, that is a privacy failure. Google DeepMind studies the question in code completion because the product is widely deployed, and because code is full of false positives: a model that is simply good at syntax can look like it memorized a specific file.
The workhorse model is Gemini Nano-1 (1.8B). To separate memory from competence, Gemini Ultra generated 6,554 Python snippets that all call a fake library, legitimatecrypto, which does not exist in pre-training. If the model writes that name, the source is this training run.
The synthetic set splits two ways:
An example counts as memorized only under 0.1-approximate counterfactual memorization: greedy decoding, normalized edit distance ≤ 0.1 against the candidate, and > 0.1 against a control that never saw the example. Edit distance uses a sliding window so a memorized span buried in extra tokens still counts. Targets shorter than 10 tokens, and prompt-target pairs with distance below 0.5, are dropped. Completing range(10 with ) should not count as memorization.
Two comparison axes sit next to a sweep on the KL coefficient α. One axis uses the same data as ordinary fine-tuning, an upper bound on memorization. The other trains IPO, a special case of ΨPO, on Gemma 2B and 7B so preference pairs update the policy directly, with no reward model in the middle. Larger α keeps the RL policy closer to its initialization.
What the fine-tune stage memorizes mostly survives RL. On filtered SD.Base, 319 of 3,526 examples were memorized after fine-tuning; 43%-47% of those still matched after RLHF. Smaller α lets the policy drift. PII-path emission on SD.Links falls from 54.8% at α=0.5 to 12.6% at α=0.005. The drop is real. Most of the memorized set remains.
Reward-model data barely transfers:
| Setup | SD.Base memorized | SD.Links PII emitted |
| Direct fine-tuning | 17.6% | 50% |
| Data only in the reward model, then RL | 0.9% | 0% |
| IPO (Gemma 2B / 7B) | 18.2% / 19.5% | 0.6% / 0.6% |
RL fine-tuning sees prompts only and discards targets. After 70 epochs with α=0.005, fewer than 0.5% of those prompts are memorized.
Scaling the reward model from T5-Base (220M) to Gemini Nano-1 and the policy from Nano-1 to Gemini Pro, then repeating on CodeXGLUE line completion, LIMA, and Anthropic HH, does not change the picture. Reward-model examples stay almost unmemorized. RL-prompt memorization stays under 1%.
A scalar reward is a low-bandwidth channel, and the KL term pins the policy near its start point. Together they make it hard for reward-model text to copy into the policy. IPO applies gradients to the preference text itself, and memorization returns to fine-tuning levels.
Washing out fine-tune memory by driving α very small has an alignment tax: RLFT.1 at α=0.005 drops HumanEval single-line infilling from 0.507 after fine-tuning to 0.217. In that regime, "no memorization" and "the run diverged" travel together.
If a team has expensive, sensitive preference data, putting it in the reward model, not in SFT and not in IPO-style direct preference learning, is the lowest-regurgitation option measured here, by about an order of magnitude. For a code-completion product, accept/reject logs can train the reward model without an automatic assumption that those strings will appear in later completions.
This is not a general safety theorem for RLHF. It is a controlled comparison on code completion, synthetic sensitive strings, and Gemini/Gemma-scale models. DPO was not evaluated, only IPO. The practical takeaway still holds: memorization risk does not transfer automatically from a claim that two alignment objectives are equivalent.
The sensitive strings were invented for the experiment, with 100 fixed path templates. Production logs do not arrive with a labeled list of what is sensitive, nor with that kind of controlled duplication. The true counterfactual definition would retrain a control model per example; the paper uses one fixed control that never saw the evaluation set, which is cheaper and slightly looser on false positives.
The IPO comparison switches model families (Gemma vs Gemini Nano-1), so the 18% vs 0.9% gap mixes algorithm with architecture and recipe. At α=0.005, a reward model trained on SD.Base collapsed the policy into always emitting #, because comments were 40%-50% of lines and every example was labeled positive.
Code requires the counterfactual filter. Without it, a strong completer inflates the memorization rate. Whether the same filters and sliding-window distance transfer to chat, long documents, or multi-turn settings is left open.