Resumo Técnico: Pré-treinamento Eficiente com Superposição de Tokens

Declaração do Problema

O pré-treinamento de Modelos de Linguagem de Grande Escala (LLMs) está se tornando cada vez mais proibitivamente caro e ineficiente em escala. Os paradigmas atuais frequentemente exigem modificações complexas e invasivas em estratégias de paralelismo, otimizadores, tokenizadores ou arquiteturas de modelos para alcançar alta vazão de dados. Além disso, embora pesquisas recentes tenham se concentrado em melhorar a eficiência no momento da inferência ou em desacoplar a eficiência do treinamento da da inferência, há uma falta de métodos que aumentem significativamente a vazão durante o treinamento sem alterar a arquitetura final do modelo ou as dinâmicas de inferência. As abordagens existentes frequentemente confundem eficiência de treinamento e de inferência ou dependem de perdas auxiliares que adicionam parâmetros e hiperparâmetros sem necessariamente aumentar a taxa bruta de processamento de tokens por FLOP.

Metodologia: Treinamento com Superposição de Tokens (TST)

Os autores propõem o Treinamento com Superposição de Tokens (TST), um método "plug-and-play" projetado para melhorar a vazão de dados por FLOP durante o pré-treinamento sem modificar o paralelismo, o otimizador, o tokenizador, os dados ou a arquitetura do modelo. O TST opera em duas fases distintas:

Fase 1: Fase de Superposição

Nesta fase, o modelo processa um "saco" de tokens contíguos como uma única unidade latente, denominada "s-token".

• Superposição de Entrada: Uma sequência contígua de $s$ tokens é segmentada em sacos não sobrepostos. Os embeddings desses $s$ tokens são mediados para criar um único embedding latente de token superposto. Isso reduz o comprimento da sequência por um fator de $s$ (de $L$ para $\lfloor L/s \rfloor$).
• Mecanismo de Vazão: Para garantir uma comparação justa com o treinamento de base (FLOPs iguais), o comprimento da sequência de dados $L$ é aumentado por um fator de $s$ durante esta fase. Consequentemente, o modelo ingere $s$ vezes mais tokens de dados pelo mesmo custo computacional por etapa.
• Superposição de Saída: Em vez de prever um único próximo token, o modelo prevê o próximo "saco" de $s$ tokens. A perda padrão de entropia cruzada (CE) one-hot é substituída por uma perda de Entropia Cruzada Multi-Hot (MCE). O objetivo MCE trata o saco alvo como um rótulo multi-hot, empurrando a probabilidade de cada token válido no saco para $1/s$. A perda é formulada como a média das perdas de CE padrão para cada token no saco, menos um termo constante ($\log |y|$) que não afeta os gradientes.
• Causalidade: Para preservar a causalidade, os rótulos padrão de previsão do próximo token são deslocados para a esquerda por $s-1$ antes de serem divididos em sacos não sobrepostos.

Fase 2: Fase de Recuperação

Um modelo treinado exclusivamente com TST produz saídas sem sentido porque perde a ordenação dos tokens dentro dos sacos e carece de capacidades padrão de amostragem autoregressiva. Para remediar isso:

• O treinamento reverte para o objetivo padrão de previsão causal de próximo token.
• O modelo retoma o treinamento a partir do ponto de verificação salvo no final da fase de superposição.
• O código TST é totalmente removido para garantir que a arquitetura final do modelo e as dinâmicas de inferência sejam idênticas à base.
• A proporção de etapas gastas na fase de superposição é denotada como $r$ (por exemplo, $r=0.3$ significa que 30% das etapas são TST e 70% são padrão).

Principais Contribuições

1. Novo Paradigma de Treinamento: Introdução do TST, que aumenta a vazão de tokens em $s$ vezes durante um regime de treinamento específico sem alterar a arquitetura final do modelo ou o comportamento de inferência.
2. Design de Duas Fases: Um mecanismo simples combinando superposição de entrada (medição de embeddings) e superposição de saída (previsão multi-hot) seguido por uma fase de recuperação para alinhar o modelo de volta ao comportamento autoregressivo padrão.
3. Ortogonalidade aos Métodos Existentes: Os autores posicionam o TST como ortogonal aos métodos de perda auxiliar (como Previsão de Múltiplos Tokens). Enquanto a MTP adiciona parâmetros e direciona tokens futuros sem aumentar tokens por FLOP, o TST aumenta estritamente a vazão e deixa a arquitetura intacta.
4. Validação Empírica: Avaliação extensa em escalas de modelos (270M, 600M, 3B e 10B Mixture-of-Experts) demonstrando robustez em diferentes configurações de hiperparâmetros.

Resultados

O artigo relata as seguintes descobertas baseadas em experimentos usando os conjuntos de dados DCLM e FineWeb-Edu:

• Redução de Perda: O TST supera consistentemente os modelos de base em termos de perda final de treinamento sob restrições de FLOPs iguais.
• Desempenho a Montante: Modelos treinados com TST mostram desempenho aprimorado em benchmarks a montante (HellaSwag, ARC-E/C, MMLU, etc.) em comparação com bases treinadas pelo mesmo número de etapas ou mesmo custo computacional.
• Ganhos de Eficiência:
  • Sob configurações de perda igual, o TST resulta em até uma redução de 2,5x no tempo total de pré-treinamento na escala de 10B parâmetros.
  • No experimento 10B A1B (10B total, 1B ativo) MoE, o modelo TST viu 2T tokens de dados em comparação com os 1,05T tokens da base, enquanto alcançava uma perda final menor (2,236 vs. 2,252) em significativamente menos horas B200 (4768 vs. 12311).
• Sensibilidade a Hiperparâmetros: O método é robusto dentro de uma faixa razoável de tamanhos de sacos de superposição ($s \in [4, 8]$) e proporções de etapas ($r \in [0.2, 0.4]$).
• Estudos de Ablação:
  • Entrada vs. Saída: Tanto a superposição apenas de entrada quanto apenas de saída melhoram o desempenho sobre a base, mas a combinação de ambas produz os melhores resultados, sugerindo que são mecanismos ortogonais.
  • Alinhamento de Representação: Re-inicializar aleatoriamente os embeddings de entrada e as cabeças LM de saída entre as fases de superposição e recuperação elimina completamente os ganhos do TST, indicando que manter o alinhamento de representação entre as fases é crítico.

Significado e Alegações

O artigo afirma que o TST oferece uma melhoria significativa na eficiência do pré-treinamento ao maximizar a vazão durante o treinamento independentemente da arquitetura de inferência do modelo. Os autores argumentam que essa abordagem desafia o paradigma monolítico de pré-treinamento, demonstrando que um cronograma de granularidade "de grosso para fino" é um ingrediente reutilizável para pré-treinamento eficiente.

As principais alegações incluem:

• Redução de Custos: Alcançar a mesma qualidade de modelo (perda) a aproximadamente metade do custo computacional.
• Simplicidade: O método não requer mudanças arquiteturais complexas, novos otimizadores ou parâmetros adicionais, tornando-o uma solução prática "plug-and-play".
• Desacoplamento de Eficiência: Desacopla com sucesso a eficiência do tempo de treinamento da eficiência do tempo de inferência, abordando uma lacuna na literatura atual onde muitos métodos confundem os dois.
• Escalabilidade: O método escala efetivamente de modelos densos pequenos (270M) até grandes modelos Mixture-of-Experts (10B), sugerindo ampla aplicabilidade.

Os autores permanecem modestos quanto aos mecanismos subjacentes, observando que, embora hipoteticamente a superposição de entrada atue como uma forma de "pré-pré-treinamento" ou regularização implícita da geometria de embeddings, mais trabalho de interpretabilidade é necessário. Eles também reconhecem limitações, como a suposição de que o pré-treinamento é atualmente limitado por computação e não por dados, e a falta de ablações em grande escala devido a restrições de computação.