OLMo-Twin

本仓库包含论文的代码

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).

我们重新运行了 OLMo 2 1B 官方的 500 亿 token 中期训练阶段，并在训练流中插入了一组文档，这些文档从匹配对中随机选取。AI2 发布的同一阶段运行作为对照组。两次运行以相同顺序读取相同的训练数据，唯一区别在于我们插入文档的位置。随后，两个模型都经过了 AI2 的后训练流程。每个文档在每个阶段都在两个模型下进行了评分，因此我们可以衡量对单个文档的训练使模型发生了多大变化，以及成员推断攻击能够检测到其中多少变化。

问题

成员推断旨在判断给定文档是否属于模型的训练数据。这对于版权与数据溯源、训练数据透明度以及基准污染都很重要。大多数针对语言模型的成员推断攻击评估都是在事后收集成员和非成员，例如在训练截止日期之前和之后发布的文档。这样两组不仅在成员身份上不同，在主题、风格和日期上也不同。攻击可以通过检测这种分布偏移而成功。当攻击失败时，这并不能表明训练没有留下痕迹。

为了区分这些解释，我们需要一个成员身份已知、且其分配与文档的任何属性无关、并且不改变训练其余部分的设置。论文提出的问题是：当一个真实文档被一次性纳入一个成熟的训练配方时，模型会发生多大变化，以及现有攻击能够利用其中多少变化？

设计

                    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)

• 官方训练阶段的重跑。 OLMo 2 1B 分两个阶段训练：约 4 万亿 token 的预训练，然后是 500 亿 token 的中间训练阶段。AI2 发布了中间训练开始时的检查点，包括其优化器状态。AI2 还发布了中间训练数据、其配置，以及三个数据顺序种子的完整中间训练运行。我们从发布的检查点开始，使用 AI2 的配置和未经修改的 OLMo 训练器重新运行了该阶段。我们精确重建了 seed-42 的数据顺序。
• 配对与抛硬币。 我们构建了 7,814 对文档。一对中的两个文档来自相同的来源和队列，并在长度和难度上匹配：起始检查点下的每字节比特数，在训练前测量。对于每一对，抛硬币选择其中一个文档用于插入；另一个被保留。插入文档的副本在分布在整个阶段的计划位置上替换整个训练块：12,414 个副本替换了 12,212,224 个块位置中的 24,081 个（0.197%），插入的文本总计 6730 万 token（占数据流的 0.135%）。其他每个位置都保留 AI2 的块。
• 发布的运行作为对照。 AI2 发布的 seed-42 运行在其他每个位置都以相同顺序看到了相同的块，因此它作为反事实。论文称之为孪生。两次运行在执行设置（GPU 数量、注意力和分片内核）上不同，因此它们并非逐位相同。AI2 的 seed-42069 运行（相同配方，不同数据顺序）作为重新训练的参考。阶段开始时发布的检查点作为较弱的参考。
• 文档已验证未见。 大多数文档是模型数据截止日期之后出现的 arXiv 论文、维基百科页面和 StackExchange 帖子，或者是训练数据快照中缺失的维基百科页面。十三个队列在来源、长度、副本数量（最多 32 个）、副本的时间以及文本形式（LaTeX 派生与 PDF 派生）上有所不同。训练数据中维基百科页面的编辑版本用于测量再暴露，合成文本跨越一系列惊奇度。对旨在未见的队列的候选文档，针对发布的预训练和中间训练语料库的所有 13 个组件（8.3 TiB）进行了穷举搜索。该搜索检测到所有 40 词或以上的逐字重叠，具有此类重叠的候选文档被排除。搜索的最后部分在配对形成之后完成。它又发现了 156 个具有此类重叠的候选文档；它们被移除，其配对在训练前被重新填充。
• 后训练。 处理模型和对照模型都经过了 AI2 的 OLMo 2 1B 后训练配方：监督微调（SFT）、直接偏好优化（DPO），以及两个阶段的可验证奖励强化学习（RLVR）。两个模型获得了相同的数据、顺序和种子。
• 测量。 22,681 个评估文档（每对的两个文档以及 7,053 个从未插入的文档）在以下模型下逐 token 评分：
  • 处理模型、对照、重新训练的参考和起始检查点；
  • 三次中间训练运行每 1,000 次更新保存的 24 个检查点；
  • 两个分支每个后训练阶段的模型。

主要发现

除非另有说明，以下数字针对 4,832 对已验证未见且插入一次的文档，与对照进行比较。论文给出了置信区间和进一步的注意事项。

• 一次曝光会留下微小但一致的痕迹。 对每篇文档而言，其增益是指从对照模型到处理模型之间，平均每 token 对数似然的变化。平均而言，一篇被插入的文档比其留出的配对文档多获得 0.0221 nats/token。在 85.4% 的配对中，被插入文档的增益更大。
• 该痕迹很少能识别出单篇文档。 以对照模型为参照，参考比值的 AUC 达到 0.767。在跨域使用单一阈值时，它以 1% 的假阳性率检测出 3.8% 的成员。按每篇文档的 token 级噪声对其比值进行标准化后，这一比例提升至 8.8%。为每个域分别设定阈值可将这两个比率提升至 16.5% 和 31.4%，但这需要来自同一域的非成员。
• 汇总证据可检测出对某个集合的训练。 以对照模型为参照，对 13 个匹配配对进行配对检验，在 5% 显著性水平下估计功效为 80%。该检验需要为每个成员匹配一个非成员，因此这并非任意集合级审计的性能。
• 重新训练的参照几乎与对照模型一样好；无参照攻击则接近随机水平。 作为参照，AI2 的 seed-42069 运行恢复了对照模型相对于随机水平的 AUC 优势的 92%，并需要 14 个配对才能达到相同功效。发布的起始检查点需要 87 个配对。针对难度匹配的非成员，无参照攻击在跨域汇总时与随机水平的 AUC 差距保持在 0.07 以内。
• 从未被插入的文档也会发生变化。 从未插入的 arXiv 论文获得 0.0421 nats/token，这超过了被插入 arXiv 论文的配对内效应（0.0232）。因此，没有匹配非成员的比较可能被整个域层面的偏移所主导。
• 重复、遗忘与后训练。 该效应随副本数量近乎线性增长（拟合指数为 1.00，最多 32 个副本），在 1% 假阳性率下，大多数成员在四到八个副本时变得可检测。在前 4,000 次更新中插入的文档，在中期训练结束时最多保留其测得峰值效应的 21%。经过完整的后训练流程后，在对应于发布的 OLMo 2 1B Instruct 的检查点上，89% 的配对内效应和 74% 的相对于随机水平的 AUC 优势仍然存在。

这些结果涉及一个 1B 参数模型、一个替换了 0.135% token 的中期训练阶段，以及我们所评估的攻击统计量。参考比值是在分析之前指定的；其他分析是后来添加的，属于探索性分析。

仓库布局

OLMo-Twin/
├── olmo_twin/        small shared package; paths.py resolves the data root
├── collection/       candidate documents and their search against the OLMo 2 training data
├── design/           matched pairs, coin flips, copy positions: the frozen design
├── training/
│   ├── midtrain/     AI2's inputs, the training stream with the inserted documents, the launcher
│   ├── checks/       checks of the finished run (data order, inserted copies, configuration)
│   └── posttrain/    SFT, DPO and RLVR with AI2's open-instruct, for both models
├── scoring/          per-token log-probabilities under every model and checkpoint
├── analysis/         membership statistics, attacks, trajectories, post-training analysis
├── figures/          the paper's tables, numbers and figures
├── env/              software environments (requirements files, Dockerfiles)
└── docs/             reproduction guide and data documentation

每个目录都有一个 README，详细描述其对应的步骤。 docs/REPRODUCING.md 按顺序列出各步骤，并说明每一步的输入和输出。docs/data.md 描述数据布局、计划中的数据发布以及许可证。 docs/paper_outputs.md 将论文中的每个图和表映射到生成它的命令。

快速开始

安装

分析、图表和检查均在 CPU 上运行。env/analysis/requirements.txt 中固定的版本需要 Python 3.12 或更高版本。

git clone https://github.com/ali-bana/OLMo-Twin.git
cd OLMo-Twin
python3 -m venv .venv && . .venv/bin/activate
pip install -r env/analysis/requirements.txt
pip install -e .

将仓库作为检出使用。你可以在不安装任何东西的情况下运行脚本，或以可编辑模式安装该包（pip install -e .）。不支持常规的 pip install .，因为该包会从自身位置推导出仓库根目录，并由此推导出默认数据根目录。文档集合、训练运行和评分使用其他环境；env/README.md 描述了这些环境，包括用于 GPU 镜像的 Dockerfile。

数据所在位置

没有任何路径是硬编码的。每个脚本都在一个数据根目录下查找其输入和输出，解析顺序如下：

1. 脚本自身的命令行选项（例如 --out），前提是脚本提供了这些选项；
2. 环境变量 OLMO_TWIN_DATA；
3. 仓库中的 data/（被 git 忽略；它可能是一个符号链接）。

export OLMO_TWIN_DATA=/path/to/olmo-twin-data
python -m olmo_twin.paths              # prints the resolved layout and which directories exist

数据根目录具有固定布局：bundle/ 用于存放冻结的设计，analysis/ 用于存放分析表，run/ 用于存放分数和检查点，assets/ 用于存放 AI2 的公开输入，等等。Shell 脚本通过 eval "$(python3 -m olmo_twin.paths --shell)" 获取与 OLMO_TWIN_<NAME> 变量相同的路径（从仓库根目录运行，或在执行 pip install -e . 之后运行）。docs/data.md 列出了每个目录、写入该目录的步骤以及读取该目录的步骤。

重新生成论文的表格和图表

在数据根目录中准备好数据发布版（一旦可用）的分析表和分数后，以下任一命令均可重新生成论文的表格、宏和图表：

bash figures/make_all.sh               # from the analysis outputs -> $OLMO_TWIN_DATA/paper/
bash analysis/run_analysis.sh          # re-runs the analysis first, then the tables and figures,
                                       # Table 4 on all pairs and the RLVR results

两个脚本都会写入数据根目录，并就地覆盖先前的输出。 如果你想保留原始文件，请在数据发布副本上操作。 run_analysis.sh 会跳过需要文档文本或 token id 的步骤（并给出提示），这些内容并非为每个文档都重新分发。其他步骤则在已发布的数据上运行。

测试

python -m pytest -p no:cacheprovider design/tests training/midtrain/tests
# collection/: in the collection environment (numba, beautifulsoup4, wikiextractor 3.0.8)
# with pytest added, one test file per process
for f in collection/tests/test_*.py; do python -m pytest -q -p no:cacheprovider "$f"; done

复现实验

完整流程从文档收集一直运行到论文图表：收集、设计、中期训练、检查、后训练、评分、分析和图表。docs/REPRODUCING.md 描述了每个阶段、其输入和输出、命令，以及无法仅凭本仓库复现的部分。脚本中不包含调度器或集群命令。多节点启动器是普通的 torchrun、accelerate launch 和 ray 命令，它们从环境变量中获取节点排名和第一个节点的地址，因此可以由任何调度器或手动启动。

阶段	论文使用的计算资源
收集与重叠搜索	CPU、网络访问、一个 GROBID 服务器；OLMo 2 训练语料库的本地副本（8.3 TiB）；一个 GPU 用于嵌入搜索
中期训练	16 块 A100 80GB GPU（两个节点，每节点八块）
SFT、DPO	每个模型 8 块 A100 80GB GPU
RLVR	每个阶段和模型 16 块 A100 GPU（两个节点）
评分	每个模型和检查点一块 GPU（在一个 A100 上每个模型约 46 分钟）
分析、图表、检查	CPU