Resumen Técnico: Autoencoders Esparsos Temporales Residualizados para la Interpretación de Modelos de Difusión

Enunciado del Problema

Los modelos de difusión de texto a imagen generan imágenes mediante un proceso iterativo de eliminación de ruido, lo que provoca que las capas neuronales internas produzcan trayectorias de activaciones en lugar de representaciones estáticas únicas. Aunque los Autoencoders Esparsos (SAE) se han aplicado recientemente para descomponer las activaciones de difusión en direcciones de características interpretables, los enfoques existentes suelen analizar las activaciones en instantes de tiempo aislados, condicionarse al tiempo o aplicar SAE a arquitecturas como los Transformadores de Difusión.

Una limitación crítica de estos métodos existentes es que no tienen en cuenta explícitamente la fuerte correlación temporal inherente a las activaciones de difusión. Las activaciones en instantes de tiempo adyacentes están altamente correlacionadas, siendo gran parte de la señal en un paso predecible linealmente a partir del paso anterior. Entrenar un SAE directamente sobre activaciones crudas concatenadas temporalmente corre el riesgo de desperdiciar la capacidad del modelo en aprender estas dinámicas redundantes y predecibles linealmente, en lugar de capturar la estructura subyacente esparsa de las características.

Metodología

Los autores proponen SAE Temporales Residualizados, un marco diseñado para aprender características esparsas a partir de trayectorias de activaciones separando explícitamente las dinámicas predecibles linealmente de la estructura residual.

1. Recolección y Preprocesamiento de Datos

• Fuente: Las activaciones se recopilan de la rama condicional del U-Net en Stable Diffusion 1.5 durante la guía sin clasificador. Específicamente, se extrae la contribución residual de la capa de atención cruzada del bloque medio.
• Construcción de Trayectorias: Para una imagen generada, las activaciones se submuestrean cada $\ell$ pasos de eliminación de ruido (por ejemplo, $\ell=10$) para formar una secuencia de tensores de activación $\{a(\tau_i)\}$.

2. Residualización Lineal

Para eliminar la redundancia temporal, el método ajusta predictores lineales entre bloques de activación normalizados consecutivos:

• Normalización: Las activaciones se normalizan por instante de tiempo utilizando estadísticas del conjunto de entrenamiento.
• Regresión Ridge: Para cada transición adyacente $i-1 \to i$, se ajusta un regresor ridge $(W_i, b_i)$ para predecir la activación actual $\bar{a}_{q,i}$ a partir de la anterior $\bar{a}_{q,i-1}$.
• Cálculo del Residual: El residual $r_{q,i} = \bar{a}_{q,i} - (W_i \bar{a}_{q,i-1} + b_i)$ captura el componente de la activación no predecible linealmente a partir del paso anterior.

3. Entrenamiento de SAE sobre Trayectorias Residualizadas

La entrada al SAE es un vector concatenado $z_q$ que consiste en la activación inicial y los residuos subsiguientes:
$$z_q = [\bar{a}_{q,0}, r_{q,1}, \dots, r_{q,T-1}]$$
Se entrena un SAE BatchTopK sobre esta representación residualizada. El objetivo es reconstruir la entrada normalizada mientras se impone esparsidad (reteniendo solo los $K$ latentes activos superiores por lote). Esto fomenta que los latentes esparsos capturen la estructura que existe más allá de la deriva temporal predecible linealmente.

4. Interpretación en el Espacio de Activaciones

Un componente clave del método es mapear las características aprendidas de vuelta al espacio de activación original para analizarlas como trayectorias:

1. Direcciones del Decodificador: Los vectores del decodificador del SAE están inicialmente en el espacio residualizado normalizado.
2. Desnormalización: Se aplican escalas para volver al espacio residual normalizado.
3. Integración de Dinámicas Lineales: Las direcciones residuales se mapean de vuelta a través de los predictores lineales ajustados ($W_i$) para reconstruir la dirección completa de la trayectoria en el espacio de activación normalizado.
4. Mapeo Final: El resultado se desnormaliza para obtener la trayectoria del decodificador $\{\phi_{k,i}\}$ en las unidades originales de activación. Esto permite analizar cada latente como una trayectoria de características que evoluciona a lo largo del tiempo de eliminación de ruido.

Contribuciones Clave

1. SAE Temporales Residualizados: La introducción de un método que separa las trayectorias de activación de difusión en componentes predecibles linealmente y residuos, entrenando SAE específicamente sobre estos últimos para evitar modelar dinámicas temporales redundantes.
2. Interpretación del Decodificador en el Espacio de Activaciones: Una derivación que mapea las direcciones del decodificador residualizado de vuelta al espacio de activación original, permitiendo el análisis de cada latente como una trayectoria de característica espacio-temporal.
3. Caracterización Empírica: Una evaluación exhaustiva en Stable Diffusion 1.5 que demuestra que estas representaciones capturan direcciones estructuradas temporalmente y localizadas espacialmente con efectos generativos semánticamente significativos.

Resultados y Evaluación

Rendimiento de Reconstrucción

• Superioridad de la Residualización: Bajo un presupuesto fijo de esparsidad a nivel de trayectoria, las variantes residualizadas superan consistentemente a sus contrapartes no residualizadas (tanto concatenadas como por instante de tiempo) en el Error Cuadrático Medio (MSE) y la Varianza Explicada (EV).
• Eficiencia: El modelo Concatenado Residualizado logra un rendimiento de reconstrucción competitivo mientras utiliza un único código esparsa para toda la trayectoria, whereas las líneas base por instante de tiempo requieren códigos distribuidos.
• Comparación con Línea Base: El método supera a una línea base de SAE Matryoshka (que utiliza trayectorias no residualizadas y normalizadas por componente) en la calidad de reconstrucción.

Análisis de la Estructura Espacio-Temporal

• Especialización Temporal: Los latentes exhiben perfiles temporales distintos. Las características "tempranas" tienden a tener un soporte temporal amplio y alta auto-similitud entre instantes de tiempo, lo que sugiere que capturan conceptos estructurales persistentes. Las características "tardías" son más específicas del instante de tiempo y menos estables temporalmente.
• Localización Espacial:
  • Las características tempranas muestran un soporte espacial positivo más amplio (mayor entropía espacial), cubriendo a menudo regiones difusas o similares a bordes.
  • Las características medias y tardías están más concentradas espacialmente (menor entropía), localizándose en áreas específicas similares a objetos o regiones salientes.

Experimentos de Dirección (Steering)

• Dirección de Característica Única: La intervención con latentes específicos produce efectos generativos visibles. Las características tempranas afectan la estructura global, mientras que las características medias pueden localizar cambios (por ejemplo, añadir joyas a un torso). Las características tardías mostraron efectos semánticos menos claros en algunos ejemplos.
• Transferencia de Características: Utilizando el conjunto de datos RIEBench, los autores realizaron transferencia de características de origen a destino. El modelo Concatenado Residualizado logró el mejor equilibrio entre la eficiencia de edición (mejorando la similitud CLIP con el prompt objetivo) y la preservación de la estructura de la imagen original (minimizando la distancia LPIPS). Superó a las variantes no residualizadas y por instante de tiempo en eficiencia de edición.

Significado y Afirmaciones

El artículo afirma que los SAE Temporales Residualizados proporcionan un marco útil para estudiar las activaciones de difusión estructuradas temporalmente. Al eliminar explícitamente la deriva predecible linealmente, el método permite que los latentes esparsos se centren en la estructura no redundante y semánticamente significativa del proceso de generación.

Los autores enfatizan que su enfoque difiere del trabajo anterior al aprender diccionarios esparsos sobre trayectorias de activación residualizadas en lugar de instantes de tiempo aislados o representaciones temporales no residualizadas. La capacidad de mapear estos latentes de vuelta al espacio de activación como trayectorias permite una comprensión más profunda de cuándo y dónde se expresan características específicas durante el proceso de eliminación de ruido.

Limitaciones señaladas por los autores:

• Los experimentos se centran en una única fuente de activación de U-Net (atención cruzada del bloque medio) en Stable Diffusion 1.5.
• Los experimentos de dirección son cualitativos; el artículo no afirma un rendimiento optimizado de edición de imágenes, sino que utiliza la dirección para validar la interpretabilidad y la influencia causal de las características aprendidas.