Exploring Diffusion Models' Corruption Stage in Few-Shot Fine-tuning and Mitigating with Bayesian Neural Networks

Este artículo identifica y explica la etapa de corrupción en el ajuste fino de pocos ejemplos de modelos de difusión, proponiendo el uso de redes neuronales bayesianas para mitigar este fenómeno y mejorar la fidelidad, calidad y diversidad de las imágenes generadas sin incurrir en costos adicionales de inferencia.

Lo esencial

¡Hola! Imagina que tienes un artista genial llamado "Difusión" (Diffusion). Este artista ha pasado años viendo millones de cuadros, fotos y pinturas. Por eso, sabe pintar de todo: desde gatos voladores hasta paisajes futuristas. Es un maestro.

Pero, ¿qué pasa si quieres que este artista aprenda a pintar solo a tu perro, "Firulais", y solo tienes tres fotos de él para enseñarle?

Aquí es donde entra este paper, que cuenta una historia muy interesante sobre lo que sucede cuando intentamos enseñarle al artista con tan poca información.

1. El Problema: La "Fase de Corrupción" (El momento en que todo se pone feo)

Normalmente, cuando entrenas a una IA con pocas fotos, esperas que mejore poco a poco. Pero los autores descubrieron algo extraño y sorprendente:

1. Al principio: ¡Funciona genial! El artista empieza a entender quién es Firulais. Las fotos salen bonitas.
2. El momento extraño (La Corrupción): De repente, las fotos empiezan a salir mal. Aparecen manchas extrañas, ruido, como si el artista hubiera bebido demasiado café y empezara a alucinar. Las imágenes pierden calidad y se ven "sucias".
3. El final: Luego, el artista se recupera, pero se vuelve un obseso. Ya no pinta a Firulais en diferentes situaciones; solo copia exactamente las tres fotos que le diste. Si le pides "Firulais durmiendo", te da la foto exacta de "Firulais despierto". Ha perdido su creatividad y se ha quedado "atascado".

Los autores llaman a este momento feo y ruidoso la "Fase de Corrupción".

2. ¿Por qué pasa esto? (La analogía del mapa)

Imagina que el conocimiento del artista es un mapa gigante de todo el mundo.

• El modelo original: Tiene el mapa completo del planeta. Sabe dónde está París, Tokio y tu casa.
• El entrenamiento con pocas fotos: Cuando le das solo 3 fotos de Firulais, el artista intenta redibujar el mapa, pero borra todo lo demás y solo deja un punto muy pequeño donde está Firulais.

El problema: Como el mapa es tan pequeño y estrecho, cuando el artista intenta pintar algo nuevo (por ejemplo, "Firulais bajo la lluvia"), se pierde. Como no tiene un mapa amplio para guiarse, empieza a inventar cosas al azar (esas manchas ruidosas o "corrupción").

Es como si un chef que solo ha cocinado un tipo de pasta intentara hacer una pizza. Al no tener ingredientes variados en su "mapa" de cocina, la pizza sale quemada y con forma de pasta.

3. La Solución: Las Redes Neuronales Bayesianas (BNNs)

Para arreglar esto, los autores proponen usar una técnica llamada Redes Neuronales Bayesianas (BNNs).

La analogía de la duda saludable:

• El modelo normal: Es como un estudiante que cree que sabe la respuesta exacta al 100%. Si le das 3 fotos, cree que "Firulais es exactamente así" y no se permite imaginar nada más.
• El modelo con BNNs: Es como un estudiante que dice: "Bueno, creo que Firulais es así, pero también podría ser un poco diferente, o tener el pelo un poco más largo, o estar en otro lugar".

Las BNNs le dicen al modelo: "No seas tan seguro de ti mismo".

En lugar de aprender un punto fijo en el mapa, el modelo con BNNs aprende un área un poco más grande. Le permite tener "dudas" o variaciones. Esto hace que, cuando intenta pintar algo nuevo, no se pierda en el caos (la corrupción), porque tiene un mapa un poco más amplio para guiarse.

4. ¿Qué logran con esto?

Al aplicar esta "duda saludable" (BNNs):

1. Desaparecen las manchas: La "Fase de Corrupción" se reduce drásticamente. Las imágenes no salen ruidosas.
2. Más variedad: El artista puede pintar a Firulais durmiendo, corriendo o con un sombrero, y no solo copiar las fotos originales.
3. Calidad: Las imágenes salen más bonitas y realistas.

En resumen

Este paper descubre que cuando enseñamos a una IA con muy pocas fotos, a veces se vuelve loca y hace imágenes feas antes de volverse aburrida. Para arreglarlo, usan una técnica que le enseña a la IA a no ser tan rígida, permitiéndole imaginar un poco más allá de las pocas fotos que tiene. Es como darle al artista un mapa un poco más grande para que no se pierda, logrando resultados más creativos y de mejor calidad.

¡Y lo mejor de todo! Esto funciona con cualquier método de entrenamiento actual y no hace que la IA sea más lenta al momento de crear las imágenes.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "Exploring Diffusion Models' Corruption Stage in Few-Shot Fine-tuning and Mitigating with Bayesian Neural Networks", traducido y estructurado en español:

1. El Problema: La Etapa de Corrupción en el Ajuste Fino (Few-Shot)

Los autores identifican un fenómeno inesperado y crítico durante el proceso de ajuste fino con pocas muestras (few-shot fine-tuning) de Modelos de Difusión (DMs), como Stable Diffusion.

• La Dinámica Observada: Al ajustar un modelo preentrenado con un conjunto pequeño de imágenes (p. ej., 1 a 6 imágenes), la fidelidad de las imágenes generadas no mejora de manera monótona. En su lugar, se observa una curva de tres fases:
  1. Mejora inicial: La similitud con las imágenes de entrenamiento aumenta rápidamente.
  2. Etapa de Corrupción: La fidelidad cae abruptamente y aparecen patrones ruidosos y distorsionados en las imágenes generadas.
  3. Recuperación con Sobreajuste: La fidelidad vuelve a subir, pero el modelo entra en un estado de sobreajuste severo, perdiendo la capacidad de generar diversidad y produciendo solo copias exactas de las imágenes de entrenamiento.
• Causa Raíz: Mediante un modelado heurístico, los autores determinan que esta corrupción surge de una distribución de aprendizaje excesivamente estrecha. En el ajuste fino con pocas muestras, el modelo intenta aprender una distribución muy específica (un conjunto pequeño de puntos en el espacio latente), lo que provoca que, al intentar denoizar imágenes con ruido no visto, el modelo amplifique errores en lugar de eliminarlos, generando patrones corruptos.

2. Metodología: Integración de Redes Neuronales Bayesianas (BNN)

Para mitigar este problema, los autores proponen incorporar Redes Neuronales Bayesianas (BNN) en el proceso de ajuste fino de los DMs.

• Enfoque Conceptual: En lugar de tratar los parámetros del modelo como valores fijos, las BNN los modelan como variables aleatorias. Esto obliga al modelo a aprender una distribución de probabilidad sobre los parámetros en lugar de un único conjunto óptimo.
• Mecanismo de Acción:
  • La aleatoriedad inherente de las BNN actúa como una aumentación de datos implícita. Al muestrear diferentes parámetros durante el entrenamiento, el modelo se ve obligado a aprender una distribución más amplia y robusta ($I_\theta$), evitando que se "encalle" en la distribución estrecha que causa la corrupción.
  • Esto permite que el modelo maneje casos de baja confianza sin colapsar en patrones ruidosos.
• Formulación Matemática:
  • Se utiliza inferencia variacional para aproximar la distribución posterior de los parámetros $P(\theta|D)$.
  • La función de pérdida se descompone en dos términos:
    1. Pérdida de Difusión Esperada ($L_{DM}$): La expectativa de la pérdida de difusión estándar sobre la distribución de parámetros variacional.
    2. Término de Regularización ($L_r$): Una divergencia KL que restringe la distribución variacional a estar cerca de la distribución previa (el modelo preentrenado original).
  • La pérdida total es una combinación lineal: $W^* = \arg \min_W E_{\theta \sim Q_W(\theta)}[L_{DM}] + \lambda L_r$.
• Eficiencia:
  • Sin costo de inferencia: Durante la inferencia (generación), los parámetros se reemplazan por sus valores medios ($\mu_\theta$), por lo que no se añade ningún costo computacional extra en comparación con los métodos tradicionales.
  • Compatibilidad: El método es compatible con técnicas existentes como DreamBooth, LoRA y OFT. Solo se modelan como variables aleatorias un subconjunto de parámetros (por ejemplo, capas lineales o capas de normalización) para mantener los costos de entrenamiento bajos.

3. Contribuciones Clave

1. Descubrimiento y Definición: Son los primeros en identificar y nombrar la "Etapa de Corrupción" en el ajuste fino de DMs, caracterizándola por la aparición temporal de patrones ruidosos antes del sobreajuste final.
2. Modelado Teórico: Proporcionan un modelado heurístico basado en distribuciones gaussianas que explica matemáticamente por qué ocurre la corrupción (debido a una varianza alta y una distribución de aprendizaje limitada) y por qué desaparece al entrar en sobreajuste.
3. Solución Innovadora: Introducen el uso de BNNs en DMs para ajustar la distribución de aprendizaje de forma implícita, logrando un equilibrio entre fidelidad y diversidad sin costos adicionales de inferencia.
4. Validación Exhaustiva: Demuestran la eficacia del método en múltiples escenarios, incluyendo generación impulsada por objetos y sujetos, y su compatibilidad con diferentes arquitecturas de ajuste fino.

4. Resultados Experimentales

Los experimentos se realizaron en conjuntos de datos como DreamBooth (objetos) y CelebA-HQ (rostros), utilizando modelos como Stable Diffusion v1.4, v1.5 y v2.0.

• Métricas Cuantitativas:
  • Fidelidad de Texto (Clip-T) e Imagen (Dino, Clip-I): Las BNNs mejoraron consistentemente la alineación con el prompt y la similitud con el sujeto/objeto en todos los métodos base (DreamBooth, LoRA, OFT).
  • Calidad de Imagen (Clip-IQA): Se observó una mejora significativa en la calidad percibida, atribuida directamente a la reducción de los patrones corruptos.
  • Diversidad (Lpips): La introducción de aleatoriedad mejoró la diversidad de las imágenes generadas, evitando el colapso modal.
• Estudio de Usuarios: En una evaluación humana con 101 participantes, los métodos con BNNs fueron preferidos abrumadoramente (en la mayoría de los casos >60-70%) sobre los métodos sin BNNs, especialmente en alineación de texto y calidad de imagen.
• Análisis de Ablación:
  • Se encontró que una desviación estándar inicial moderada ($\sigma_\theta$) es óptima; valores demasiado altos causan inestabilidad.
  • El hiperparámetro $\lambda$ permite ajustar el equilibrio entre diversidad y fidelidad.
  • Aplicar BNNs solo en las capas de normalización (LN/GN) reduce el número de parámetros modificados a ~0.02% manteniendo un rendimiento fuerte.

5. Significado e Impacto

Este trabajo es fundamental para el campo de la generación de imágenes personalizada:

• Estabilidad en Entrenamiento: Resuelve un problema de inestabilidad crítica que afecta a las aplicaciones de IA generativa personalizadas, permitiendo un entrenamiento más seguro y predecible con pocos datos.
• Eficiencia: Ofrece una solución que no penaliza el tiempo de inferencia, lo cual es crucial para la implementación en producción.
• Nueva Perspectiva: Cambia la comprensión de la dinámica de entrenamiento de los DMs, sugiriendo que la aleatoriedad controlada (vía BNN) es una herramienta poderosa para regularizar modelos generativos en regímenes de pocos datos, en lugar de ser un obstáculo.

En resumen, la propuesta transforma la vulnerabilidad de los modelos de difusión ante el ajuste fino con pocas muestras en una oportunidad para mejorar la robustez, la calidad y la diversidad de las imágenes generadas mediante la aplicación de principios bayesianos.