"Denoising diffusion" en una dimensión (entre otras simplificaciones)

2023-1-10

I. Motivación e introducción

Denoising diffusion —DD en lo que sigue— es uno de los principales ingredientes del archipopular stable diffusion. Es un algoritmo que se usa fundamentalmente para generar imágenes y que funciona, a grandes rasgos así:

Se parte de un catálogo de imágenes, que son vectores en un espacio (de dimensión alta).
Esos vectores se difuminan utilizando un proceso concreto —piénsese en una especie de movimiento Browniano— hasta que su distribución es aproximadamente una normal (en ese espacio de dimensión elevada).
A partir de valores aleatorios de esa distribución normal, invirtiendo el proceso de difusión, se obtienen muestras del espacio original (de las fotos).

Subyace a todo este tinglado la conocida como hipótesis de la subvariedad. Todas las fotos son, en el fondo, vectores en $R^N$ donde si las fotos son, digamos, $1000 \times 1000$, $N$ es 3M (número de píxeles por el número de canales). La hipótesis de la subvariedad dice que la distribución de las fotos que reconocemos como tales —piénsese que la mayoría de las fotos de $R^N$ no dejan de ser manchas grises— residen en una subvariedad de dimensión baja incrustada en $R^N$. Generar imágenes equivale entonces a muestrear dicha subvariedad, con el problema de que no sabemos ni qué forma tiene ni dónde está. Lo que proporciona DD es un caminito para llegar a ella desde un punto cualquiera del espacio.

Esta entrada está basada en Denoising diffusion probabilistic models from first principles, una entrada en la que se discute el problema de la DD ejemplificándolo en un caso en dimensión dos. En él, la subvariedad es una espiral y en la parte final se pueden observar los caminitos (o aproximaciones a ellos) a los que me refiero en el párrafo anterior.

Nótese que en esa entrada se ha simplificado enormemente el problema al llamar fotos a puntos en una espiral en el plano. Yo voy a simplificar todavía mucho más la presentación anterior, comenzando por reducir al mínimo extremo la dimensionalida del problema, dejándola en 1 y llamando fotos a unos cuantos puntos en la recta numérica.

photos <- c(-1, 0, 1)

Como necesito muchos puntos de partida, en realidad voy a tomar muestras repetidas de ellas:

n_photos <- 1000
orig_photos <- sample(photos, n_photos, replace = TRUE)

(Este es un problema que ocurre solo en dimensión uno. De todos modos, casi nada cambiaría en lo que sigue si las fotos, en lugar de vivir en un conjunto discreto de puntos fuesen muestras de una distribución de probabilidad muy concentrada en un entorno pequeño de ellos.)

II. Difusión

A cada punto lo vamos a someter a una especie de movimiento browniano discretizado (en varios saltitos). En concreto:

n_jumps <- 60
means_jumps <- 0.01 * (1:n_jumps)
sd_jumps <- means_jumps

diffusion <- function(x, i){
  rnorm(n_photos, x - x * means_jumps[i], sd_jumps[i])
}

Vale, no es un movimiento browniano porque tiene un drift. De hecho, el drift, que es

$$x \longmapsto (1 - \mu_i) x,$$

donde los $\mu_i$ son valores muy pequeños:

Tiene un valor distinto en cada $x$.
Contrae, es decir, atrae centrípetamente cada $x$ hacia el origen.

Nota: Esas son las propiedades más importantes de la difusión.

Así que el drift contrae y, por otro lado, la pequeña $\sigma$ difumina ligeramente las trayectorias e impide que todo colapse lineal y aburridamene hacia el cero.

Así que difuminamos —guardando los pasos intermedios—

diffusion_data <- matrix(orig_photos, n_photos, n_jumps + 1)

for (i in 2:ncol(diffusion_data)) {
  diffusion_data[,i] <- diffusion(diffusion_data[, i-1], i - 1)
}

y obtenemos, entre otras cosas, un histograma del resultado final:

hist(diffusion_data[,ncol(diffusion_data)], breaks = 50)

III. Inversión de la difusión

Para generar una foto hay que:

Elegir un punto al azar de esa distribución final (la representada por histograma).
Invertir el proceso de difusión.

Lo primero es sencillo. De hecho, no nos vamos a contentar con una única foto sino que vamos a generar 100 de ellas. Para ello, vamos a tomar una muestra de puntos (o semillas) con la distribución que aproximadamente refleja el histograma:

tmp <- diffusion_data[,ncol(diffusion_data)]
seeds <- rnorm(100, mean(tmp), sd(tmp))

Para invertir el proceso de difusión hay que invertir cada uno de los saltitos por separado. Por ejemplo, el saltito 10 tiene esta pinta:

plot(diffusion_data[, 10],
     diffusion_data[, 11],
     xlab = "jump from",
     ylab = "jump to")

La inversa, por lo tanto, tiene este aspecto:

plot(diffusion_data[, 11],
     diffusion_data[, 10],
     xlab = "points at",
     ylab = "came from")

Nótese que podemos representar fácilmente esas funciones por trabajar en una dimensión. Las funciones correspondientes en dimensión $N$ lo serían entre vectores de $R^N$ y $R^N$. La ventaja de trabajar en una dimensión, además, es que tenemos mil maneras de aproximar esa función; en dimensiones más altas, todo el mundo usa redes neuronales. Usando —¿por qué no?— árboles (de cierta profundidad/complejidad):

inv_diff <- function(x, i) {
  my_x <- diffusion_data[,i+1]
  my_y <- diffusion_data[,i]

  mod <- ctree(my_y ~ my_x,
               data.frame(my_x = my_x, my_y = my_y),
               controls = ctree_control(
                 mincriterion = .1))

  predict(mod, data.frame(my_x = my_x))
}

Cada foo_i(x) <- inv_diff(x, i) representa una aproximación al saltito i, de manera que

out <- seeds
for (i in n_jumps:2) {
  print(i)
  out <- inv_diff(out, i)
}

debería ser una muestra —aproximada— de nuestras fotos originales. Y en efecto,

IV. Comentarios finales

En el fondo, por construcción, la mejor representación de la inversa de cada saltito sería una dilatación (la operación contraria a la contracción que representa el drift del movimiento seudobrowniano). Eso, al menos, para los primeros saltitos.
Es instructivo contemplar la función que invierte los saltitos como definiendo un campo vectorial en $R^N$ que apunta hacia la subvariedad donde residen las fotos.
En todo lo anterior hay cierto trabajo de calibración del tamaño de los saltos, etc. sine qua non.
No sé hasta qué punto todas las referencias a Langevin, etc. en la literatura sobre el tema no es puro ejercicio de erudición académica irrelevante en la práctica. Tal vez fueron relevantes en el momento de la ideación del proceso, pero podría vivirse sin ellos (sí, me estoy refiriendo aquí muy concretamente a los contextos de descubrimiento y de justificación).