Resumen Técnico: El Preentrenamiento Modular Permite el Control de Acceso

Planteamiento del Problema

Los desarrolladores de IA enfrentan un "dilema de doble uso": las capacidades que permiten aplicaciones beneficiosas (por ejemplo, curar enfermedades, fortalecer la ciberseguridad) a menudo se solapan con aquellas que permiten daños (por ejemplo, sintetizar armas biológicas, ejecutar ciberataques). Los mecanismos actuales de control de acceso son insuficientes. El "estándar de oro" para el control de acceso sería servir modelos distintos, entrenados por separado con diferentes perfiles de capacidad, a diferentes usuarios (por ejemplo, un modelo con conocimientos de virología para investigadores pero sin ellos para el público en general). Sin embargo, entrenar y desplegar $N$ modelos separados para $N$ perfiles de capacidad es prohibitivamente costoso en escalas de frontera.

Las alternativas existentes no logran cumplir con los requisitos de robustez:

• El entrenamiento de rechazo, el filtrado de salidas y los clasificadores restringen las salidas pero dejan la capacidad subyacente intacta, lo que los hace vulnerables a la elicitación adversaria.
• El desaprendizaje post-hoc (por ejemplo, MaxEnt) intenta eliminar capacidades de un modelo ya entrenado, pero es frágil; las capacidades a menudo pueden recuperarse rápidamente mediante el ajuste fino (fine-tuning) con unos pocos ejemplos.
• La ramificación de modelos (entrenar un modelo base y realizar ajustes finos de variantes separadas) reduce los costos de preentrenamiento, pero sigue requiriendo múltiples despliegues, lo cual es ineficiente para patrones de uso dispersos.

Metodología: Módulos Auxiliares de Enrutamiento de Gradiente (GRAM)

El artículo propone GRAM, un método de preentrenamiento que permite a un solo modelo soportar múltiples perfiles de capacidad mediante la inducción de especialización a través de una arquitectura modular y enrutamiento de gradientes.

Arquitectura

GRAM aumenta un Transformer denso estándar añadiendo módulos MLP auxiliares más pequeños a cada capa junto al MLP "núcleo" principal.

• Módulo Núcleo (Core): Un MLP grande compartido por todos los datos, que captura el conocimiento general.
• Módulos Auxiliares: MLPs más pequeños, cada uno correspondiente a una capacidad auxiliar específica (por ejemplo, virología, ciberseguridad).
• Recuento de Parámetros: El recuento total de parámetros es comparable al de un modelo denso estándar, donde el módulo núcleo iguala el tamaño de la línea base y los módulos auxiliares comprenden una pequeña fracción de los parámetros.

Dinámica de Entrenamiento (Enrutamiento de Gradiente)

A diferencia de la Mezcla de Expertos (MoE) que utiliza enrutadores aprendidos para el enrutamiento por token, GRAM utiliza un enrutamiento de gradiente dependiente de los datos basado en las etiquetas del conjunto de datos:

1. Paso hacia adelante (Forward Pass):
   • Datos del Núcleo: El MLP del núcleo siempre está activo. Con una probabilidad $p_{cr}$, también se activa un módulo auxiliar aleatorio para asegurar que el núcleo sea robusto ante la presencia del módulo.
   • Datos Auxiliares: Tanto el MLP del núcleo como el módulo auxiliar específico correspondiente a esos datos están activos.
2. Paso hacia atrás (Backward Pass - Actualización de Gradientes):
   • Parámetros del Núcleo: Se actualizan con una probabilidad $p_{as}$ cuando se procesan datos auxiliares. Esto controla el grado de interferencia entre capacidades.
   • Parámetros Auxiliares: El módulo auxiliar específico para el lote (batch) actual siempre se actualiza. Otros módulos auxiliares no reciben actualizaciones.
   • Optimizadores Separados: Cada subconjunto de parámetros (el núcleo y cada módulo auxiliar) se optimiza utilizando un optimizador AdamW independiente.

Inferencia y Control de Acceso

En el momento de la inferencia, se sirve un perfil de capacidad específico activando el MLP del núcleo y el subconjunto de módulos auxiliares correspondientes a las capacidades autorizadas. La ablación (desactivación) de un módulo auxiliar elimina su capacidad asociada de la red. El artículo postula que esto aproxima un modelo entrenado con datos filtrados, ya que el módulo ablacionado nunca aprendió la capacidad específica de forma aislada del núcleo.

Contribuciones Clave y Experimentos

1. Aproximación del Filtrado de Datos

Los autores demuestran que un único modelo entrenado con GRAM puede aproximar el rendimiento de múltiples modelos entrenados mediante el filtrado de datos (donde cada modelo se entrena en un subconjunto de datos).

• Evaluación Sintética (Historias Simples): En un conjunto de datos de historias infantiles particionadas en temas de núcleo y auxiliares, GRAM igualó a los modelos de datos filtrados tanto en la retención de las capacidades deseadas como en la eliminación de las no deseadas.
• Evaluación de Doble Uso en el Mundo Real: Entrenado en un modelo de 800 millones de parámetros con datos del núcleo (texto web, código, artículos generales) y cuatro dominios auxiliares (virología, ciberseguridad, física nuclear, código Lisp especializado), GRAM logró:
  • Retención: El rendimiento en las capacidades retenidas casi igualó al de la línea base de todos los datos.
  • Eliminación: El rendimiento en las capacidades eliminadas cayó significativamente, igualando el rendimiento de "olvido" de los modelos de datos filtrados.
  • Robustez: La eliminación de capacidades por parte de GRAM fue más robusta frente al ajuste fino malicioso (elicitación de olvido) que los métodos de desaprendizaje post-hoc (MaxEnt).

2. Escalabilidad y Composibilidad

• Leyes de Escala: Un análisis de escala óptimo para Chinchilla de 50 millones a 5 mil millones de parámetros mostró que GRAM sigue de cerca el perfil de capacidad del filtrado de datos. A medida que el tamaño del modelo aumentó, la brecha entre los modelos de datos filtrados y los de datos completos se amplió para las capacidades eliminadas, y GRAM mantuvo esta tendencia.
• Composibilidad: GRAM soporta subconjectos arbitrarios de capacidades desde una única ejecución de entrenamiento. A diferencia de la ramificación basada en LoRA (FT-LoRA), que se degrada cuando se suman múltiples adaptadores, GRAM mantuvo el rendimiento cuando múltiples módulos auxiliares estaban activos simultáneamente, permitiendo $2^{N-1}$ perfiles de capacidad de un solo modelo.
• Etiquetado Parcial: En un escenario donde solo el 50% de los datos estaban etiquetados, GRAM superó significativamente tanto al filtrado de datos como a FT-LoRA en la eliminación de capacidades. Los autores atribuyen esto a un "efecto de absorción", donde las actualizaciones de los datos no etiquetados se concentran en los módulos especializados en lugar de en el núcleo compartido.

3. Eficiencia de Costos

• Costo de Entrenamiento: GRAM requiere un único proceso de entrenamiento, reduciendo el costo computacional de entrenamiento en un factor de 5 en comparación con el entrenamiento de 5 modelos de datos filtrados separados.
• Costo de Despliegue: A diferencia de la ramificación de modelos o el ajuste fino completo, GRAM no requiere el despliegue de múltiples copias del modelo; sirve un único modelo con activación dinámica de módulos.

Significado y Reivindicaciones

El artículo sostiene que GRAM ofrece una vía práctica para un control de acceso robusto para las capacidades de IA de doble uso sin los costos prohibitivos de entrenar múltiples modelos.

• Seguridad: Al eliminar físicamente los parámetros responsables de una capacidad (mediante ablación) en lugar de simplemente suprimir las salidas, GRAM proporciona garantías más fuertes contra el uso indebido y el jailbreaking en comparación con la dirección conductual o el desaprendizaje post-hoc.
• Viabilidad Económica: El método desacopla el costo de entrenamiento del número de perfiles de capacidad soportados, haciendo que el control de acceso diferenciado sea factible en escalas de frontera.
• Alineación con el Filtrado de Datos: Los resultados establecen que el filtrado de datos en sí mismo escala favorablemente (los modelos más grandes son mejores para olvidar datos específicos mientras retienen el conocimiento general) y que GRAM sigue con éxito este comportamiento.

Los autores mantienen la modestia respecto a las limitaciones, señalando que:

1. Entrelazamiento: Si las capacidades están altamente entrelazadas, una separación limpia puede ser imposible independientemente del método.
2. Escala de Frontera: Aunque las tendencias se mantienen hasta los 5 mil millones de parámetros, el rendimiento en escalas de frontera (billones de parámetros) no ha sido verificado.
3. Implementación: La integración de GRAM en los flujos de producción introduce una complejidad que los desarrolladores podrían preferir evitar.
4. Post-entrenamiento: Se desconoce la interacción entre la separación de preentrenamiento de GRAM y el posterior ajuste de instrucciones o RLHF.

En conclusión, el artículo posiciona a GRAM no como una solución perfecta, sino como un paso significativo hacia el despliegue seguro y flexible de sistemas de IA altamente capaces, permitiendo a los proveedores servir perfiles de capacidad autorizados desde un único modelo entrenado de manera eficiente.