Toward Localizing and Repairing Bias in Transformer Attention Heads

Resumen Técnico: Hacia la localización y reparación del sesgo en las cabezas de atención de los Transformers

Declaración del Problema
Los modelos de lenguaje Transformer se despliegan cada vez más como componentes de software, sin embargo, sus salidas sesgadas siguen siendo difíciles de localizar y reparar. Las intervenciones de equidad existentes suelen operar a nivel de entrada-salida o requieren el reentrenamiento de todo el modelo, fallando al no identificar las estructuras internas específicas responsables de los fallos de equidad durante la inferencia. Aunque investigaciones recientes indican que el comportamiento relacionado con el sesgo en los modelos de lenguaje con máscara puede concentrarse en un pequeño conjunto de cabezas de atención, los métodos actuales para modificar estas cabezas suelen depender de la "anulación de la cabeza completa" (whole-head zeroing). Este enfoque asume que una cabeza seleccionada codifica principalmente el sesgo, lo que podría eliminar información lingüística útil y degradar el rendimiento del modelo. Existe la necesidad de un método de depuración de caja blanca que pueda localizar componentes susceptibles al sesgo y aplicar una modificación dirigida y acotada sin reentrenamiento.

Metodología: ROBIN
El artículo presenta ROBIN (Repair Of Bias via INtervention), un método de depuración de equidad a nivel de cabeza y de caja blanca. ROBIN opera en dos etapas utilizando un modelo transformer entrenado $f_\theta$ y un conjunto de pares de sondas de equidad (pares de texto que difieren solo en un atributo protegido, como los pronombres de género).

1. Etapa 1: Clasificación de Cabezas (Localización)
   ROBIN clasifica las cabezas de atención basándose en su sensibilidad a la diferencia entre los pares de sondas ($x_a$ y $x_b$).

   • Puntuación: Para cada par, el método calcula el gradiente de la pérdida con respecto a los pesos de atención de cada cabeza.
   • Agregación: Calcula dos métricas por cabeza: la magnitud media del gradiente ($G_h$) y la magnitud media del cuadrado del gradiente ($C_h$). $G_h$ captura la sensibilidad constante a través de los pares, mientras que $C_h$ captura cabezas con respuestas agudas y concentradas.
   • Selección: Estas puntuaciones se estandarizan y se suman para producir un ranking combinado ($F_h$). Se seleccionan las $k$ cabezas superiores para la intervención.

2. Etapa 2: Intervención en el Tiempo de Inferencia (Reparación)
   En lugar de anular las cabezas seleccionadas, ROBIN elimina únicamente el componente de sesgo de sus salidas.

   • Estimación del Subespacio: Para cada cabeza seleccionada, el método estima un subespacio de sesgo analizando los vectores de diferencia entre las salidas de la cabeza en los pares de sondas. Calcula los $r$ vectores singulares derechos superiores de estos vectores de diferencia para definir la dirección del sesgo.
   • Proyección: Durante la inferencia, la salida de una cabeza seleccionada ($x_h$) se modifica proyectando fuera el componente que reside dentro del subespacio de sesgo estimado ($V_h$):
     $$x'_h = x_h - V_h^\top V_h x_h$$
   • Esta operación conserva los componentes ortogonales de la salida de la cabeza, preservando el comportamiento lingüístico general mientras se elimina el sesgo estimado.

Contribuciones Clave

• Encuadre: Los autores enmarcan la mitigación de la equidad a nivel de cabeza como una tarea de depuración de software de caja blanca que separa explícitamente la localización de los componentes candidatos de la modificación de dichos componentes.
• Algoritmo: La propuesta de ROBIN, que identifica cabezas de atención sensibles al sesgo y elimina solo el subespacio sesgado de sus salidas en el tiempo de inferencia, evitando la naturaleza destructiva de la ablación de la cabeza completa.
• Evidencia Empírica: Un estudio piloto en cuatro modelos transformer preentrenados (BERT, DistilBERT, GPT-2, DistilGPT-2) que demuestra que ROBIN reduce la brecha de WinoBias mientras limita el aumento en la pérdida de modelado de lenguaje.

Resultos
Evaluado en un estudio piloto con un presupuesto de $k=20$ cabezas modificadas:

• Mejora de la Equidad: ROBIN redujo la brecha de WinoBias en los cuatro modelos. La reducción absoluta más significativa se observó en BERT, donde la brecha cayó de 45.97 a 19.14.
• Preservación del Rendimiento: A diferencia de la anulación de la cabeza completa (métodos Gradient y SquaredGrad), que causaron aumentos sustanciales en la pérdida de modelado de lenguaje (por ejemplo, hasta 8 veces de aumento en GPT-2), ROBIN preservó la calidad del modelo. En BERT, ROBIN logró la brecha más baja con una pérdida de modelado de lenguaje cercana a la línea base. En GPT-2, el aumento de la pérdida fue mínimo en comparación con la línea base.
• Costo de Inferencia: La intervención añade una multiplicación de matrices por cada cabeza elegida. Con $k=20$, la sobrecarga de latencia osciló aproximadamente entre un 7% para modelos de tamaño base y un 11% para modelos destilados.
• StereoSet: El puntaje de estereotipo de StereoSet mostró movimientos inconsistentes a través de modelos y categorías, lo que llevó a los autores a tratar la brecha de WinoBias como la métrica principal de equidad.

Significancia y Reivindicaciones
El artículo sostiene que la reparación de sesgo a nivel de cabeza debe considerar no solo qué cabezas se seleccionan, sino cómo se modifican. Los resultados sugieren que la eliminación de subespacios es un operador superior a la anulación de la cabeza completa porque reconoce que las cabezas de atención a menudo apoyan múltiples funciones (tanto de sesgo como lingüísticas generales).

Los autores posicionan a ROBIN como una herramienta para la localización de candidatos e intervención medible, no como una explicación causal completa o una garantía de equidad en todos los contextos de despliegue. Establecen explícitamente que el método proporciona una reparación acotada sin reentrenamiento, permitiendo a los profesionales seleccionar un presupuesto de intercambio ($k$) basado en la degradación medida del rendimiento. El estudio reconoce limitaciones, incluyendo el uso de datos superpuestos para el ranking y la evaluación, así como el enfoque en benchmarks específicos (WinoBias) en lugar de la equidad en tareas de aplicación general. Se identifica que es necesario un trabajo futuro para evaluar modelos más grandes, variantes ajustadas por instrucciones (instruction-tuned) y distribuciones de datos disjuntas.