Investigadores de Sakana AI y la Universidad de Tokio proponen DiffusionBlocks. Entrena redes basadas en transformadores un bloque a la vez. La memoria de entrenamiento se reduce en un factor de B, donde B es el número de bloques. El rendimiento se mantiene en diversas arquitecturas.
El problema de la memoria en el entrenamiento de redes neuronales
La retropropagación de un extremo a otro requiere almacenar activaciones intermedias en cada capa. El consumo de memoria crece linealmente con la profundidad de la red. A medida que los modelos se vuelven más profundos, esto se convierte en un importante cuello de botella en el entrenamiento.
Una técnica existente, los puntos de control de activación, reduce la memoria de activación al recalcular las activaciones bajo demanda. Sin embargo, no reduce la memoria para parámetros, gradientes o estados del optimizador. Con el optimizador Adam, cada capa requiere memoria para parámetros, gradientes y dos estados del optimizador (impulso y varianza). Esto suma 4 veces el tamaño del parámetro por capa, sin cambios mediante puntos de control de activación.
El entrenamiento por bloques ofrece un enfoque diferente. Dividir una red en bloques B y entrenar cada uno de forma independiente reduce la memoria a aproximadamente 1/B. La reducción es proporcional al número de bloques. El desafío es definir un objetivo local basado en principios para cada bloque que aún produzca un modelo globalmente coherente.
Los enfoques anteriores, como el algoritmo Forward-Forward de Hinton y el codicioso entrenamiento por capas, se basan en objetivos locales ad hoc. Constantemente tienen un rendimiento inferior al de la formación de un extremo a otro y se limitan en gran medida a tareas de clasificación.
DiffusionBlocks aborda tanto la brecha teórica como la aplicabilidad limitada de los métodos anteriores.
La idea central: conexiones residuales como pasos de Euler
La idea clave se basa en una conexión establecida en la literatura. Las redes residuales actualizan cada entrada de capa a través de zℓ=zℓ−1+fθℓ(zℓ−1)zℓ = zℓ−1 + fθℓ (zℓ−1). Esto corresponde a la discretización de Euler de ecuaciones diferenciales ordinarias.
El equipo de investigación muestra que estas actualizaciones corresponden específicamente a la EDO de flujo de probabilidad en modelos de difusión basados en puntuaciones. En la formulación de Variance Exploding (VE), el proceso de difusión inversa es el siguiente:
d𝐳σdσ=−σ∇𝐳logpσ(𝐳σ) frac{mathrm{d}mathbf{z}_sigma}{mathrm{d}sigma} = -sigma nabla_{mathbf{z}} log p_sigma(mathbf{z}_sigma)
La aplicación de la discretización de Euler a esta ecuación produce una regla de actualización que coincide estructuralmente con la actualización de la conexión residual. Una pila de bloques residuales se puede interpretar como pasos de eliminación de ruido discretizados. Los pasos abarcan un rango de niveles de ruido [𝞂min, 𝞂max].
En los modelos de difusión basados en puntuaciones, el objetivo de coincidencia de puntuaciones se puede optimizar de forma independiente en cada nivel de ruido. Esto significa que cada bloque se puede entrenar de forma independiente, utilizando únicamente su propio objetivo local. No se necesita comunicación entre bloques durante el entrenamiento.
Convertir una red: tres pasos
Convertir una red residual estándar a DiffusionBlocks requiere tres modificaciones:
Partición de bloques: divida la red de capa L en bloques B. Cada bloque contiene un grupo contiguo de capas. Asignación de rango de ruido: defina una distribución de ruido pruido y un rango de ruido [𝞂min, 𝞂max]. Divida este rango en intervalos B y asigne un intervalo a cada bloque. El equipo de investigación recomienda una distribución log-normal para el ruido. Acondicionamiento de ruido: amplíe la entrada de cada bloque para incluir una versión ruidosa del objetivo. Agregue acondicionamiento del nivel de ruido a través de AdaLN (Normalización de capa adaptativa). Cada bloque aprende a predecir el objetivo limpio a partir de su versión ruidosa dentro de su rango de ruido asignado.
Durante el entrenamiento, se muestrea un solo bloque por iteración. Los demás bloques no se computan. El consumo de memoria corresponde a las capas L/B, no a todas las capas L.
Partición de equiprobabilidad
Una partición uniforme simple divide [𝞂min, 𝞂max] en intervalos iguales. Esto ignora la diferente dificultad de eliminar el ruido según los niveles de ruido. Los niveles de ruido intermedios contribuyen más a la calidad de la generación bajo la distribución de entrenamiento log-normal.
En su lugar, DiffusionBlocks utiliza particiones de equiprobabilidad. Los límites se eligen de modo que cada bloque maneje exactamente 1/B de la masa de probabilidad total bajo ruido. Los bloques asignados a niveles de ruido intermedios reciben intervalos más estrechos. Los bloques que manejan regiones de ruido extremo reciben intervalos más amplios.
En estudios de ablación en CIFAR-10 utilizando DiT-S/2, se deshabilitó la superposición de bloques para aislar cada componente. La partición equiprobabilidad logró un FID de 38,03 frente a 43,53 para la partición uniforme (cuanto más bajo, mejor). Ambos utilizaron una distribución de capas uniforme de [4,4,4] en 3 bloques.
Resultados experimentales
El equipo de investigación evaluó DiffusionBlocks en cinco arquitecturas que abarcan tres categorías de tareas. Todos los resultados comparan DiffusionBlocks (entrenados por bloques) con la misma arquitectura entrenada con retropropagación de un extremo a otro.
Comparación adelante-adelante: en CIFAR-100, el algoritmo adelante-adelante logró solo un 7,85 % de precisión bajo la misma arquitectura ViT. Esto resalta la brecha entre los objetivos contrastantes ad hoc y el objetivo de coincidencia de puntajes utilizado por DiffusionBlocks.
Eficiencia de inferencia DiT: para los modelos de difusión, cada paso de eliminación de ruido durante la inferencia activa solo un bloque. Un DiT de 12 capas con B=3 utiliza solo evaluaciones de 4 capas por paso de eliminación de ruido. Esta es una reducción de cálculo de inferencia de 3 veces en comparación con la ejecución de las 12 capas.
Entrenamiento de Huginn: Huginn aplica el mismo bloque recurrente de 4 capas de forma recurrente. Utiliza una profundidad de recurrencia estocástica con un promedio de 32 iteraciones. El entrenamiento utiliza propagación hacia atrás truncada a través del tiempo (BPTT) de 8 pasos. DiffusionBlocks reemplaza esto con un único pase hacia adelante por paso de entrenamiento. El procedimiento de inferencia de K-iteración se mantiene sin cambios. La reducción de iteraciones de 32 veces supera el programa de entrenamiento 3 veces más largo. DiffusionBlocks entrena durante 15 épocas frente a las 5 épocas de Huginn. El cálculo total se reduce aproximadamente 10 veces.
Resultados de OpenWebText: en OpenWebText, DiffusionBlocks MAUVE fue 0,82 frente a 0,85. La perplejidad generativa bajo Llama-2 fue de 14,99 frente a 15,05. Los resultados de este conjunto de datos fueron mixtos, con algunas métricas ligeramente peores que la línea de base.
Partición de difusión enmascarada: para los modelos de difusión enmascarada, la partición de bloques apunta al programa de enmascaramiento en lugar de a los niveles de ruido continuo. Cada bloque maneja una disminución igual en la probabilidad alfa de desenmascaramiento.
Comparación con NoProp
NoProp es un trabajo concurrente que utiliza un marco de difusión para la capacitación sin retropropagación. Se evalúa únicamente en tareas de clasificación utilizando una arquitectura personalizada basada en CNN. No proporciona un procedimiento para aplicar el método a otras arquitecturas o tareas.
DiffusionBlocks es el único método que combina una formulación de tiempo continuo con entrenamiento por bloques. Se mantiene dentro de 1 punto porcentual de la línea base de retropropagación de un extremo a otro.
Fortalezas y debilidades
Fortalezas:
Base teórica basada en principios a través de coincidencia de puntajes, no objetivos locales ad hoc Funciona en cinco arquitecturas distintas sin modificaciones específicas de la tarea Reducción de memoria de entrenamiento de B×, proporcional al número de bloques Para los modelos de difusión, el cálculo de inferencia también se reduce en B× durante la generación La partición de equiprobabilidad supera significativamente la partición uniforme (FID 38.03 vs 43.53 en CIFAR-10) Reemplaza BPTT de iteración K en modelos de profundidad recurrente con un solo paso hacia adelante Los bloques se pueden entrenar en paralelo entre GPU sin sobrecarga de comunicación. Los recuentos de bloques moderados (B=2 o B=3) a veces mejoran la FID con respecto al entrenamiento de un extremo a otro.
Debilidades:
Requiere dimensiones de entrada y salida coincidentes; actualmente no se puede aplicar a arquitecturas estilo U-Net. Validado solo en modelos entrenados desde cero; el ajuste fino de los modelos previamente entrenados no se ha probado No hay ningún método basado en principios para seleccionar el recuento de bloques óptimo para una arquitectura y una tarea determinadas Agrega una sobrecarga de acondicionamiento de ruido: el tiempo de pared agregado es 0,0543 s versus 0,0507 s bajo el entrenamiento estándar En OpenWebText, algunas métricas son marginalmente peores que la línea de base autorregresiva
Explicador visual de Marktechpost
Bloques de difusión · Sakana AI
ICLR 2026 · Entrenamiento por bloques
01 / 10 Siguiente →
Conclusiones clave
DiffusionBlocks divide las redes residuales en B bloques entrenables de forma independiente, lo que reduce la memoria de entrenamiento en un factor de B. Las conexiones residuales en los transformadores se asignan a los pasos de Euler del proceso de difusión inversa, lo que proporciona un objetivo de entrenamiento local basado en principios para cada bloque. La partición de equiprobabilidad asigna una masa de probabilidad igual por bloque, no intervalos de ruido iguales, mejorando significativamente la generación de imágenes FID sobre la partición uniforme. Validado en cinco arquitecturas: ViT, DiT, difusión enmascarada, transformadores autorregresivos y de profundidad recurrente. Los modelos de profundidad recurrente como Huginn reemplazan el BPTT de iteración K con un solo paso hacia adelante, lo que reduce el cálculo total de entrenamiento en aproximadamente 10 veces.
Consulte el artículo de investigación, el repositorio y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 150.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros