Micro Diffusion - un pequeño modelo de Diffusion para estudiar

Hola, hace como un mes implementé desde cero un modelo de Diffusion para texto y lo publiqué. Hace poco, viendo GeekNews, se me vino a la mente y recién ahora lo comparto aquí.

---

Cómo surgió

Al ver MicroGPT de Andrej Karpathy, me quedó la impresión de que "se puede explicar el núcleo de GPT con un código tan corto". Justo también quería estudiar Diffusion, así que empecé este proyecto educativo pensando: "sería divertido hacer código con el que se pueda entender Diffusion de la misma manera".

---

AR vs Diffusion: ¿en qué se diferencian?

En la era de los LLM, casi toda la generación de texto usa el enfoque Autoregressive (AR). Es decir, se predicen tokens uno por uno de izquierda a derecha.

Discrete Diffusion hace lo contrario. Toma toda la secuencia de una vez y la restaura gradualmente a partir del ruido (enmascaramiento).

Tomando el nombre "emma" como ejemplo:  
  
Forward (entrenamiento - agregar ruido):  
  t=0 :  e m m a   ← original  
  t=25:  e _ m a   ← parcialmente enmascarado  
  t=50:  _ _ m _   ← más enmascarado  
  t=100: _ _ _ _   ← completamente enmascarado  
  
Reverse (generación - eliminar ruido):  
  t=100: _ _ _ _   ← empieza desde un estado vacío  
  t=75:  _ m _ _   ← restaura primero las posiciones con mayor certeza  
  t=50:  e m _ a  
  t=0 :  e m m a   ← completado  

Si AR es como "escribir una palabra letra por letra", Diffusion se parece más a "resolver un crucigrama".

---

Estructura de la implementación

Hay tres versiones, así que puedes elegir cuál ver según el nivel de dificultad.

• train_minimal.py — MLP de 2 capas / usa solo NumPy (la más simple)
• train_pure.py — MLP de 3 capas + skip connection / usa solo NumPy
• train.py — Transformer de 4 capas / usa PyTorch

Las tres versiones comparten el mismo loop de diffusion. Solo cambia la arquitectura del denoiser.
Los datos de entrenamiento son 32,000 nombres en inglés, y los parámetros están en el rango de 170K~239K.

Puedes ejecutarlo de inmediato:

pip install numpy  # versión mínima  
python train_minimal.py  

---

¡Se agradece cualquier feedback o PR!