De texto a imagen

La pérdida contrastiva de CLIP

En un lote de $N$ pares, el codificador de imágenes $f$ y el de texto $g$ producen vectores unitarios $u_i = f(x_i)$ y $v_j = g(t_j)$. Su similitud coseno es $S_{ij} = u_i \cdot v_j$. Se dividen entre una temperatura $\tau$ (que también se aprende) y se pide, con entropía cruzada, que cada imagen elija su propia frase y que cada frase elija su propia imagen:

Es la pérdida simétrica del artículo de Radford y colaboradores (figura 3): sube la diagonal de la tabla $S$ y baja el resto de su fila y de su columna. CLIP usó lotes de 32 768 pares, es decir, tablas de más de mil millones de casillas en cada paso, y unos 400 millones de pares de imagen y texto reunidos de fuentes públicas de internet. La ficha del modelo es de enero de 2021.

En nuestro juguete, el codificador de imágenes recibe 720 números (12 × 12 píxeles, cada uno con 5 casillas one-hot de color) y tiene una capa oculta de 96; el de texto lee 5 posiciones de palabra, cada palabra con un vector de 16 más un vector de su posición, y tiene una capa oculta de 64. Ambos entregan $D = 32$. Se entrenó 2 500 pasos con lotes de 96; la temperatura pasó de $1/\tau \approx 14.3$ a $22.0$. En el lote fijo de seis pares que muestra la escena, la diagonal termina entre 0.78 y 0.88. El error del lote no llega a cero con lotes de 96: muchas frases incompletas («círculo») se repiten dentro del mismo lote y no pueden distinguir a sus imágenes.

Condicionar y guiar

La red de difusión predice el ruido $\epsilon_\theta(z_t, t, c)$ de un latente $z_t$ de 16 números, en el paso $t$, dado el vector $c$ de la frase. Al entrenar, la frase se sustituye por la frase vacía $\varnothing$ el 10 % de las veces (como en Stable Diffusion v1-4). Al muestrear, con escala de guiado $s$:

Con $s = 0$ la frase no cuenta; con $s = 1$ es la red condicionada tal cual; con $s > 1$ se exagera. Es la convención de la biblioteca diffusers (escala por omisión 7.5, «activa cuando es mayor que 1»). Ho y Salimans escriben lo mismo como $(1+w)\,\epsilon_c - w\,\epsilon_\varnothing$, con $w = s - 1$, y lo presentan como un equilibrio entre calidad de las muestras y diversidad. GLIDE (Nichol y colaboradores, 2021) ya lo usaba con texto, y los evaluadores humanos lo prefirieron al guiado con CLIP.

Nuestro muestreo sigue la receta de la lección 17: $T = 40$ pasos, $\beta_t$ lineal de 0.02 a 0.30 y varianza del paso inverso $\sigma_t^2 = \beta_t$. El latente sale de un VAE de 16 números y se decodifica al final. Stable Diffusion condiciona de otra manera (atención cruzada sobre los vectores de cada palabra del codificador de texto de CLIP); aquí, para que quepa en el navegador, el vector de la frase se pega a la entrada de una red de tres capas.

Lo que midió la escena

• Escala de guiado con «círculo» (256 ruidos iniciales): círculos 36 % con $s = 0$, 80 % con $s = 1$, 95 % con $s = 2$, 93 % con $s = 4$ y 63 % con $s = 8$. Las veinticuatro imágenes de la escena (semillas 1000 a 1023, prompt «círculo») dan 11, 19, 21 y 16 círculos con $s = 0, 1, 3, 8$, y 18, 12, 7 y 7 combinaciones distintas de forma, color y fondo (con el prompt en inglés, «circle»: 11, 20, 23 y 13 círculos; 18, 11, 6 y 8 combinaciones).
• Sesgo: en los datos, el 64.9 % de los círculos son magenta (una de cada dos imágenes lleva el color típico de su forma y la otra, uno al azar). Entre las veinticuatro imágenes, 15 son magenta con $s = 1$ y 22 con $s = 3$.
• Lo nunca visto (diez combinaciones fuera del entrenamiento, 16 ruidos por combinación, $s = 3$): 86.4 % correctas en las 70 combinaciones vistas y 75.6 % en las 10 nunca vistas. Con $s = 1$ la diferencia es mayor: 85.6 % frente a 45.6 %.
• Quién va con quién (60 parejas con los colores intercambiados, en los dos idiomas): el codificador que lee el orden elige la frase correcta el 100 % de las veces; la bolsa de palabras da exactamente el mismo vector a las dos frases (0 %: siempre empate).

Las cuentas de forma y color las hace un lector de reglas fijas, no aprendido: el fondo es el color más común en el borde, la figura es el resto, y la forma es la plantilla (círculo, cuadrado, triángulo o cruz, en varias posiciones y tamaños) que mejor se superpone; si ninguna supera 0.6 de intersección sobre unión, la imagen cuenta como «sin forma». Lee bien las 20 imágenes reales de la lección. Con los cuadrados es menos fiable: el modelo los dibuja a menudo con esquinas redondeadas y el lector los toma por círculos, por eso la escena no los usa para contar.

En código

import numpy as np

def perdida_contrastiva(U, V, tau=0.07):
    """U, V: vectores unitarios de N imágenes y N frases (fila i = par i)."""
    S = U @ V.T / tau
    S = S - S.max(axis=1, keepdims=True)
    filas = -np.mean(np.diag(S) - np.log(np.exp(S).sum(axis=1)))
    S2 = (U @ V.T / tau).T
    S2 = S2 - S2.max(axis=1, keepdims=True)
    cols = -np.mean(np.diag(S2) - np.log(np.exp(S2).sum(axis=1)))
    return (filas + cols) / 2

rng = np.random.default_rng(0)
U = rng.normal(size=(4, 8)); U /= np.linalg.norm(U, axis=1, keepdims=True)
V_azar = rng.normal(size=(4, 8)); V_azar /= np.linalg.norm(V_azar, axis=1, keepdims=True)
V_bien = U + 0.1 * rng.normal(size=(4, 8)); V_bien /= np.linalg.norm(V_bien, axis=1, keepdims=True)
print(round(perdida_contrastiva(U, V_azar), 3), round(perdida_contrastiva(U, V_bien), 6))

def guiado(e_sin, e_con, s):
    return e_sin + s * (e_con - e_sin)

e_sin, e_con = np.array([0.2, -0.4]), np.array([0.5, 0.1])
for s in (0, 1, 3):
    print(s, guiado(e_sin, e_con, s))
# 6.064 0.000141
# 0 [ 0.2 -0.4]
# 1 [0.5 0.1]
# 3 [1.1 1.1]

Con frases al azar el error es alto; con frases casi iguales a su imagen, la diagonal domina y el error cae a casi cero. La segunda función es toda la escala de guiado.

Lo que viene

Aquí el texto solo entraba: se convertía en un vector. En la lección 19 el texto sale: un modelo de lenguaje asigna probabilidades a la siguiente palabra y escribe una palabra tras otra.