SHIFT: Steering Hidden Intermediates in Flow Transformers

Resumen Técnico: SHIFT - Dirección de Intermediarios Ocultos en Flow Transformers

1. El Problema

Los modelos de difusión, especialmente los basados en Transformers de Difusión (DiT) como FLUX y SD3.5, han establecido nuevos estándares en la generación de imágenes de alta fidelidad. Sin embargo, a medida que estas capacidades escalan, aumentan los riesgos de generar contenido dañino, protegido por derechos de autor o no deseado.

El desafío principal abordado es la Eliminación de Conceptos (Concept Erasure - CE): la tarea de suprimir semánticamente conceptos específicos dentro del manifold generativo del modelo sin reentrenarlo.

• Limitaciones de los métodos actuales: Las estrategias existentes (como ESD, Concept Ablation, EAP) dependen de la optimización basada en gradientes y el reentrenamiento. Estos métodos son computacionalmente prohibitivos para modelos DiT masivos (ej. FLUX con 12B parámetros).
• Incompatibilidad con modelos destilados: Muchos métodos requieren guía negativa durante la optimización, lo cual es problemático para modelos destilados (como Flux.1[schnell]) que están diseñados para operar con una escala de guía fija.
• Arquitectura Unificada: A diferencia de las arquitecturas U-Net anteriores que separaban la atención cruzada (texto) y auto-atención (imagen), los DiT modernos utilizan bloques de atención unificados (MM-DiT) donde texto e imagen comparten un espacio latente, haciendo que las intervenciones tradicionales (como manipular solo la atención cruzada) sean insuficientes o inestables.

2. Metodología: SHIFT

SHIFT es un marco de trabajo ligero y sin reentrenamiento que aplica dirección de activación (activation steering) directamente en los intermediarios ocultos del modelo durante la inferencia.

• Mecanismo Central: En lugar de modificar los pesos del modelo, SHIFT inyecta vectores de dirección específicos en las activaciones de los tokens de texto dentro de los bloques del Transformer de Difusión.
• Puntos de Intervención:
  1. Embedding agrupado del codificador de texto: Se manipula la representación global del prompt (ej. el token CLIP en Flux).
  2. Tokens de texto en los bloques DiT: Se inyecta el vector de dirección en las salidas de los bloques de atención después de la capa de atención compartida.
• Construcción del Vector de Dirección:
  • Se crea un pequeño conjunto de datos con pares de prompts contrastantes (con y sin el concepto objetivo).
  • Se recogen las activaciones correspondientes.
  • Se calcula el vector de dirección utilizando dos enfoques:
    1. Diferencia de medias: La diferencia promedio entre las activaciones de los prompts positivos y negativos.
    2. SVM (Máquinas de Vectores de Soporte): Se entrena un SVM lineal sobre las activaciones agrupadas para encontrar el hiperplano separador; el vector normal de este hiperplano se utiliza como vector de dirección.
• Aplicación en Inferencia:
  • La activación $a$ se modifica como: $\tilde{a} = a + \alpha \cdot v$, donde $v$ es el vector de dirección y $\alpha$ es la fuerza de dirección.
  • Regularización Dinámica: Para evitar la "sobre-dirección" (degradar conceptos no objetivo), la fuerza $\alpha$ se ajusta dinámicamente en función de la probabilidad predicha por un clasificador ligero (SVM) de que el concepto objetivo aún esté presente en la activación actual.
• Invarianza Temporal: Un hallazgo clave es que un único vector de dirección (independiente del tiempo) puede aplicarse eficazmente a través de todos los pasos de difusión, simplificando enormemente el proceso.

3. Contribuciones Clave

1. Marco Unificado de Dirección: SHIFT es el primer marco integral capaz de dirigir tanto modelos DiT originales como destilados (ej. Flux.1[dev] y Flux.1[schnell]) mediante cambios en las activaciones latentes, sin reentrenamiento.
2. Análisis de Dinámicas Espaciales y Temporales: El estudio revela una consistencia temporal notable en los espacios de activación de los DiT. Se demuestra que los conceptos semánticos se codifican como direcciones globales estables, permitiendo el uso de un solo vector para todo el proceso de generación.
3. Transferencia de Vectores entre Destilación: Se valida que los vectores de dirección aprendidos en un modelo destilado (Flux.1[schnell]) pueden transferirse y ser efectivos en el modelo no destilado (Flux.1[dev]), lo cual es crucial para entornos con recursos limitados.
4. Versatilidad de Tareas: Además de la eliminación de conceptos, SHIFT permite:
   • Estilización: Cambiar el dominio de generación (ej. a estilo cyberpunk o boceto).
   • Adición de Objetos: Insertar objetos específicos en la imagen.
   • Cambio de Concepto: Sustituir un objeto por otro (ej. gato $\to$ perro).

4. Resultados

Los experimentos se realizaron principalmente en los modelos Flux.1[dev] y Flux.1[schnell].

• Eliminación de Conceptos Abstractos (Seguridad):
  • En el benchmark I2P (enfocado en contenido inapropiado/nudez), SHIFT supera significativamente a las líneas base (ESD, EAP, CA, UCE).
  • Logra una supresión de contenido explícito 3 a 4 veces mayor que los métodos anteriores.
  • Mantiene una alta alineación con el prompt (CLIP score estable) y una degradación mínima en la calidad de imagen (FID), demostrando un equilibrio superior entre seguridad y fidelidad.
• Eliminación de Conceptos Concretos:
  • En tareas de eliminar objetos específicos (ej. "Snoopy") preservando conceptos relacionados (ej. "Mickey", "Perro"), SHIFT muestra una supresión más fuerte del objetivo con menor impacto en los conceptos preservados en comparación con CA y UCE.
• Estilización y Adición:
  • El método logra eliminar estilos de artistas famosos (ej. Van Gogh) preservando otros estilos, y añade objetos pequeños (gafas, sombreros) o cambia atributos (edad, género) con éxito.
• Eficiencia:
  • No requiere reentrenamiento ni optimización de gradientes durante la inferencia.
  • El proceso de cálculo de vectores es rápido y requiere solo un pequeño conjunto de datos de prompts (20-135 pares).

5. Significado e Impacto

SHIFT representa un cambio de paradigma en el control de modelos de difusión de última generación:

• Escalabilidad: Ofrece una solución viable para la seguridad y el control en modelos masivos (12B+ parámetros) donde el reentrenamiento es inviable.
• Independencia de la Guía: Al no depender de la guía negativa durante la optimización, es compatible con modelos destilados de alta velocidad, que son la tendencia actual en la industria.
• Interpretabilidad: Refuerza la idea de que los conceptos semánticos en los DiT están organizados en direcciones lineales estables en el espacio latente, facilitando la manipulación precisa sin alterar la estructura global de la imagen.
• Accesibilidad: Al ser un método de inferencia, democratiza el control de modelos avanzados, permitiendo a los usuarios eliminar conceptos dañinos o modificar estilos sin necesidad de recursos computacionales masivos para el fine-tuning.

En conclusión, SHIFT demuestra que la manipulación directa de las activaciones internas en los Transformers de Difusión es una estrategia robusta, eficiente y flexible para el control semántico, superando las limitaciones de los métodos basados en optimización en la era de los modelos de difusión modernos.