Micro Diffusion - um pequeno modelo de Diffusion para estudo

Olá, há cerca de um mês implementei e publiquei um modelo de Diffusion para texto do zero. Recentemente, vendo o GeekNews, isso me veio à cabeça e só agora estou postando aqui.

---

O que me levou a criar isso

Ao ver o MicroGPT do Andrej Karpathy, fiquei com a impressão de que "dá para explicar a essência do GPT com um código tão curto". Como eu também queria estudar Diffusion, comecei este projeto educacional pensando: "seria interessante criar um código para entender Diffusion da mesma forma".

---

AR vs Diffusion: qual é a diferença?

Na era dos LLMs, a geração de texto é quase toda feita no modo Autoregressive (AR). É um método que prevê os tokens um a um, da esquerda para a direita.

Discrete Diffusion é o oposto. Ele pega a sequência inteira de uma vez e a restaura gradualmente a partir do ruído (mascaramento).

Usando o nome "emma" como exemplo:  
  
Forward (treinamento - adicionando ruído):  
  t=0 :  e m m a   ← original  
  t=25:  e _ m a   ← parcialmente mascarado  
  t=50:  _ _ m _   ← mais mascarado  
  t=100: _ _ _ _   ← completamente mascarado  
  
Reverse (geração - removendo ruído):  
  t=100: _ _ _ _   ← começa de um estado vazio  
  t=75:  _ m _ _   ← restaura primeiro as posições de maior confiança  
  t=50:  e m _ a  
  t=0 :  e m m a   ← concluído  

Se AR é como "escrever uma palavra letra por letra", Diffusion é mais próximo de "resolver palavras cruzadas".

---

Estrutura da implementação

Há três versões, então você pode escolher de acordo com o nível de dificuldade.

• train_minimal.py — MLP de 2 camadas / usa apenas NumPy (o mais simples)
• train_pure.py — MLP de 3 camadas + skip connection / usa apenas NumPy
• train.py — Transformer de 4 camadas / usa PyTorch

As três versões compartilham o mesmo loop de diffusion. Só a arquitetura do denoiser é diferente.
Os dados de treinamento são 32.000 nomes em inglês, e os parâmetros ficam na faixa de 170K~239K.

Você pode executar na hora:

pip install numpy  # versão mínima  
python train_minimal.py  

---

Comentários, feedback ou PRs são bem-vindos!