OLMo-Twin

Questo repository contiene il codice per il 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 (di prossima pubblicazione).

Abbiamo rieseguito la fase ufficiale di mid-training da 50 miliardi di token di OLMo 2 1B con un insieme di documenti inseriti nel flusso di addestramento, scelti casualmente da coppie appaiate. La run rilasciata da AI2 della stessa fase funge da controllo. Le due run hanno letto gli stessi dati di addestramento nello stesso ordine, tranne che nelle posizioni in cui abbiamo inserito i documenti. Entrambi i modelli sono stati poi sottoposti alla ricetta di post-training di AI2. Ogni documento è stato valutato con entrambi i modelli in ogni fase, così possiamo misurare quanto l'addestramento su un singolo documento cambi il modello, e quanto di questo cambiamento gli attacchi di membership inference riescano a rilevare.

La domanda

La membership inference si chiede se un dato documento abbia fatto parte dei dati di addestramento di un modello. Questo è rilevante per il copyright e la provenienza dei dati, per la trasparenza sui dati di addestramento e per la contaminazione dei benchmark. La maggior parte delle valutazioni degli attacchi di membership inference sui modelli linguistici raccoglie membri e non-membri a posteriori, ad esempio documenti pubblicati prima e dopo una data di cutoff dell'addestramento. I due gruppi differiscono quindi per argomento, stile e data oltre che per l'appartenenza. Un attacco può avere successo rilevando questo spostamento di distribuzione. Quando un attacco fallisce, ciò non dimostra che l'addestramento non abbia lasciato traccia.

Per separare queste spiegazioni, abbiamo bisogno di un contesto in cui l'appartenenza sia nota, sia assegnata indipendentemente da ogni proprietà del documento e lasci invariato il resto dell'addestramento. Il paper si chiede: quando un documento reale viene incluso una volta in una ricetta di addestramento consolidata, quanto cambia il modello, e quanto di questo cambiamento possono sfruttare gli attacchi esistenti?

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

• Una ripetizione di una fase di addestramento ufficiale. OLMo 2 1B è addestrato in due fasi: circa 4 trilioni di token di pre-addestramento, poi una fase di mid-training da 50 miliardi di token. AI2 ha rilasciato il checkpoint in cui inizia il mid-training, incluso il suo stato dell'optimizer. AI2 ha anche rilasciato i dati del mid-training, la sua configurazione e le esecuzioni complete del mid-training per tre seed di ordinamento dei dati. Siamo partiti dal checkpoint rilasciato e abbiamo eseguito la fase di nuovo con la configurazione di AI2 e il trainer OLMo non modificato. Abbiamo ricostruito esattamente l'ordinamento dei dati del seed 42.
• Coppie abbinate e un lancio di moneta. Abbiamo costruito 7.814 coppie di documenti. I due documenti di una coppia provengono dalla stessa fonte e coorte e sono abbinati per lunghezza e difficoltà: bit per byte sotto il checkpoint iniziale, misurati prima dell'addestramento. Per ogni coppia, un lancio di moneta ha scelto un documento da inserire; l'altro è stato escluso. Copie dei documenti inseriti sostituiscono interi chunk di addestramento in posizioni pianificate distribuite lungo la fase: 12.414 copie sostituiscono 24.081 delle 12.212.224 posizioni dei chunk (0,197%), e il testo inserito ammonta a 67,3 milioni di token (0,135% del flusso). Ogni altra posizione contiene il chunk di AI2.
• L'esecuzione rilasciata come controllo. L'esecuzione rilasciata di AI2 con seed 42 ha visto gli stessi chunk nello stesso ordine in ogni altra posizione, quindi funge da controfattuale. Il paper la chiama il gemello. Le due esecuzioni differiscono nelle impostazioni di esecuzione (numero di GPU, kernel di attenzione e sharding), quindi non sono identiche bit per bit. L'esecuzione di AI2 con seed 42069 (la stessa ricetta in un diverso ordinamento dei dati) funge da riferimento riaddestrato. Il checkpoint rilasciato da cui inizia la fase funge da riferimento più debole.
• Documenti verificati come non visti. La maggior parte dei documenti sono paper arXiv, pagine Wikipedia e post StackExchange apparsi dopo la data di cutoff dei dati del modello, o pagine Wikipedia mancanti dallo snapshot nei dati di addestramento. Tredici coorti variano la fonte, la lunghezza, il numero di copie (fino a 32), la tempistica delle copie e la forma del testo (derivata da LaTeX versus derivata da PDF). Versioni modificate di pagine Wikipedia presenti nei dati di addestramento misurano la riesposizione, e i testi sintetici coprono una gamma di surprisal. I candidati delle coorti destinate a essere non visti sono stati cercati in modo esaustivo contro tutte le 13 componenti dei corpora di pre-addestramento e mid-training rilasciati (8,3 TiB). La ricerca rileva ogni sovrapposizione verbatim di 40 parole o più, e i candidati con tali sovrapposizioni sono stati esclusi. L'ultima parte della ricerca è terminata dopo che le coppie erano state formate. Ha trovato altri 156 candidati con tali sovrapposizioni; sono stati rimossi, e le loro coppie sono state riempite prima dell'addestramento.
• Post-addestramento. Sia il modello trattato sia il controllo sono stati sottoposti alla ricetta di post-addestramento di AI2 per OLMo 2 1B: supervised fine-tuning (SFT), direct preference optimization (DPO) e due fasi di reinforcement learning con reward verificabili (RLVR). Entrambi i modelli hanno ricevuto gli stessi dati, ordinamento e seed.
• Misurazioni. I 22.681 documenti di valutazione (entrambi i documenti di ogni coppia e 7.053 documenti che non sono mai stati inseriti) sono valutati token per token sotto:
  • il modello trattato, il controllo, il riferimento riaddestrato e il checkpoint iniziale;
  • i 24 checkpoint salvati ogni 1.000 aggiornamenti da ciascuna delle tre esecuzioni di mid-training;
  • i modelli di ogni fase di post-addestramento di entrambi i rami.

Risultati principali

Salvo diversa indicazione, i numeri seguenti si riferiscono alle 4.832 coppie di documenti verificate come non viste e inserite una volta, misurate rispetto al controllo. Il paper fornisce intervalli di confidenza e ulteriori avvertenze.