OLMo-Twin

Ce dépôt contient le code de l'article

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. Prépublication arXiv (à paraître).

Nous avons réexécuté l'étape officielle de mi-entraînement de 50 milliards de tokens d'OLMo 2 1B avec un ensemble de documents insérés dans le flux d'entraînement, choisis au hasard parmi des paires appariées. L'exécution publiée par AI2 de la même étape sert de contrôle. Les deux exécutions ont lu les mêmes données d'entraînement dans le même ordre, sauf aux positions où nous avons inséré des documents. Les deux modèles ont ensuite été soumis à la recette de post-entraînement d'AI2. Chaque document a été évalué sous les deux modèles à chaque étape, ce qui nous permet de mesurer dans quelle mesure l'entraînement sur un document modifie le modèle, et quelle part de ce changement les attaques d'inférence d'appartenance peuvent détecter.

La question

L'inférence d'appartenance consiste à déterminer si un document donné faisait partie des données d'entraînement d'un modèle. Cela importe pour le droit d'auteur et la provenance des données, pour la transparence des données d'entraînement et pour la contamination des benchmarks. La plupart des évaluations des attaques d'inférence d'appartenance sur les modèles de langage collectent les membres et les non-membres a posteriori, par exemple des documents publiés avant et après une date limite d'entraînement. Les deux groupes diffèrent alors par le sujet, le style et la date ainsi que par l'appartenance. Une attaque peut réussir en détectant ce décalage de distribution. Lorsqu'une attaque échoue, cela ne prouve pas que l'entraînement n'a laissé aucune trace.

Pour séparer ces explications, nous avons besoin d'un cadre dans lequel l'appartenance est connue, est attribuée indépendamment de toute propriété du document, et laisse le reste de l'entraînement inchangé. L'article pose la question : lorsqu'un document réel est inclus une fois dans une recette d'entraînement établie, dans quelle mesure le modèle change-t-il, et quelle part de ce changement les attaques existantes peuvent-elles exploiter ?

La conception

                    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)

• Une réexécution d'une étape d'entraînement officielle. OLMo 2 1B est entraîné en deux étapes : environ 4 billions de tokens de pré-entraînement, puis une étape de mid-training de 50 milliards de tokens. AI2 a publié le checkpoint auquel commence le mid-training, y compris l'état de son optimiseur. AI2 a également publié les données de mid-training, leur configuration et les exécutions complètes de mid-training pour trois graines d'ordre des données. Nous sommes partis du checkpoint publié et avons réexécuté l'étape avec la configuration d'AI2 et l'entraîneur OLMo non modifié. Nous avons reconstruit exactement l'ordre des données de la graine 42.
• Des paires appariées et un tirage au sort. Nous avons construit 7 814 paires de documents. Les deux documents d'une paire proviennent de la même source et de la même cohorte et sont appariés sur la longueur et sur la difficulté : bits par octet sous le checkpoint de départ, mesurés avant l'entraînement. Pour chaque paire, un tirage au sort a choisi un document à insérer ; l'autre a été mis de côté. Des copies des documents insérés remplacent des chunks d'entraînement entiers à des positions planifiées réparties sur l'étape : 12 414 copies remplacent 24 081 des 12 212 224 positions de chunk (0,197 %), et le texte inséré totalise 67,3 millions de tokens (0,135 % du flux). Toutes les autres positions contiennent le chunk d'AI2.
• L'exécution publiée comme contrôle. L'exécution publiée d'AI2 pour la graine 42 a vu les mêmes chunks dans le même ordre à toutes les autres positions, elle sert donc de contrefactuel. L'article l'appelle le jumeau. Les deux exécutions diffèrent par leurs paramètres d'exécution (nombre de GPU, kernels d'attention et de sharding), elles ne sont donc pas identiques bit à bit. L'exécution d'AI2 pour la graine 42069 (la même recette dans un ordre de données différent) sert de référence réentraînée. Le checkpoint publié auquel commence l'étape sert de référence plus faible.
• Documents vérifiés comme non vus. La plupart des documents sont des articles arXiv, des pages Wikipédia et des posts StackExchange apparus après la date de coupure des données du modèle, ou des pages Wikipédia absentes de l'instantané des données d'entraînement. Treize cohortes font varier la source, la longueur, le nombre de copies (jusqu'à 32), le moment des copies et la forme du texte (dérivé de LaTeX versus dérivé de PDF). Des versions modifiées de pages Wikipédia présentes dans les données d'entraînement mesurent la réexposition, et des textes synthétiques couvrent une gamme de surprisal. Les candidats des cohortes censées être non vus ont été recherchés de manière exhaustive dans les 13 composantes des corpus de pré-entraînement et de mid-training publiés (8,3 TiB). La recherche détecte tout chevauchement mot pour mot de 40 mots ou plus, et les candidats présentant de tels chevauchements ont été exclus. La dernière partie de la recherche s'est terminée après la formation des paires. Elle a trouvé 156 candidats supplémentaires présentant de tels chevauchements ; ils ont été retirés, et leurs paires ont été recomplétées avant l'entraînement.
• Post-entraînement. Le modèle traité et le contrôle ont tous deux été soumis à la recette de post-entraînement d'AI2 pour OLMo 2 1B : fine-tuning supervisé (SFT), optimisation directe des préférences (DPO) et deux étapes d'apprentissage par renforcement avec récompenses vérifiables (RLVR). Les deux modèles ont reçu les mêmes données, le même ordre et les mêmes graines.
• Mesures. Les 22 681 documents d'évaluation (les deux documents de chaque paire et 7 053 documents jamais insérés) sont notés token par token sous :
  • le modèle traité, le contrôle, la référence réentraînée et le checkpoint de départ ;
  • les 24 checkpoints sauvegardés toutes les 1 000 mises à jour par chacune des trois exécutions de mid-training ;
  • les modèles de chaque étape de post-entraînement des deux bras.

Principaux résultats

Sauf indication contraire, les chiffres ci-dessous concernent les 4 832 paires de documents vérifiés comme non vus et insérés une fois, mesurés par rapport au contrôle. L'article donne les intervalles de confiance et d'autres mises en garde.