OLMo-Twin

Este repositorio contiene el código del artículo

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. Preprint de arXiv (próxima aparición).

Repetimos la etapa oficial de mid-training de 50 mil millones de tokens de OLMo 2 1B con un conjunto de documentos insertados en el flujo de entrenamiento, elegidos al azar de pares emparejados. La ejecución publicada por AI2 de la misma etapa sirve como control. Las dos ejecuciones leyeron los mismos datos de entrenamiento en el mismo orden, excepto en las posiciones donde insertamos documentos. Después, ambos modelos se sometieron a la receta de post-entrenamiento de AI2. Cada documento fue puntuado bajo ambos modelos en cada etapa, de modo que podemos medir cuánto cambia el modelo al entrenar con un documento, y cuánto de ese cambio pueden detectar los ataques de inferencia de pertenencia.

La pregunta

La inferencia de pertenencia pregunta si un documento dado formó parte de los datos de entrenamiento de un modelo. Esto importa para los derechos de autor y la procedencia de los datos, para la transparencia de los datos de entrenamiento y para la contaminación de benchmarks. La mayoría de las evaluaciones de ataques de inferencia de pertenencia sobre modelos de lenguaje recopilan miembros y no miembros a posteriori, por ejemplo documentos publicados antes y después de una fecha de corte de entrenamiento. Los dos grupos difieren entonces en tema, estilo y fecha además de en la pertenencia. Un ataque puede tener éxito detectando este cambio de distribución. Cuando un ataque falla, eso no demuestra que el entrenamiento no haya dejado rastro.

Para separar estas explicaciones, necesitamos un entorno en el que la pertenencia sea conocida, se asigne de forma independiente de toda propiedad del documento y deje el resto del entrenamiento sin cambios. El artículo plantea: cuando un documento real se incluye una vez en una receta de entrenamiento establecida, ¿cuánto cambia el modelo, y cuánto de ese cambio pueden aprovechar los ataques existentes?

El diseño

                    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 repetición de una etapa de entrenamiento oficial. OLMo 2 1B se entrena en dos etapas: unos 4 billones de tokens de preentrenamiento, y luego una etapa de entrenamiento intermedio de 50 mil millones de tokens. AI2 publicó el checkpoint en el que comienza el entrenamiento intermedio, incluido su estado del optimizador. AI2 también publicó los datos del entrenamiento intermedio, su configuración y las ejecuciones completas del entrenamiento intermedio para tres semillas de orden de datos. Partimos del checkpoint publicado y ejecutamos la etapa de nuevo con la configuración de AI2 y el entrenador OLMo sin modificaciones. Reconstruimos exactamente el orden de datos de la semilla 42.
• Pares emparejados y una tirada de moneda. Construimos 7.814 pares de documentos. Los dos documentos de un par provienen de la misma fuente y cohorte y están emparejados en longitud y en dificultad: bits por byte bajo el checkpoint inicial, medidos antes del entrenamiento. Para cada par, una tirada de moneda eligió un documento para la inserción; el otro se reservó. Copias de los documentos insertados reemplazan fragmentos completos de entrenamiento en posiciones planificadas distribuidas a lo largo de la etapa: 12.414 copias reemplazan 24.081 de las 12.212.224 posiciones de fragmentos (0,197%), y el texto insertado suma 67,3 millones de tokens (0,135% del flujo). Todas las demás posiciones contienen el fragmento de AI2.
• La ejecución publicada como control. La ejecución publicada de la semilla 42 de AI2 vio los mismos fragmentos en el mismo orden en todas las demás posiciones, por lo que sirve como contrafactual. El artículo la llama el gemelo. Las dos ejecuciones difieren en ajustes de ejecución (número de GPUs, kernels de atención y de sharding), por lo que no son idénticas bit a bit. La ejecución de la semilla 42069 de AI2 (la misma receta en un orden de datos diferente) sirve como referencia reentrenada. El checkpoint publicado en el que comienza la etapa sirve como referencia más débil.
• Documentos verificados como no vistos. La mayoría de los documentos son artículos de arXiv, páginas de Wikipedia y publicaciones de StackExchange que aparecieron después del corte de datos del modelo, o páginas de Wikipedia ausentes de la instantánea en los datos de entrenamiento. Trece cohortes varían la fuente, la longitud, el número de copias (hasta 32), el momento de las copias y la forma del texto (derivado de LaTeX frente a derivado de PDF). Las versiones editadas de páginas de Wikipedia que están en los datos de entrenamiento miden la reexposición, y los textos sintéticos abarcan un rango de sorpresa. Los candidatos de las cohortes destinadas a ser no vistas se buscaron exhaustivamente contra los 13 componentes de los corpus publicados de preentrenamiento y entrenamiento intermedio (8,3 TiB). La búsqueda detecta toda superposición literal de 40 palabras o más, y se excluyeron los candidatos con tales superposiciones. La última parte de la búsqueda terminó después de que se hubieran formado los pares. Encontró 156 candidatos más con tales superposiciones; se eliminaron y sus pares se rellenaron antes del entrenamiento.
• Post-entrenamiento. Tanto el modelo tratado como el control se sometieron a la receta de post-entrenamiento de AI2 para OLMo 2 1B: ajuste fino supervisado (SFT), optimización directa de preferencias (DPO) y dos etapas de aprendizaje por refuerzo con recompensas verificables (RLVR). Ambos modelos recibieron los mismos datos, orden y semillas.
• Mediciones. Los 22.681 documentos de evaluación (ambos documentos de cada par y 7.053 documentos que nunca se insertaron) se puntúan token por token bajo:
  • el modelo tratado, el control, la referencia reentrenada y el checkpoint inicial;
  • los 24 checkpoints guardados cada 1.000 actualizaciones por cada una de las tres ejecuciones de entrenamiento intermedio;
  • los modelos de cada etapa de post-entrenamiento de ambas ramas.

Hallazgos principales

Salvo que se indique lo contrario, las cifras a continuación corresponden a los 4.832 pares de documentos verificados como no vistos e insertados una vez, medidos frente al control. El artículo proporciona intervalos de confianza y más advertencias.