Qual é o modelo de IA mais poderoso que pode ser treinado em um notebook em 5 minutos?

• Em um MacBook Pro, foi treinado em 5 minutos um modelo transformer no estilo GPT com cerca de 1,8M de parâmetros usando aproximadamente 20M de tokens do TinyStories, alcançando cerca de 9,6 de perplexidade
• As principais limitações do treinamento em 5 minutos são o tamanho do modelo e a quantidade de tokens processáveis; modelos maiores convergem mais lentamente e ficam menos eficazes com menos dados
• Na otimização de desempenho, usar MPS, compilar/quantizar/acumular gradientes ou substituir o PyTorch é menos eficaz do que escolher um modelo pequeno
• Um conjunto de dados simples e consistente como o TinyStories tem impacto mais positivo no desempenho de modelos pequenos do que dados enciclopédicos
• A arquitetura transformer apresentou resultados superiores a abordagens de LSTM ou diffusion em condições de tamanho reduzido e tempo curto de treinamento

---

Visão geral

Este texto apresenta os resultados de um experimento para descobrir o modelo de linguagem de IA com melhor desempenho que pode ser treinado em 5 minutos em um notebook (MacBook Pro), além de trazer insights sobre a estratégia ideal de treinamento, escolha de dataset e arquitetura de modelo.

Resumo dos resultados do experimento

• Foi treinado um modelo transformer no estilo GPT com cerca de 1,8M de parâmetros em aproximadamente 20M de dados do TinyStories, registrando 9,6 de perplexidade
• Os exemplos gerados são curtos, mas no formato de histórias consistentes, com gramática inglesa em geral corretamente mantida
• Destaca-se que os resultados do modelo, em um nível prático dentro de 5 minutos, foram melhores do que o esperado

Contexto e limitações do experimento

• O experimento começou a partir da curiosidade pouco realista de treinar rapidamente um modelo poderoso em um ambiente de notebook
• Na prática, seria possível treinar modelos mais poderosos na nuvem com GPUs de alto desempenho (como H100), mas a condição limitante do experimento é o tempo: 5 minutos
• Quanto maior o modelo, mais lenta fica a velocidade de processamento de tokens, dificultando obter bons resultados em 5 minutos
  • Modelos pequenos demais (ex.: 10K parâmetros) não conseguem aprender complexidade suficiente
  • A faixa prática é de aproximadamente 1M a 2M de parâmetros

Otimização de throughput

• O uso de MPS (Metal Performance Shaders da Apple) é o mais eficaz
• Diversas otimizações matemáticas como torch.compile, float16 e MLX tiveram ganho de desempenho menor do que o esperado ou até pioraram o resultado
• O acúmulo de gradientes ajuda na gestão de memória, mas na prática provoca forte queda de velocidade
• Para ser eficiente, o modelo precisa conseguir fazer updates de pesos rapidamente na memória interna

Escolha do dataset

• Com uma quantidade limitada de tokens (cerca de 10~20M), ao usar primeiro dados simples em inglês da wiki, como Simple English Wikipedia, houve consistência gramatical, mas faltou consistência semântica
  • Como o conteúdo é centrado em nomes próprios e em enumerações de fatos que parecem forçadas, há limite para gerar conteúdo realmente significativo
• Ao usar o dataset TinyStories, os resultados foram muito mais consistentes e significativos, porque a estrutura narrativa é clara e a linguagem é simples
  • Como são histórias em nível de uma criança de 4 anos, até modelos pequenos conseguem aprender bem

Tokenizer e tokenização

• O treinamento do tokenizer não está incluído nos 5 minutos e, como a escala de dados é pequena, a necessidade de otimização é baixa
• Aprender tokens multibyte é mais fácil para o treinamento do modelo

Experimentos com arquitetura de modelo

• Foi usada a arquitetura transformer (estilo GPT-2)

  • Ajuste de hiperparâmetros como 2~3 camadas, funções de ativação como SwiGLU e positional embedding
  • O LSTM teve desempenho próximo, mas o transformer foi superior em termos de perplexidade
  • Dropout e mixture-of-experts são ineficientes nesse porte pequeno
  • Curriculum learning teve pouco efeito porque o tempo de treinamento era curto demais

• Tentativa com modelo diffusion (D3PM)

  • Como linguagem natural é composta por tokens discretos, o processo de difusão acabou gerando apenas tokens aleatórios sem sentido e fracassou
  • É mais difícil formar rapidamente estruturas de frase do que com transformer ou LSTM

Relação entre tamanho do modelo e throughput em tokens/segundo

• Modelos com 1M~2M de parâmetros são o sweet spot mais ideal
  • Se forem grandes demais, não convergem em 5 minutos; se forem pequenos demais, atingem o limite de desempenho logo no começo do treinamento
• A Chinchilla scaling law foi, em geral, consistente com os resultados do experimento
  • O tamanho ideal do modelo é o total de tokens de treinamento dividido por 20, e isso também foi confirmado neste experimento

Conclusão e implicações

• Mesmo com tempo muito curto e hardware pequeno, é possível treinar um modelo de storytelling consistente
• Treinar por 5 minutos não é adequado para desenvolver modelos poderosos, mas tem valor para experimentos de design de modelos pequenos e ultraleves, além de otimização de hardware e arquitetura
• Com a evolução futura das GPUs para notebook e das estruturas de modelo, existe potencial para avanços no desempenho de modelos treináveis em apenas alguns minutos