Attention, May I Have Your Decision? Localizing Generative Choices in Diffusion Models

Resumen Técnico: Localización de Decisiones Generativas en Modelos de Difusión

1. El Problema

Los modelos de difusión de texto a imagen (T2I) poseen capacidades generativas notables, pero sus mecanismos internos siguen siendo una "caja negra". Este problema se agrava cuando los prompts (instrucciones de texto) no son totalmente descriptivos (ej. "una foto de flores" o "una foto de una persona"). En estos casos, el modelo debe tomar decisiones implícitas para rellenar los detalles faltantes (como el color, la edad, el género o la raza) basándose en patrones aprendidos de los datos de entrenamiento.

Estas decisiones implícitas a menudo perpetúan sesgos sociales (ej. asumir que un "doctor" es hombre o que un "presidente de EE. UU." es Donald Trump). La literatura previa se ha centrado en localizar conocimientos mediante la inyección de prompts (modificar el texto de entrada para ver su impacto), pero los autores argumentan que este enfoque no revela los mecanismos internos que el modelo utiliza cuando no hay una instrucción explícita. El desafío es identificar dónde y cómo el modelo decide instantiar conceptos generales en imágenes específicas sin alterar el prompt original.

2. Metodología: ICM (Implicit Choice-Modification)

Los autores proponen una nueva técnica llamada ICM, que se basa en la localización precisa de las capas del modelo responsables de estas decisiones implícitas, evitando la confusión de la ingeniería de prompts. El proceso consta de tres etapas principales:

• A. Extracción de Activaciones y Etiquetado Pseudo:

  • Se generan imágenes utilizando prompts generales (ej. "una foto de una persona").
  • Se extraen las activaciones intermedias de todas las capas del modelo durante el proceso de generación.
  • En lugar de forzar una etiqueta en el prompt, se utiliza un clasificador externo (como CLIP o FairFace) para etiquetar a posteriori los atributos de las imágenes generadas (ej. determinar si la imagen resultante es de un hombre o una mujer). Esto crea un conjunto de datos de activaciones con etiquetas reales de la decisión tomada por el modelo.

• B. Localización mediante Sondas Lineales (Linear Probing):

  • Se entrenan sondas lineales (regresión logística) sobre las activaciones de cada capa y paso de tiempo para predecir el atributo (ej. género).
  • El objetivo es medir la separabilidad de atributos: ¿En qué capa las activaciones se vuelven linealmente separables para distinguir entre diferentes instantiaciones del mismo concepto?
  • Hallazgo Clave: Las activaciones justo después de las capas de Auto-Atención (Self-Attention) muestran una separabilidad significativamente mayor que las de las capas de Atención Cruzada (Cross-Attention). Esto sugiere que, mientras la atención cruzada maneja el texto explícito, la auto-atención es la responsable de resolver la ambigüedad y tomar decisiones basadas en el ruido inicial y el contexto espacial.

• C. Control de Generación (Intervención):

  • Una vez identificadas las capas óptimas (un subconjunto pequeño de capas de auto-atención), se aplican intervenciones dirigidas:
    1. Steering de Activación: Se calcula un vector de dirección basado en los pesos del clasificador y se añade a las activaciones de las capas seleccionadas durante la inferencia ($H' = H + \alpha \cdot s$).
    2. Fine-tuning (Ajuste Fino): Se ajusta el modelo utilizando técnicas como LoRA, pero restringiendo el entrenamiento únicamente a las capas seleccionadas.

3. Contribuciones Clave

1. Localización de Decisiones Implícitas: Demuestran que es posible localizar las capas críticas para las decisiones no especificadas en el texto mediante probing, sin necesidad de inyectar atributos en el prompt.
2. Rol de la Auto-Atención: Identifican que las capas de Auto-Atención son el locus principal para la resolución de conceptos ambiguos y la toma de decisiones implícitas, superando a las capas de Atención Cruzada en este contexto específico.
3. Método ICM: Presentan un método de intervención precisa que modifica solo un subconjunto pequeño de capas, logrando un control granular sobre los atributos generados.
4. Validación en Múltiples Arquitecturas: El método se valida en modelos basados en UNet (Stable Diffusion v1.5, SDXL) y en arquitecturas Transformer (SANA), demostrando su generalidad.

4. Resultados Experimentales

Los experimentos se centraron en la eliminación de sesgos de género, edad y raza:

• Rendimiento de Desviación (Debiasing):

  • La variante ICM (Steering) logró una reducción superior en la Discrepancia de Equidad (Fairness Discrepancy - FD) en comparación con métodos state-of-the-art como Finetuning completo, Latent Editing o DIFFLENS.
  • Por ejemplo, para el sesgo de género, ICM redujo el FD a 0.087, superando a otros métodos que a menudo degradaban la calidad de la imagen al intentar corregir el sesgo.

• Calidad de Imagen y Alineación:

  • A diferencia de métodos que modifican todas las capas o usan fine-tuning global (que suelen aumentar el FID y degradar la imagen), ICM mantuvo una alta calidad visual (FID bajo) y una fuerte alineación con el prompt original (CLIP-T alto).
  • La intervención en capas de Auto-Atención fue significativamente más efectiva para el desviado que la intervención en Atención Cruzada, confirmando la hipótesis de que las decisiones implícitas se originan en la auto-atención.

• Comparación de Localización:

  • Se demostró que las sondas entrenadas con prompts generales capturan mejor los mecanismos intrínsecos del modelo que las entrenadas con prompts específicos (inyección), evitando un desplazamiento de distribución que podría llevar a intervenciones subóptimas.

• Escalabilidad:

  • En modelos grandes como SDXL y SANA, seleccionar solo las mejores capas (ej. las 20 mejores de 70) permitió un control preciso (ej. cambiar la imagen de un presidente de EE. UU. a Barack Obama) sin destruir la coherencia de la imagen, algo que fallaba al intervenir en todas las capas.

5. Significado e Impacto

Este trabajo es fundamental porque cambia el paradigma de cómo entendemos y controlamos los modelos de difusión:

• Interpretabilidad: Revela que la toma de decisiones "creativa" o "implícita" en los modelos de difusión no es difusa, sino que está computacionalmente localizada en capas específicas de auto-atención.
• Eficiencia y Precisión: Proporciona una herramienta para la mitigación de sesgos que es mucho más eficiente (requiere menos recursos de entrenamiento) y precisa que el fine-tuning completo o la manipulación global de activaciones.
• Aplicabilidad Práctica: Ofrece un método robusto para generar imágenes más justas y diversas sin sacrificar la calidad artística o la fidelidad al texto, lo cual es crucial para la adopción ética de la IA generativa en entornos reales.

En resumen, el artículo demuestra que para controlar las decisiones ocultas de una IA generativa, no debemos forzar la entrada con más texto, sino intervenir quirúrgicamente en las capas internas donde el modelo "toma la decisión" de llenar los vacíos de información.