Mixture of Training: Recombining Small-Scale Scaffolded Pretraining Runs into a Larger Language Model

Resumo Técnico: Mixture of Training (MoT)

Definição do Problema
O pré-treinamento de grandes modelos de linguagem (LLMs) é atualmente organizado como uma única execução de otimização monolítica de ponta a ponta. Esta arquitetura cria um gargalo de escalonamento onde todas as camadas devem ser treinadas simultaneamente, tornando o processo frágil a falhas (onde um único erro interrompe toda a execução) e inflexível para pesquisa iterativa (onde melhorias exigem o reinício ou a continuação de um sistema massivo e acoplado). O artigo questiona se o pré-treinamento pode ser decomposto em tarefas menores, treináveis de forma independente, que possam posteriormente ser recompostas em um modelo maior e coerente sem sacrificar a qualidade.

Metodologia: Mixture of Training (MoT)
Os autores introduzem o Mixture of Training (MoT), um procedimento de pré-treinamento modular e estruturado projetado para particionar um Transformer alvo em blocos de camadas contíguos. O método opera em três estágios:

1. Estágio 0 (Preparação): Um alinhador ($A$) pré-treinado é selecionado ou preparado. O Transformer alvo ($F$) e o alinhador são ambos particionados em $K$ blocos de camadas contíguos. O alinhador é compatível em forma com o alvo, compartilhando largura global, dimensionalidade de cabeça de atenção, largura de feed-forward, embeddings de tokens e cabeças de saída.
2. Estágio 1 (Treinamento Estruturado): O modelo alvo é decomposto em $K$ submodelos ($f_1, \dots, f_K$). Para cada submodelo $f_i$, uma rede estruturada $S_i$ é construída inserindo $f_i$ na estrutura do alinhador congelado:
   $$S_i = a_K \circ \dots \circ a_{i+1} \circ f_i \circ a_{i-1} \circ \dots \circ a_1$$
   Nesta configuração, apenas o bloco alvo $f_i$ é treinável; todos os cortes do alinhador estão congelados. Cada scaffold é otimizado independentemente usando o padrão de perda de previsão do próximo token (next-token prediction loss). Crucialmente, nenhum gradiente é trocado entre os blocos alvo durante esta fase, mas todos aprendem dentro do contexto representacional estável fornecido pelo alinhador compartilhado.
3. Estágio 2 (Recomposição e Adaptação): O alinhador é descartado. Os blocos alvo treinados são recompostos no modelo completo $\hat{F} = f_K \circ \dots \circ f_1$. Este modelo de "composição a frio" (cold-composed) pode ser avaliado diretamente ou pode passar por uma breve passagem de adaptação de ponta a ponta para resolver quaisquer discrepâncias representacionais remanescentes.

Principais Contribuições

1. Treinamento Estruturado Modular: Um procedimento que treina fatias de profundidade alvo de forma independente enquanto preserva uma interface representacional compartilhada via um alinhador congelado, permitindo a recomposição sem a necessidade de costurar camadas post-hoc.
2. Prova de Mecanismo: Um estudo em pequena escala demonstrando que fatias de profundidade treinadas independentemente podem ser recompostas com sucesso em um modelo de linguagem de 1,3 bilhão de parâmetros utilizável.
3. Contabilidade de Computação e Escalonamento: Uma análise detalhada das compensações entre computação (FLOPs), exposição de tokens e o caminho crítico ideal equivalente de camadas. Isso inclui o tratamento explícito do custo do alinhador, distinguindo entre custos de treinamento "diretos" e custos "totalmente carregados" (atribuindo o custo total do alinhador a uma única execução) versus custos amortizados em múltiplas execuções.
4. Estudos de Ablação: Evidência empírica mostrando a necessidade do alinhador, o impacto de fluxos de dados distintos e a compensação entre qualidade e eficiência associada ao aumento do número de divisões ($K$).

Resultos Experimentais
Os autores avaliaram o MoT em um modelo estilo Gemma de 12 camadas e 1,3 bilhão de parâmetros treinado no dataset C4.

• Paridade de Qualidade: Um cronograma de "paridade de qualidade", que estendeu o treinamento do submodelo e incluiu uma passagem de adaptação de 30k passos, alcançou a mesma perplexidade (15,0) que o baseline monolítico.
• Compensações de Computação:
  • O cronograma de paridade de qualidade processou mais tokens agregados (47,1B vs. 33,6B) e teve um caminho crítico ideal mais curto (potencial de aceleração de 1,7x), mas exigiu mais FLOPs totais (285,0 EFLOPs vs. 268,4 EFLOPs) quando o custo do alinhador foi totalmente carregado para uma única execução.
  • No entanto, se o alinhador for reutilizado em $R$ execuções independentes, o custo efetivo por execução cai. O artigo observa que para $R \ge 3$, o custo efetivo fica abaixo do baseline monolítico.
  • Um cronograma de menor computação (MoT + 15k de adaptação) atingiu uma perplexidade de 15,9 em 189,4 EFLOPs (totalmente carregados), o que está abaixo do orçamento do baseline, embora nenhum baseline monolítico de igual computação tenha sido fornecido para comparação direta neste orçamento específico.
• Ablações:
  • Necessidade do Alinhador: Sem o alinhador, a qualidade da composição a frio colapsou (PPL 38,9), confirmando que o alinhador evita o desajuste de interface.
  • Fluxos de Dados: Fluxos de dados distintos para os submodelos melhoraram a qualidade da composição a frio (PPL 19,3) em comparação com dados compartilhados (PPL 20,3) quando um alinhador estava presente.
  • Contagem de Divisões: Aumentar as divisões de $K=2$ para $K=4$ reduziu a computação, mas degradou a qualidade da composição a frio (PPL 24,8), destacando uma compensação entre paralelismo e complexidade de interface.

Significância e Alegações
O artigo posiciona o MoT não como um substituto geral para o pré-treinamento monolítico, mas como um framework de pequena escala para estudar se sub-execuções estruturadas podem atuar como unidades de treinamento reutilizáveis.

Os autores afirmam três conclusões primárias:

1. Submodelos treinados independentemente podem ser recompostos em um modelo de linguagem coerente se treinados dentro de um scaffold de alinhador compartilhado; sem ele, a composição a frio falha.
2. O desajuste entre blocos treinados independentemente é amplamente recuperável; uma curta passagem de adaptação fecha a maior parte da lacuna, e um cronograma mais longo pode alcançar a paridade com o treinamento monolítico.
3. O MoT altera a forma de engenharia do pré-treinamento ao criar tarefas menores, independentemente agendáveis, reiniciáveis e abláveis. Isso é particularmente adequado para pesquisa de treinamento em pequena escala, tornando-o potencialmente mais barato, mais reprodutível e mais fácil de iterar.

Os autores declaram explicitamente que a evidência atual demonstra a composibilidade modular e identifica um regime de escalonamento de reutilização amortizada, em vez de estabelecer uma vantagem uniforme de computação ou de hardware. O estudo é intencionalmente limitado em escopo (uma família de modelos, um dataset, métricas de apenas perplexidade) para servir como uma prova de mecanismo para um espaço de design concreto de sub-execuções estruturadas reutilizáveis.