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 (게재 예정).

우리는 OLMo 2 1B의 공식 500억 토큰 중간 학습(mid-training) 단계를 다시 실행했으며, 이때 일치하는 쌍(matched pairs)에서 무작위로 선택한 일련의 문서들을 학습 스트림에 삽입했습니다. 동일한 단계에 대한 AI2의 공개 실행이 대조군 역할을 합니다. 두 실행은 문서를 삽입한 위치를 제외하고는 동일한 학습 데이터를 동일한 순서로 읽었습니다. 그런 다음 두 모델 모두 AI2의 사후 학습(post-training) 레시피를 거쳤습니다. 모든 문서는 모든 단계에서 두 모델 모두에 대해 채점되었으므로, 하나의 문서에 대한 학습이 모델을 얼마나 변화시키는지, 그리고 그 변화 중 얼마나 많은 부분을 멤버십 추론 공격이 탐지할 수 있는지 측정할 수 있습니다.

질문

멤버십 추론은 주어진 문서가 모델의 학습 데이터에 포함되었는지를 묻습니다. 이는 저작권과 데이터 출처, 학습 데이터 투명성, 그리고 벤치마크 오염 문제에서 중요합니다. 언어 모델에 대한 멤버십 추론 공격의 대부분의 평가는 사후에 멤버와 비멤버를 수집합니다. 예를 들어 학습 컷오프 이전과 이후에 게시된 문서를 사용하는 식입니다. 그러면 두 그룹은 멤버십뿐만 아니라 주제, 문체, 날짜에서도 차이가 납니다. 공격은 이러한 분포 이동을 탐지함으로써 성공할 수 있습니다. 공격이 실패한다고 해서 학습이 아무 흔적도 남기지 않았다는 것을 보여주지는 않습니다.

이러한 설명들을 분리하기 위해서는 멤버십이 알려져 있고, 모든 문서 속성과 독립적으로 할당되며, 나머지 학습은 변경하지 않는 설정이 필요합니다. 이 논문은 다음과 같이 묻습니다. 실제 문서가 확립된 학습 레시피에 한 번 포함될 때, 모델은 얼마나 변화하며, 그 변화 중 얼마나 많은 부분을 기존 공격이 활용할 수 있는가?

설계

                    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조 토큰의 사전 학습, 그다음 500억 토큰의 중간 학습 단계다. AI2는 중간 학습이 시작되는 체크포인트를 옵티마이저 상태를 포함해 공개했다. AI2는 또한 중간 학습 데이터, 그 구성, 그리고 세 가지 데이터 순서 시드에 대한 전체 중간 학습 실행을 공개했다. 우리는 공개된 체크포인트에서 시작해 AI2의 구성과 수정되지 않은 OLMo 트레이너로 해당 단계를 다시 실행했다. 우리는 시드 42의 데이터 순서를 정확히 재구성했다.
• 짝지은 쌍과 동전 던지기. 우리는 7,814쌍의 문서를 만들었다. 한 쌍의 두 문서는 동일한 출처와 코호트에서 나오며 길이와 난이도(학습 전에 측정한 시작 체크포인트 기준 바이트당 비트 수)가 일치하도록 짝지어졌다. 각 쌍마다 동전 던지기로 한 문서를 삽입용으로 선택했고, 다른 문서는 제외했다. 삽입된 문서의 사본은 단계 전반에 걸쳐 계획된 위치에서 전체 학습 청크를 대체한다. 12,414개의 사본이 12,212,224개 청크 위치 중 24,081개(0.197%)를 대체하며, 삽입된 텍스트는 총 6,730만 토큰(스트림의 0.135%)에 달한다. 다른 모든 위치에는 AI2의 청크가 들어 있다.
• 대조군으로서의 공개된 실행. AI2가 공개한 시드 42 실행은 다른 모든 위치에서 동일한 순서로 동일한 청크를 보았으므로 반사실적 대조군 역할을 한다. 논문에서는 이를 트윈이라고 부른다. 두 실행은 실행 설정(GPU 수, 어텐션 및 샤딩 커널)이 달라 비트 단위로 동일하지는 않다. AI2의 시드 42069 실행(동일한 레시피를 다른 데이터 순서로 실행)은 재학습된 참조로 사용된다. 단계가 시작되는 공개된 체크포인트는 더 약한 참조로 사용된다.
• 보지 못한 것으로 검증된 문서. 대부분의 문서는 모델의 데이터 컷오프 이후에 등장한 arXiv 논문, 위키백과 문서, StackExchange 게시물이거나, 학습 데이터의 스냅샷에서 누락된 위키백과 문서다. 13개의 코호트는 출처, 길이, 사본 수(최대 32), 사본의 시점, 텍스트 형식(LaTeX 기반 대 PDF 기반)을 다양하게 한다. 학습 데이터에 포함된 위키백과 문서의 편집본은 재노출을 측정하며, 합성 텍스트는 다양한 서프라이절 범위를 포괄한다. 보지 못한 것으로 의도된 코호트의 후보들은 공개된 사전 학습 및 중간 학습 코퍼스의 13개 구성 요소 전체(8.3 TiB)에 대해 철저히 검색되었다. 이 검색은 40단어 이상의 모든 축자적 중복을 탐지하며, 그러한 중복이 있는 후보는 제외되었다. 검색의 마지막 부분은 쌍이 형성된 후에 완료되었다. 이 검색에서 그러한 중복이 있는 후보 156개가 추가로 발견되었고, 이들은 제거되었으며 그 쌍들은 학습 전에 다시 채워졌다.
• 사후 학습. 처리된 모델과 대조군 모두 OLMo 2 1B에 대한 AI2의 사후 학습 레시피를 거쳤다. 지도 미세 조정(SFT), 직접 선호 최적화(DPO), 그리고 검증 가능한 보상을 사용한 두 단계의 강화 학습(RLVR)이다. 두 모델은 동일한 데이터, 순서, 시드를 받았다.
• 측정. 22,681개의 평가 문서(모든 쌍의 두 문서와 한 번도 삽입되지 않은 7,053개 문서)는 다음 모델들에서 토큰 단위로 채점된다.
  • 처리된 모델, 대조군, 재학습된 참조, 시작 체크포인트;
  • 세 번의 중간 학습 실행 각각이 1,000회 업데이트마다 저장한 24개의 체크포인트;
  • 두 갈래 모두의 모든 사후 학습 단계 모델.

주요 발견

달리 명시되지 않는 한, 아래 수치는 보지 못한 것으로 검증되어 한 번 삽입된 4,832쌍의 문서에 대한 것이며 대조군과 비교해 측정한 것이다. 논문은 신뢰 구간과 추가적인 주의 사항을 제시한다.

• 한 번의 노출은 작고 일관된 흔적을 남긴다. 각 문서에 대해 *이득(gain)*은 대조 모델에서 처리 모델로 갈 때 평균 토큰당 로그 우도(log-likelihood)의 변화량이다. 평균적으로 삽입된 문서는 보류된(held-out) 짝 문서보다 0.0221 nats/token 더 큰 이득을 얻는다. 삽입된 문서가 더 큰 이득을 보이는 비율은 짝의 85.4%이다.
• 그 흔적이 단일 문서를 식별하는 경우는 드물다. 대조 모델을 기준으로 할 때, 참조 비율(reference ratio)은 AUC 0.767에 도달한다. 도메인 전반에 걸쳐 하나의 임계값을 적용하면, 1% 위양성률에서 멤버의 3.8%를 탐지한다. 각 문서의 비율을 토큰 수준 노이즈로 표준화하면 이 수치는 8.8%로 올라간다. 도메인별로 임계값을 따로 설정하면 두 비율이 16.5%와 31.4%로 올라가지만, 이는 같은 도메인의 비멤버를 필요로 한다.
• 풀링된 증거는 컬렉션에 대한 학습을 탐지한다. 대조 모델을 기준으로 할 때, 13개의 짝지어진 쌍에 대한 대응 검정(paired test)은 5% 수준에서 추정 검정력 80%를 가진다. 이 검정은 모든 멤버에 대해 짝지어진 비멤버를 필요로 하므로, 이는 임의의 컬렉션 수준 감사의 성능이 아니다.
• 재학습된 참조 모델은 대조 모델만큼 거의 좋으며, 참조 없는 공격은 우연 수준에 가깝다. 참조 모델로서 AI2의 seed-42069 실행은 대조 모델이 우연 대비 가지는 AUC 이점의 92%를 회복하며, 같은 검정력을 위해 14쌍이 필요하다. 공개된 시작 체크포인트는 87쌍이 필요하다. 난이도가 매칭된 비멤버에 대해서는, 참조 없는 공격이 도메인 전반에 걸쳐 풀링될 때 우연 대비 AUC 0.07 이내에 머문다.
• 삽입되지 않은 문서들도 움직인다. 삽입된 적 없는 arXiv 논문들은 0.0421 nats/token의 이득을 얻는데, 이는 삽입된 arXiv 논문들의 쌍 내 효과(0.0232)보다 크다. 따라서 매칭된 비멤버가 없는 비교는 전체 도메인 수준의 이동에 의해 지배될 수 있다.
• 반복, 망각, 그리고 사후 학습(post-training). 이 효과는 복사본 수에 거의 선형으로 증가하며(적합된 지수 1.00, 최대 32개 복사본), 대부분의 멤버는 4~8개의 복사본으로 1% 위양성률에서 탐지 가능해진다. 처음 4,000회 업데이트에서 삽입된 문서들은 중간 학습(mid-training) 종료 시점에 측정된 최대 효과의 최대 21%만을 유지한다. 전체 사후 학습 레시피를 거친 후, 공개된 OLMo 2 1B Instruct에 해당하는 체크포인트에서는 쌍 내 효과의 89%와 우연 대비 AUC 이점의 74%가 남아 있다.

이 결과들은 하나의 1B 파라미터 모델, 토큰의 0.135%가 교체된 하나의 중간 학습 단계, 그리고 우리가 평가한 공격 통계에 관한 것이다. 참조 비율은 분석 전에 지정되었으며, 다른 분석들은 나중에 추가된 탐색적 분석이다.

저장소 구조

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