¿Las Bayesian Flow Networks son modelos de difusión?

No exactamente. Generan de forma parecida al proceso inverso de la difusión, pero la red recibe los parámetros de una distribución que se actualiza con la regla de Bayes, no una versión ruidosa del dato, y no necesitan definir un proceso hacia adelante.

¿Por qué las Bayesian Flow Networks sirven para texto?

Porque para datos discretos la red recibe probabilidades por clase, que son números continuos. La entrada es diferenciable aunque el texto sea discreto, a diferencia de la difusión discreta, que cambia símbolos completos al azar.

¿Cuántos pasos necesita una Bayesian Flow Network para generar?

Se entrena en tiempo continuo, así que se puede correr con cualquier número de pasos. En text8, con 100 pasos llegó a 1.43 bits por carácter, frente a 1.41 en el límite continuo.

¿Las Bayesian Flow Networks le ganan a un Transformer autorregresivo en texto?

No. En text8 la BFN obtuvo 1.41 bits por carácter y un Transformer autorregresivo estándar 1.23. Lo que sí supera es a los modelos de difusión discreta que citan los autores.

¿Hay código de las Bayesian Flow Networks?

Sí. Los autores publicaron código y modelos entrenados en el repositorio nnaisense/bayesian-flow-networks de GitHub, y el paper incluye pseudocódigo para entrenar y generar con datos continuos, discretizados y discretos.

Bayesian Flow Networks explicado: difusión con Bayes

Respuesta corta: qué logran las Bayesian Flow Networks

Las Bayesian Flow Networks (BFN) de Graves, Srivastava, Atkinson y Gomez (NNAISENSE, 2023) generan datos actualizando con inferencia bayesiana los parámetros de distribuciones simples y pasándolos a una red neuronal, sin proceso de ruido hacia adelante. En text8 logran 1.41 bits por carácter, mejor que la difusión discreta que citan; un Transformer autorregresivo llega a 1.23.

Motivación de las Bayesian Flow Networks: difusión que funcione con datos discretos

Casi todos los modelos generativos modernos parten un problema difícil en piezas fáciles. Un modelo autorregresivo genera un token después de otro. Un modelo de difusión (diffusion) parte de ruido puro y lo limpia poco a poco hasta llegar a una imagen.

Los autores cuentan esto como una transmisión. Alice tiene los datos y le manda mensajes a Bob; Bob intenta adivinar cada mensaje antes de recibirlo. Entre mejor adivina, menos bits cuesta mandarlo. La pérdida es el total de bits, así que entrenar por máxima verosimilitud equivale a entrenar para comprimir.

Los modelos autorregresivos dominan el lenguaje, pero tienen un costo: generar requiere tantas pasadas de la red como variables tiene el dato. Una imagen de MNIST tiene 784 pixeles, así que pide 784 pasos. La difusión desacopla el número de pasos del número de variables y construye la imagen de lo grueso a lo fino, que según los autores explica su ventaja en imágenes.

El problema aparece con datos discretos, como texto. Ahí el ruido de la difusión también es discreto: cada símbolo se cambia al azar o se queda igual. Bob empieza con un texto revuelto y cada mensaje altera letras completas. Para los autores, ese ruido discontinuo es un obstáculo de fondo, y creen que un proceso continuo funcionaría mejor con texto.

Las Bayesian Flow Networks cambian qué recibe la red. En vez de una versión ruidosa del dato, la red recibe los parámetros de una distribución sobre el dato: medias para datos continuos, probabilidades por clase para texto. Esas probabilidades son números reales, así que la entrada es continua y diferenciable aunque el dato sea discreto.

Los autores proponen que esa continuidad abriría la puerta a guiar la generación con gradientes y a generar en pocos pasos, como ya se hace en difusión continua. Es una motivación, no algo que midan en este paper. Otra ventaja que sí muestran en el método: una BFN no necesita definir ni invertir un proceso hacia adelante, solo una regla de actualización bayesiana.

Cómo funcionan las Bayesian Flow Networks: cuatro distribuciones y una actualización

Una BFN maneja cuatro distribuciones. La distribución de entrada (input distribution) modela cada variable por separado: una normal por pixel o una categórica por carácter. Empieza en un prior simple, una normal estándar o una categórica uniforme. La distribución de salida (output distribution) la produce la red neuronal a partir de todos los parámetros de entrada, así que sí usa contexto, como los pixeles vecinos o las palabras cercanas.

La distribución del emisor (sender) es el dato con ruido, con una precisión $\alpha$ que controla qué tanto revela. La distribución del receptor (receiver) es la salida de la red con ese mismo ruido encima: lo que Bob esperaría recibir si su predicción fuera correcta.

Cada paso funciona así. La red recibe los parámetros de entrada y el tiempo $t$, y devuelve la distribución de salida. Alice toma una muestra $\mathbf{y}$ del emisor. Mandarla cuesta la divergencia KL entre emisor y receptor. Bob usa la muestra para actualizar los parámetros de entrada con la regla de Bayes. Para datos continuos, la actualización es la de una normal con precisión conocida (ecuaciones 49 y 50 del paper):

$$\rho_i = \rho_{i-1} + \alpha, \qquad \boldsymbol{\mu}_i = \frac{\boldsymbol{\mu}_{i-1}\rho_{i-1} + \mathbf{y}\,\alpha}{\rho_i}$$

$\boldsymbol{\mu}_{i-1}$ y $\rho_{i-1}$: la media y la precisión (inverso de la varianza) de la distribución de entrada antes del paso $i$.
$\mathbf{y}$: la muestra ruidosa que manda Alice, centrada en el dato con varianza $1/\alpha$.
$\alpha$: la precisión de esa muestra; un $\alpha$ grande significa poco ruido.
$\boldsymbol{\mu}_i$ y $\rho_i$: la media y la precisión después de ver la muestra.

La intuición: las precisiones se suman y la nueva media es un promedio ponderado entre lo que Bob creía y lo que acaba de ver. Si $\alpha$ es casi cero, la muestra no mueve nada; si es enorme, la media salta al valor observado.

Como las precisiones se suman, varios pasos equivalen a uno solo con la precisión total. Eso permite saltar directo a cualquier tiempo $t$ con la distribución de flujo bayesiano (Bayesian flow distribution), que es la que da nombre al método. La precisión acumulada sigue un calendario de precisión $\beta(t)$ (accuracy schedule) que crece desde cero. Para texto, con $K$ clases, el flujo queda así (ecuación 185 del paper):

$$p_{_F}(\boldsymbol{\theta} \mid \mathbf{x}; t) = \mathop{\mathbb{E}}_{\mathcal{N}(\mathbf{y} \mid \beta(t)(K\mathbf{e}_{\mathbf{x}} - \mathbf{1}),\, \beta(t)K\boldsymbol{I})} \delta\big(\boldsymbol{\theta} - \mathrm{softmax}(\mathbf{y})\big)$$

$\boldsymbol{\theta}$: las probabilidades por clase que recibe la red; viven en el simplex.
$\mathbf{e}_{\mathbf{x}}$: el vector one-hot de la clase correcta; $\mathbf{1}$ es un vector de unos.
$\beta(t)$: la precisión acumulada al tiempo $t$; para datos discretos los autores usan $\beta(t) = t^2\beta(1)$ como aproximación empírica.
$\mathbf{y}$: logits con ruido gaussiano cuya media empuja hacia la clase correcta.
$\delta$ y $\mathrm{softmax}$: la entrada es exactamente el softmax de esos logits.

La intuición: al principio $\beta(t)$ vale cero y la entrada es uniforme. Conforme crece, los logits de la clase correcta dominan al ruido y $\boldsymbol{\theta}$ se concentra en ella. La entrada nunca deja de ser un vector continuo de probabilidades.

La pérdida de $n$ pasos es la suma de las divergencias KL de todos los pasos, más el costo de mandar el dato al final. Los autores toman el límite cuando el número de pasos tiende a infinito y obtienen una pérdida en tiempo continuo más simple. Para datos discretos es (ecuación 205 del paper):

$$L^{\infty}(\mathbf{x}) = K\beta(1) \mathop{\mathbb{E}}_{t \sim U(0,1),\, p_{_F}(\boldsymbol{\theta} \mid \mathbf{x}, t)} t \,\lVert \mathbf{e}_{\mathbf{x}} - \hat{\mathbf{e}}(\boldsymbol{\theta}, t) \rVert^2$$

$t \sim U(0,1)$: un tiempo al azar; con él se muestrea $\boldsymbol{\theta}$ del flujo, sin simular la cadena completa.
$\hat{\mathbf{e}}(\boldsymbol{\theta}, t)$: las probabilidades que predice la red para cada clase.
$\lVert \mathbf{e}_{\mathbf{x}} - \hat{\mathbf{e}} \rVert^2$: el error cuadrático entre el one-hot real y la predicción.
$t$ y $K\beta(1)$: pesos; los errores tardíos cuestan más porque ahí la precisión crece más rápido.

La intuición: entrenar una BFN es pedirle a la red que adivine el dato a partir de una creencia ruidosa, con una sola pasada por ejemplo. Para generar se parte del prior y se repite el ciclo con $n$ pasos: muestrear de la salida, agregar ruido, actualizar con Bayes. Una BFN entrenada en tiempo continuo se puede correr con cualquier número de pasos.

Ejemplo visual del flujo bayesiano en una Bayesian Flow Network

Para datos continuos, el flujo bayesiano tiene forma cerrada: la media de entrada sigue una normal con media $\gamma(t)\,x$ y varianza $\gamma(t)(1-\gamma(t))$, donde $\gamma(t) = 1 - \sigma_1^{2t}$. La gráfica usa los valores de la Figura 4 del paper y muestra la media esperada con una banda de una desviación estándar.

Cómo se acerca la media de entrada al dato en el flujo bayesiano continuo. Gráfica ilustrativa, no son datos del paper. Fuente: elaboración propia con la distribución de flujo de la ecuación (77) y el calendario de la ecuación (80) de Graves et al. (2023), con los valores de su Figura 4.

La banda explica una diferencia clave con la difusión. Al inicio la entrada es el prior, sin ruido. Muy pronto la incertidumbre llega a su máximo y después se cierra alrededor del dato. En un modelo de difusión, en cambio, la red empieza recibiendo ruido puro. Los autores comparan esta varianza de entrada con los calendarios lineal y coseno de difusión en su Figura 5, y la de la BFN es mucho menor. Plantean como hipótesis, sin probarla, que esa entrada menos ruidosa podría acelerar el aprendizaje en datasets grandes.

Resultados de las Bayesian Flow Networks en text8, CIFAR-10 y MNIST

Los autores entrenan todas las redes con la pérdida en tiempo continuo y evalúan con la pérdida de $n$ pasos para varios valores de $n$. Reportan verosimilitud, no calidad visual: no calcularon FID. Las cifras de casi todos los modelos, incluida la BFN, son cotas superiores. Los competidores salen de la literatura; no los reentrenaron.

Texto: text8

En text8 (Wikipedia en minúsculas con un alfabeto de 27 símbolos y secuencias de 256 caracteres) usaron un Transformer de 24 capas con unos 170M parámetros. La BFN obtuvo 1.41 bits por carácter en prueba. Eso es mejor que todos los modelos de difusión discreta de su tabla, con D3PM mask en 1.45 como el más cercano, y casi igual a MAC, un modelo de orden arbitrario, con 1.40.

Bits por carácter en text8 con secuencias de 256 caracteres; menos es mejor. Fuente: Graves et al. (2023), Tabla 4.

Los autores reconocen que un Transformer autorregresivo estándar y un modelo de flujo discreto quedan bastante mejor, ambos con 1.23. Con pocos pasos la BFN aguanta: con 100 pasos llega a 1.43, frente a 1.41 en el límite continuo. Con 10 pasos sube a 1.70.

Imágenes: CIFAR-10 y MNIST

En CIFAR-10 con 256 niveles por canal, la mejor BFN da 2.66 bits por dimensión, cerca del mejor resultado sin aumento de datos que citan, 2.64. La comparación natural es con Variational Diffusion Models (VDM), que comparte arquitectura y llega a 2.65. Sin embargo, VDM entrenó 10M actualizaciones y la BFN solo 5M por falta de tiempo; su validación seguía mejorando.

La Tabla 3 compara la pérdida continua contra una pérdida discretizada que sabe que los pixeles caen en niveles fijos. Con pocos pasos la discretizada gana con claridad: con 10 pasos da 3.91 frente a 6.18. Con muchos pasos se emparejan, y en el límite la continua queda ligeramente mejor, 2.66 frente a 2.68.

Bits por dimensión en CIFAR-10 (256 bins) según los pasos de evaluación; menos es mejor. Fuente: Graves et al. (2023), Tabla 3.

En MNIST binarizado dinámicamente, la BFN obtiene 77.87 nats por imagen, mejor que NVAE (78.01) pero detrás de Locally Masked PixelCNN (77.58). Con 784 pasos, los mismos que necesitaría un autorregresivo, da 78.02; con 10 pasos sube a 95.21.

Limitaciones de las Bayesian Flow Networks: qué no prueba el paper

Solo verosimilitud en tres benchmarks pequeños. No hay FID, evaluación humana ni pruebas en datasets grandes. Las muestras de texto generadas se muestran sin métrica.
En texto no alcanza a los autorregresivos. En text8, 1.41 frente a 1.23 de un Transformer estándar. Tampoco comparan contra difusión continua para texto basada en embeddings o en el simplex.
Calendarios heurísticos. El $\beta(t)$ discreto es una aproximación empírica. En MNIST la pérdida de reconstrucción quedó alta, 0.46 nats por imagen, y los autores concluyen que el calendario no es óptimo para datos binarios.
Comparaciones con presupuesto distinto. En CIFAR-10 la BFN entrenó la mitad de actualizaciones que VDM. Además, el texto y la Tabla 3 no coinciden sobre cuál pérdida entrenó más rápido con 16 bins.
Promesas sin medir. La guía por gradientes y la generación en pocos pasos para texto quedan como dirección futura.

Qué significa para ti: cuándo mirar las Bayesian Flow Networks

Si trabajas con datos discretos (texto, secuencias, etiquetas o datos binarios) y quieres un generador que no sea autorregresivo, las BFN son una opción con una propiedad útil: la red siempre recibe vectores continuos de probabilidades. Los autores proponen que eso facilitaría guiar la generación con gradientes, algo que la difusión discreta no permite de forma natural.

Si generas imágenes, la BFN se parece mucho a un modelo de difusión variacional. No esperes una mejora clara de calidad: el paper reporta verosimilitud muy cercana a VDM, no mejores muestras.

La parte práctica más atractiva es la pérdida. No necesitas definir un proceso hacia adelante ni matrices de transición. Basta muestrear un tiempo, armar la entrada con un softmax de logits ruidosos y medir un error cuadrático. Este es el Algoritmo 8 del paper en NumPy, con una "red" de juguete que solo copia su entrada:

import numpy as np


def perdida_bfn_discreta(x, red, K, beta1, rng):
    """Pérdida en tiempo continuo de una BFN para datos discretos (Algoritmo 8 del paper)."""
    e_x = np.eye(K)[x]                          # one-hot, forma (D, K)
    t = rng.uniform()                           # tiempo al azar en (0, 1)
    beta = beta1 * t**2                         # calendario de precisión beta(t)
    y = rng.normal(beta * (K * e_x - 1), np.sqrt(beta * K))  # logits ruidosos
    theta = np.exp(y - y.max(axis=1, keepdims=True))
    theta /= theta.sum(axis=1, keepdims=True)   # softmax: parámetros de entrada
    e_hat = red(theta, t)                       # probabilidades de salida (D, K)
    return K * beta1 * t * np.sum((e_x - e_hat) ** 2)


rng = np.random.default_rng(0)
x = rng.integers(0, 27, size=256)               # una secuencia de 256 caracteres
red_ingenua = lambda theta, t: theta            # "red" que solo copia su entrada
print(perdida_bfn_discreta(x, red_ingenua, K=27, beta1=0.75, rng=rng))

En una implementación real, como la de text8, la red sería un Transformer que recibe $\boldsymbol{\theta}$ reescalado a $[-1, 1]$ y el tiempo $t$, y devuelve un softmax por posición. Los autores publicaron su código y modelos entrenados en github.com/nnaisense/bayesian-flow-networks, y el PDF del paper tiene pseudocódigo para entrenar, evaluar y generar en los tres tipos de datos.

Lo que no cambia: para modelar lenguaje con la mejor verosimilitud, un Transformer autorregresivo sigue ganando en este benchmark. Las BFN son una línea de investigación para generación no autorregresiva, no un reemplazo listo para producción.

Bayesian Flow Networks: redes de flujo bayesiano