Modelos generativos: autoencoders, GANs y difusión

1. Autoencoder. Es la capa más chica, entre el codificador y el decodificador: obliga a resumir la imagen en pocos números, el código latente.
2. GAN. Es un clasificador que dice qué tan probable es que un dato sea real; el generador entrena para engañarlo.
3. Difusión. Dice cuánto ruido se agrega en cada paso de la ida. Es fijo: no se aprende.
4. Difusión. Es la red que estima el ruido ε que tiene encima x_t. Devuelve un vector del tamaño de x_t porque estima el ruido de cada coordenada (o de cada píxel).
5. GAN. Es el generador: aprende solo a través de lo que le responde el discriminador.

for latente in [2, 8, 32]:
    torch.manual_seed(0)
    modelo = Autoencoder(latente=latente)
    error = entrenar_ae(modelo, x_ent, x_val, epocas=5, mostrar=False)
    print(f"Cuello de {latente:2d} números: error {error[-1]:.4f} "
          f"(comprime {784 // latente} veces)")

Cuello de  2 números: error 0.0333 (comprime 392 veces)
Cuello de  8 números: error 0.0184 (comprime 98 veces)
Cuello de 32 números: error 0.0164 (comprime 24 veces)

Cuanto más ancho el cuello, menor el error: con 32 números caben más detalles que con 2. Pero la mejora se achica: pasar de 2 a 8 baja mucho el error y pasar de 8 a 32, bastante menos. Un cuello demasiado ancho, además, deja de obligar a la red a resumir: en el extremo, con 784 números, podría copiar sin aprender nada útil.

El de 2 números copia mal, pero tiene una ventaja: su espacio latente es un plano, y se puede dibujar. Si pintas el código de cada prenda de validación con el color de su clase, ves un mapa de las prendas parecido a los del tema 15 («Reducir dimensiones: PCA, t-SNE y UMAP»).

mnist_prueba = datasets.MNIST(root="datos", train=False, download=True)
x_mp = mnist_prueba.data.float().div(255).mul(2).sub(1).view(-1, 784)

torch.manual_seed(0)
clasificador = nn.Sequential(nn.Linear(784, 128), nn.ReLU(),
                             nn.Linear(128, 10))
opt = torch.optim.Adam(clasificador.parameters(), lr=1e-3)
cargador = DataLoader(TensorDataset(x_mnist, mnist.targets), batch_size=128,
                      shuffle=True, generator=torch.Generator().manual_seed(0))
for epoca in range(2):
    for x, y in cargador:
        perdida = nn.functional.cross_entropy(clasificador(x), y)
        opt.zero_grad()
        perdida.backward()
        opt.step()

with torch.no_grad():
    acierto = (clasificador(x_mp).argmax(1) == mnist_prueba.targets)
    generados = gen_dig(torch.randn(1000, ruido_dim, device=device)).cpu()
    clases = clasificador(generados).argmax(1)
print(f"Exactitud del clasificador con dígitos reales: "
      f"{acierto.float().mean():.3f}")
conteo = torch.bincount(clases, minlength=10)
for d in range(10):
    print(f"{d}: {conteo[d] / 10:4.1f} %  " + "#" * int(conteo[d] / 10))

Exactitud del clasificador con dígitos reales: 0.946
0:  1.5 %  #
1: 28.6 %  ############################
2:  9.3 %  #########
3: 12.9 %  ############
4:  9.4 %  #########
5:  6.6 %  ######
6:  1.9 %  #
7: 13.4 %  #############
8:  5.8 %  #####
9: 10.6 %  ##########

c) Aparecen los diez dígitos, así que no hay un colapso total: el generador no se quedó en un solo modo. Pero están lejos de salir parejos: el 1 se lleva casi el 29 %, y el 0 y el 6 apenas llegan al 2 %. Eso es un colapso parcial: el generador prefiere los dígitos que le resulta más fácil hacer pasar por reales (un 1 es un palito) y casi abandona los que no.

La limitación: el clasificador fue entrenado con dígitos reales y bien escritos. Ante un garabato de la GAN igual elige alguna clase, la que menos lejos le queda. Una parte del conteo son garabatos asignados a la clase que más se les parece, y no dígitos de verdad. El método sirve para ver si faltan clases enteras, no para medir la calidad de cada dígito.

a) Es la receta de cuánto ruido se agrega en cada paso de la ida: al principio poco, al final mucho, hasta que del dato no queda nada. «Queda de la señal» es cuánto del dato original sigue en la mezcla (√ᾱ_t) y «ruido» es cuánto ruido tiene encima (√(1 − ᾱ_t)). Los cuadrados de las dos cifras suman 1: lo que pierde una lo gana la otra. El modelo aprende a deshacer cada paso de esa receta.

# beta constante, con beta[0] = 0 como en la teoría
constante = torch.cat([torch.zeros(1), torch.full((T,), 0.02)])

for nombre, b in [("lineal", beta), ("constante", constante)]:
    senal = torch.cumprod(1 - b, dim=0).sqrt()
    print(f"{nombre:9s}: señal en t = 10: {senal[10]:.2f}, "
          f"t = 50: {senal[50]:.2f}, t = 100: {senal[100]:.2f}")

lineal   : señal en t = 10: 0.98, t = 50: 0.53, t = 100: 0.07
constante: señal en t = 10: 0.90, t = 50: 0.60, t = 100: 0.36

b) Con el programa constante, al principio se agrega más ruido que con el lineal (la espiral se borraría antes), y al final se agrega menos: en t = 100 todavía queda más de un tercio de la señal. Ese es el problema. Para generar, la vuelta arranca de ruido puro, torch.randn, porque se supone que en t = T ya no queda nada del dato. Con el programa constante, la red nunca vio en el entrenamiento algo tan ruidoso como ese punto de partida, y la vuelta empieza en un lugar que no conoce. Un buen programa de ruido termina en ruido puro y no borra la forma demasiado pronto. El lineal cumple lo primero, pero, como viste en la teoría, borra la espiral en los primeros 20 pasos. Los modelos de imágenes usan programas más cuidadosos, como el coseno (Nichol y Dhariwal, 2021), que reparten mejor los pasos.

g = torch.Generator().manual_seed(0)
ddim_10 = ddim(batch_size=4, num_inference_steps=10, generator=g,
               output_type="np").images
g = torch.Generator().manual_seed(0)
ddim_50 = ddim(batch_size=4, num_inference_steps=50, generator=g,
               output_type="np").images
print(f"Diferencia media por píxel, DDIM 10 contra 50 pasos: "
      f"{np.abs(ddim_10 - ddim_50).mean():.3f}")
print(f"Diferencia media por píxel, DDPM 50 contra DDIM 50: "
      f"{np.abs(imgs_ddpm - ddim_50).mean():.3f}")

Diferencia media por píxel, DDIM 10 contra 50 pasos: 0.032
Diferencia media por píxel, DDPM 50 contra DDIM 50: 0.241

Las imágenes de 10 y de 50 pasos casi coinciden: la diferencia es chica comparada con la que hay entre DDPM y DDIM con 50 pasos. Tiene sentido porque DDIM es determinista: no agrega ruido nuevo en cada paso, así que el mismo ruido inicial lleva siempre más o menos a la misma imagen, y los pasos solo deciden qué tan fino se recorre el camino. Con 10 pasos se pierde algo de detalle, pero tarda la quinta parte. Es el compromiso entre pasos y calidad: con DDIM se puede bajar mucho el número de pasos antes de que la calidad caiga.

a) Si reales y falsos son indistinguibles, cualquier dato tiene la misma probabilidad de venir de un lado o del otro. Lo mejor que puede decir es «no sé»: 0,5. Si dijera más, se equivocaría con los falsos; si dijera menos, con los reales.

b) Con D = 0,5 en todos lados, L_D = −log 0,5 − log(1 − 0,5) = 2 · log 2 y L_G = −log 0,5 = log 2:

d = 0.5
print(f"Pérdida del discriminador: {-np.log(d) - np.log(1 - d):.3f}")
print(f"Pérdida del generador:     {-np.log(d):.3f}")

Pérdida del discriminador: 1.386
Pérdida del generador:     0.693

Son las dos líneas punteadas de la figura de las pérdidas. Que la GAN del anillo se quedara ahí indica que llegó al empate. La de MNIST no llegó: su discriminador todavía distinguía.

# Paso 1: entrenar con ruido
torch.manual_seed(0)
limpiador = Autoencoder(latente=32)
historial = entrenar_ae(limpiador, x_ent, x_val, epocas=10, ruido=0.3,
                        mostrar=False)
print(f"Error de validación tras 10 épocas: {historial[-1]:.4f}")

# Paso 2: ¿cuánto mejora?
torch.manual_seed(1)
ruidosa = con_ruido(x_val, 0.3)
limpiador.eval()
with torch.no_grad():
    limpia = limpiador(ruidosa.to(device)).cpu()
print(f"Error con ruido:  {((ruidosa - x_val) ** 2).mean():.4f}")
print(f"Error limpiada:   {((limpia - x_val) ** 2).mean():.4f}")

# Paso 3: la figura
lote = torch.cat([ruidosa[:4], limpia[:4], x_val[:4]]).view(-1, 1, 28, 28)
mosaico = make_grid(lote, nrow=4, padding=2, pad_value=0)
plt.imshow(mosaico[0], cmap=tinta)
plt.yticks([15, 45, 75], ["Con ruido", "Limpiada", "Original"])
plt.xticks([])
plt.grid(False)
plt.title("Quitar ruido con un autoencoder")
plt.show()

Error de validación tras 10 épocas: 0.0167
Error con ruido:  0.0527
Error limpiada:   0.0168

# Paso 4a: error por clase
error_imagen = ((limpia - x_val) ** 2).mean(dim=1)
por_clase = [(error_imagen[y_val == k].mean().item(), nombres[k])
             for k in range(10)]
for error, nombre in sorted(por_clase, reverse=True)[:3]:
    print(f"Peor limpiada: {nombre:9s} error {error:.4f}")
for error, nombre in sorted(por_clase)[:2]:
    print(f"Mejor limpiada: {nombre:9s} error {error:.4f}")


def error_limpieza(x, ruido):
    """MSE de la imagen limpiada frente a la original."""
    with torch.no_grad():
        salida = limpiador(con_ruido(x, ruido).to(device)).cpu()
    return ((salida - x) ** 2).mean().item()


# Paso 4b: el doble de ruido del que vio al entrenar
print(f"Ruido 0,3 (el del entrenamiento): {error_limpieza(x_val, 0.3):.4f}")
print(f"Ruido 0,6 (el doble):             {error_limpieza(x_val, 0.6):.4f}")
# Paso 4c: dígitos de MNIST de prueba, que nunca vio (de 0 a 1, como
# las prendas)
mnist_prueba = datasets.MNIST(root="datos", train=False, download=True)
digitos_reales = mnist_prueba.data[:2000].float().div(255).view(-1, 784)
print(f"Dígitos de MNIST con ruido 0,3:   "
      f"{error_limpieza(digitos_reales, 0.3):.4f}")

Peor limpiada: sandalia  error 0.0238
Peor limpiada: bolso     error 0.0234
Peor limpiada: botín     error 0.0185
Mejor limpiada: pantalón  error 0.0103
Mejor limpiada: zapatilla error 0.0136
Ruido 0,3 (el del entrenamiento): 0.0168
Ruido 0,6 (el doble):             0.0286
Dígitos de MNIST con ruido 0,3:   0.0515

Un informe posible. El autoencoder quita bien el ruido del fondo: el MSE baja de 0,053 a 0,017, casi a un tercio. Pero falla en cuatro cosas. Primero, borra los detalles junto con el ruido: las prendas salen suaves, sin estampados ni rayas. Para el modelo, un estampado fino se parece al ruido, y con 32 números no alcanza a guardarlo. Segundo, le va peor con las prendas de forma variada: las sandalias, los bolsos y los botines tienen el error más alto, y los pantalones, que casi siempre tienen la misma forma, el más bajo. En la figura, un bolso de mano rectangular salió con forma de zapatilla: con tanto ruido, el modelo eligió la forma más común que se le parecía. Tercero, no generaliza a más ruido del que vio: con el doble de ruido, el error casi se duplica (0,029). Cuarto, no sabe limpiar lo que no conoce: con dígitos, el error (0,052) es tres veces el de las prendas y casi igual al de una prenda con ruido sin limpiar; el decodificador solo sabe dibujar prendas y convierte cada dígito en la prenda que más se le parece. Para mejorarlo se podría entrenar con niveles de ruido variados, usar una red convolucional como las del tema 20 («Redes convolucionales»), que respeta la vecindad de los píxeles, o entrenar más épocas con más datos.