El muro de la memoria en el aprendizaje profundo
Entrenar redes neuronales profundas siempre ha sido una tarea ávida de memoria.
La retropropagación de extremo a extremo exige que cada activación intermedia se almacene hasta el paso hacia atrás.
A medida que las redes se vuelven más profundas —especialmente con grandes transformers— la huella de memoria se infla linealmente con la profundidad.
Esto crea un “muro de memoria” que limita el tamaño y la accesibilidad del modelo.
Imagina una larga caminata en la que debes cargar desde el principio toda la comida y el agua que necesitarás para el viaje de regreso.
Así es exactamente como funciona el entrenamiento estándar: toda la pasada hacia adelante debe permanecer en memoria, lo que impide escalar más.
Los investigadores han buscado durante mucho tiempo formas de entrenar redes en piezas más pequeñas e independientes, pero hasta ahora la mayoría de los intentos sacrificaban rendimiento a cambio de ahorro de memoria.
Entrenamiento por bloques: una idea prometedora pero con defectos
El entrenamiento por bloques divide una red en fragmentos que se entrenan por separado, cada uno requiriendo memoria solo para sus propias activaciones.
En teoría, esto podría reducir drásticamente el uso máximo de memoria.
Sin embargo, los métodos anteriores consistentemente rinden menos que el entrenamiento de extremo a extremo.
Se basan en objetivos locales ad-hoc —como el preentrenamiento voraz por capas o señales de gradiente sintéticas— sin una forma fundamentada de coordinar los bloques.
Además, los enfoques existentes se han limitado en gran medida a tareas de clasificación y arquitecturas personalizadas.
Su aplicabilidad a los modelos transformer modernos y a la IA generativa permaneció prácticamente inexplorada.
Estos dos desafíos centrales —falta de fundamento teórico y alcance arquitectónico limitado— han impedido que el entrenamiento por bloques se convierta en una herramienta práctica para la corriente principal del aprendizaje profundo.
Una conexión sorprendente: redes residuales como procesos de difusión
Una idea crucial tiende un puente: las conexiones residuales dentro de los transformers pueden interpretarse desde la óptica de los modelos de difusión basados en score matching.
En un modelo de difusión, los datos se destruyen gradualmente añadiendo ruido, y un proceso inverso aprende a eliminar el ruido paso a paso.
Cada paso de eliminación de ruido se puede optimizar de forma independiente de los demás.
Resulta que las redes con capas residuales se asemejan mucho a la discretización de Euler de una EDO de flujo de probabilidad en tiempo continuo.
Cada bloque residual actúa como un paso infinitesimal de eliminación de ruido en una trayectoria de difusión continua.
Esta correspondencia ya se conocía en el contexto de las EDOs neuronales, pero DiffusionBlocks la aprovecha con un propósito novedoso: replantear un transformer profundo como una pila de etapas independientes de eliminación de ruido.
El marco de trabajo DiffusionBlocks
DiffusionBlocks divide una red transformer en bloques, a cada uno se le asigna un rango contiguo de niveles de ruido.
Durante el entrenamiento, solo el bloque correspondiente al nivel de ruido actual está activo; el resto se puede ignorar por completo.
Los gradientes fluyen únicamente dentro de ese único bloque.
El consumo de memoria disminuye en proporción al número de bloques, lo que hace viable el entrenamiento de modelos muy profundos en hardware modesto.
Y lo que es fundamental, el objetivo de cada bloque se deriva directamente de la teoría de score matching, no de una función de pérdida diseñada a mano.
Cada bloque aprende a eliminar el ruido de su rango asignado y, juntos, aproximan fielmente el proceso global de difusión inversa.
Las técnicas de recetas consolidadas de entrenamiento de difusión (por ejemplo, Karras et al., 2022) se pueden integrar sin problemas para mejorar aún más el rendimiento.
Particionamiento por equiprobabilidad para un aprendizaje equilibrado
Una partición ingenua que simplemente divide la línea temporal de ruido en segmentos iguales corre el riesgo de subutilizar algunos bloques.
La densidad de estados de ruido varía, y algunos intervalos recibirían mucha más señal de entrenamiento que otros.
DiffusionBlocks introduce el particionamiento por equiprobabilidad: los niveles de ruido se dividen de manera que cada bloque represente una cantidad igual de masa de probabilidad acumulada bajo el núcleo de adición de ruido hacia adelante.
Esto garantiza que cada bloque reciba una carga de trabajo equilibrada durante el entrenamiento.
Ningún bloque se queda sin señal y ninguno se ve sobrecargado.
El resultado es una utilización simétrica de los parámetros en toda la red, lo que promueve una convergencia estable y una equivalencia de extremo a extremo consistentemente fuerte.
Poniendo DiffusionBlocks a prueba
El marco se evaluó en una amplia familia de arquitecturas: transformers de visión, transformers de difusión, transformers autorregresivos, transformers de profundidad recurrente y transformers de difusión enmascarada.
En todos los casos, el entrenamiento con DiffusionBlocks igualó el rendimiento del entrenamiento estándar de extremo a extremo a la vez que permitió una verdadera independencia por bloques.
Los experimentos van mucho más allá de la clasificación a pequeña escala; demuestran por primera vez en la literatura de entrenamiento por bloques la viabilidad en tareas prácticas de IA generativa.
Esta es una refutación contundente de la creencia arraigada de que los métodos por bloques inevitablemente se quedan atrás respecto a la optimización conjunta.
La clave no residía en una pérdida local más ingeniosa, sino en una sólida reinterpretación teórica de la propia red.
De la eficiencia en el entrenamiento a las mejoras en la inferencia
El entrenamiento por bloques no solo ahorra memoria durante la optimización.
En los modelos de difusión, la inferencia solo requiere el bloque asignado al paso de eliminación de ruido actual, lo que reduce aún más la huella de memoria en el momento del despliegue.
Para los modelos de profundidad recurrente —donde un bloque de capas se aplica repetidamente— DiffusionBlocks sustituye el costoso entrenamiento iterativo por un procedimiento de una sola pasada, eliminando por completo la retropropagación a través del tiempo.
Este doble beneficio ilustra cómo una perspectiva centrada en la difusión transforma tanto el entrenamiento como la inferencia.
La estructura de la red ya no es una cadena monolítica, sino una pila modular de eliminadores de ruido especializados que se pueden intercambiar o reutilizar bajo demanda.
Una nueva dirección para el aprendizaje escalable
DiffusionBlocks aporta dos contribuciones fundamentales.
Primero, reinterpreta las redes basadas en transformers como instancias discretizadas de un proceso de difusión en tiempo continuo, lo que permite un entrenamiento de bloques genuinamente independiente, donde los gradientes solo se necesitan para un bloque a la vez.
Segundo, la estrategia de particionamiento por equiprobabilidad equilibra la utilización de parámetros en todos los bloques, garantizando que ninguna parte de la red quede subentrenada.
Al combinar el entrenamiento por bloques con una rigurosa teoría de la difusión, el método abre la puerta al entrenamiento de modelos cada vez más grandes sin chocar con el muro de memoria.
El código está disponible públicamente en https://github.com/SakanaAI/DiffusionBlocks, invitando a la comunidad a explorar y ampliar este nuevo paradigma.
