OLMo-Twin

This repository contains the code for the paper

From Training Exposure to Membership Evidence: A Controlled Study of OLMo 2
Ali Banaei, Muhammad Shahab Shahzad, Gerald Q. Maguire Jr., Marco Chiesa, Dejan Kostić
KTH Royal Institute of Technology. arXiv preprint (to appear).

We re-ran the official 50-billion-token mid-training stage of OLMo 2 1B with a set of documents inserted into the training stream, chosen at random from matched pairs. AI2's released run of the same stage serves as the control. The two runs read the same training data in the same order, except at the positions where we inserted documents. Both models were then put through AI2's post-training recipe. Every document was scored under both models at every stage, so we can measure how much training on one document changes the model, and how much of that change membership-inference attacks can detect.

The question

Membership inference asks whether a given document was part of a model's training data. This matters for copyright and data provenance, for training-data transparency and for benchmark contamination. Most evaluations of membership-inference attacks on language models collect members and non-members after the fact, for example documents published before and after a training cutoff. The two groups then differ in topic, style and date as well as in membership. An attack can succeed by detecting this distribution shift. When an attack fails, that does not show that training left no trace.

To separate these explanations, we need a setting in which membership is known, is assigned independently of every document property, and leaves the rest of training unchanged. The paper asks: when a real document is included once in an established training recipe, how much does the model change, and how much of that change can existing attacks use?

The design

                    AI2's released stage-1 checkpoint (4T tokens)
                                       |
            +--------------------------+---------------------------+
            |                                                      |
  AI2's released seed-42 run                     our re-run of the same stage:
  of the 50B-token stage                         same start, data and order;
  (23,852 updates)                               24,081 of 12,212,224 chunks replaced
            |                                    by copies of inserted documents
            v                                                      v
         control                                             treated model
            |                                                      |
            v                                                      v
   SFT -> DPO -> RLVR stage 1 -> RLVR stage 2  (same data, order and seeds for both)

• A re-run of an official training stage. OLMo 2 1B is trained in two stages: about 4 trillion tokens of pretraining, then a 50-billion-token mid-training stage. AI2 released the checkpoint at which mid-training starts, including its optimizer state. AI2 also released the mid-training data, its configuration, and the complete mid-training runs for three data-order seeds. We started from the released checkpoint and ran the stage again with AI2's configuration and the unmodified OLMo trainer. We reconstructed the seed-42 data order exactly.
• Matched pairs and a coin flip. We built 7,814 pairs of documents. The two documents of a pair come from the same source and cohort and are matched on length and on difficulty: bits per byte under the starting checkpoint, measured before training. For each pair, a coin flip chose one document for insertion; the other was held out. Copies of the inserted documents replace whole training chunks at planned positions spread over the stage: 12,414 copies replace 24,081 of the 12,212,224 chunk positions (0.197%), and the inserted text adds up to 67.3 million tokens (0.135% of the stream). Every other position holds AI2's chunk.
• The released run as the control. AI2's released seed-42 run saw the same chunks in the same order at every other position, so it serves as the counterfactual. The paper calls it the twin. The two runs differ in execution settings (number of GPUs, attention and sharding kernels), so they are not bitwise identical. AI2's seed-42069 run (the same recipe in a different data order) serves as a retrained reference. The released checkpoint at which the stage starts serves as a weaker reference.
• Documents verified unseen. Most documents are arXiv papers, Wikipedia pages and StackExchange posts that appeared after the model's data cutoff, or Wikipedia pages missing from the snapshot in the training data. Thirteen cohorts vary the source, the length, the number of copies (up to 32), the timing of the copies, and the text form (LaTeX-derived versus PDF-derived). Edited versions of Wikipedia pages that are in the training data measure re-exposure, and synthetic texts span a range of surprisal. The candidates of the cohorts meant to be unseen were searched exhaustively against all 13 components of the released pretraining and mid-training corpora (8.3 TiB). The search detects every verbatim overlap of 40 words or more, and candidates with such overlaps were excluded. The last part of the search finished after the pairs had been formed. It found 156 more candidates with such overlaps; they were removed, and their pairs were refilled before training.
• Post-training. Both the treated model and the control were put through AI2's post-training recipe for OLMo 2 1B: supervised fine-tuning (SFT), direct preference optimization (DPO), and two stages of reinforcement learning with verifiable rewards (RLVR). Both models got the same data, order and seeds.
• Measurements. The 22,681 evaluation documents (both documents of every pair and 7,053 documents that were never inserted) are scored token by token under:
  • the treated model, the control, the retrained reference and the starting checkpoint;
  • the 24 checkpoints saved every 1,000 updates by each of the three mid-training runs;
  • the models of every post-training stage of both arms.

Main findings

Unless stated otherwise, the numbers below are for the 4,832 pairs of documents verified unseen and inserted once, measured against the control. The paper gives confidence intervals and further caveats.