このリポジトリには、以下の論文のコードが含まれています。
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 億トークン中間トレーニング段階を再実行しました。この際、マッチしたペアからランダムに選んだ一連の文書をトレーニングストリームに挿入しました。同じ段階における AI2 の公開済み実行が対照群として機能します。2 つの実行は、文書を挿入した位置を除いて、同じトレーニングデータを同じ順序で読み込みました。その後、両方のモデルに AI2 のポストトレーニングレシピを適用しました。すべての文書は各段階で両方のモデルによってスコアリングされたため、1 つの文書に対するトレーニングがモデルをどれだけ変化させるか、そしてその変化のうちどれだけをメンバーシップ推論攻撃が検出できるかを測定できます。
メンバーシップ推論とは、ある文書がモデルのトレーニングデータに含まれていたかどうかを問うものです。これは著作権やデータの来歴、トレーニングデータの透明性、ベンチマークの汚染にとって重要です。言語モデルに対するメンバーシップ推論攻撃の評価のほとんどは、メンバーと非メンバーを事後的に収集します。たとえば、トレーニングのカットオフより前に公開された文書と後に公開された文書などです。すると 2 つのグループは、メンバーシップだけでなく、トピック、スタイル、日付においても異なります。攻撃はこの分布のずれを検出することで成功し得ます。攻撃が失敗したとしても、それはトレーニングが痕跡を残さなかったことを示すものではありません。
これらの説明を切り分けるには、メンバーシップが既知であり、あらゆる文書の属性から独立して割り当てられ、トレーニングの残りの部分を変えないような設定が必要です。本論文は次のように問います。実際の文書が確立されたトレーニングレシピに 1 回含められたとき、モデルはどれだけ変化するのか、そしてその変化のうちどれだけを既存の攻撃が利用できるのか。
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)
別段の記載がない限り、以下の数値は、未見であることを検証済みで一度だけ挿入された4,832組の文書ペアについて、対照に対して測定したものである。論文では信頼区間とさらなる注意事項が示されている。
これらの結果は、1つの1Bパラメータモデル、トークンの0.135%が置き換えられた1つの中間トレーニング段階、および我々が評価した攻撃統計に関するものである。参照比は分析前に指定された。その他の分析は後から追加されたものであり、探索的である。
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 も含まれています。
パスはハードコードされていません。すべてのスクリプトは、以下の順序で解決される単一のデータルートの下で入力と出力を見つけます。
--out)が存在する場合、それら;OLMO_TWIN_DATA;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 の公開入力用などです。シェルスクリプトは eval "$(python3 -m olmo_twin.paths --shell)" から(リポジトリルートで実行するか、pip install -e . の後に実行)、同じパスを OLMO_TWIN_<NAME> 変数として取得します。docs/data.md には各ディレクトリ、それを書き込むステップ、それを読み込むステップが一覧されています。
データリリース(利用可能になった時点)の分析テーブルとスコアをデータルートに置いた状態で、以下のいずれかのコマンドで論文のテーブル、マクロ、図を再生成できます: