Resumo Técnico: Além dos Custos Irrecuperáveis: Aumentando a Eficiência do Pré-treinamento de LLMs via Crescimento Ortogonal de Mistura de Especialistas

Declaração do Problema

As demandas computacionais para o pré-treinamento de Modelos de Linguagem de Grande Escala (LLMs) estão aumentando, mas os checkpoints pré-treinados existentes frequentemente permanecem subutilizados devido a limitações arquitetônicas. As estratégias atuais de crescimento de modelos geralmente focam na expansão dos modelos durante os estágios iniciais do treinamento, falhando em capitalizar os substanciais "custos irrecuperáveis" (investimento computacional anterior) de checkpoints bem convergidos. Além disso, a aplicação do crescimento de modelos a arquiteturas de Mistura de Especialistas (MoE), que são cada vez mais dominantes em LLMs de grande escala, não foi estudada sistematicamente. O desafio central é determinar como "reciclar" efetivamente modelos MoE convergidos para criar modelos maiores e mais capazes, sem descartar o conhecimento aprendido e os parâmetros otimizados do modelo base.

Metodologia

Os autores propõem um framework de "crescimento ortogonal" projetado para expandir modelos MoE convergidos ao longo de duas dimensões independentes: Profundidade (aumento de camadas) e Largura (aumento do número de especialistas).

1. Crescimento de Profundidade: Interposicional vs. Empilhamento

O artigo contrasta dois métodos para expandir a profundidade do modelo:

• Empilhamento: Concatenação da sequência original de camadas $k$ vezes ($l_1, \dots, l_n, l_1, \dots, l_n$).
• Crescimento Interposicional: Duplicação de cada camada no local ($l_1, l_1, \dots, l_2, l_2, \dots$).

Insight Chave: A análise de LLMs pré-treinados revela uma distribuição característica de norma de pesos por camada: as normas são pequenas nas camadas iniciais, aumentam monotonicamente através do meio e diminuem ligeiramente perto do final. Este padrão é uma assinatura estrutural de uma convergência saudável.

• Empilhamento interrompe essa tendência monotônica ao colocar normas de camadas iniciais após normas de camadas tardias, criando descontinuidades abruptas que exigem que o modelo "repare" durante o treinamento contínuo.
• Crescimento Interposicional preserva a trajetória de norma de pesos aprendida ao duplicar camadas no local, mantendo a progressão estrutural suave.
• Limiar: A superioridade do crescimento interposicional torna-se crítica assim que o modelo base foi treinado além dos FLOPs ótimos de computação definidos pelas leis de escala Chinchilla (aproximadamente $1\times F_c$).

2. Crescimento de Largura: Duplicação de Especialistas Ruidosa

Para expandir a largura do modelo, os autores dobram tanto o número total de especialistas ($E \to 2E$) quanto o número de especialistas ativados ($k \to 2k$).

• Estratégia de Inicialização: Novos especialistas são inicializados copiando os pesos dos especialistas existentes e adicionando uma pequena quantidade de ruído Gaussiano ($\epsilon \sim \mathcal{N}(0, (\alpha\sigma_{orig})^2)$ com $\alpha \approx 0.01$).
• Racional: A duplicação exata cria simetria perfeita, impedindo a divergência funcional. Uma pequena perturbação quebra essa simetria o suficiente para permitir que o roteador diferencie os especialistas durante o treinamento contínuo, enquanto preserva as representações aprendidas. Isso contrasta com o "upcycling" a partir de modelos densos, que frequentemente requer níveis altos de ruído ($\ge 50\%$) e inicialização aleatória do roteador.

3. Ortogonalidade

O artigo demonstra que o crescimento de profundidade e largura são ortogonais.

• Disjunção Estrutural: O crescimento de profundidade adiciona novas camadas; o crescimento de largura adiciona novos especialistas dentro das camadas existentes. Eles operam em conjuntos de parâmetros não sobrepostos.
• Dinâmicas de Otimização: A análise de gradientes mostra que as direções de otimização para modelos crescidos em profundidade e em largura são quase ortogonais à direção pré-crescimento e entre si.
• Independência de Ordem: O desempenho final é invariante à sequência de expansão (profundidade-então-largura vs. largura-então-profundidade), confirmando que as estratégias podem ser integradas flexivelmente.

Resultados Chave

O framework foi validado através de experimentos extensivos em modelos variando de 3B a 70B parâmetros, utilizando um conjunto de dados de 1 trilhão de tokens.

• Eficácia do Crescimento de Profundidade: Em um modelo MoE de 3B crescido para 6B, o método interposicional superou consistentemente o empilhamento tanto na perda de treinamento quanto na precisão downstream. A lacuna de desempenho aumentou à medida que o modelo base se tornava mais convergido.
• Correlação com Custos Irrecuperáveis: Há uma forte correlação positiva entre a quantidade de "FLOPs irrecuperáveis" (etapas de treinamento anteriores ao crescimento) e a precisão final do modelo crescido. Modelos crescidos a partir de checkpoints posteriores e mais convergidos alcançaram desempenho final mais alto do que aqueles crescidos a partir de checkpoints anteriores, mesmo com orçamentos computacionais adicionais idênticos.
• Comparação com Treinamento do Zero:
  • Sob um orçamento computacional adicional fixo, o crescimento de modelos superou significativamente o treinamento de um modelo grande do zero (por exemplo, uma melhoria de 10,6% na precisão no experimento de 70B).
  • Sob um orçamento computacional total fixo (irrecuperável + adicional), o crescimento de modelos produziu resultados comparáveis ou ligeiramente superiores ao treinamento do zero, desde que o crescimento fosse iniciado a partir de um checkpoint com custo irrecuperável suficiente.
• Escalabilidade: O método escalou com sucesso um modelo MoE de 17B (treinado em ~600B tokens) para 70B parâmetros (treinado em um total de 1T de tokens). O modelo final de 70B alcançou uma precisão média de 64,17, superando a linha de base de 17B em 5,62 pontos e a linha de base de 70B treinada do zero em um ganho relativo de 10,6% sob o mesmo orçamento extra de FLOPs.

Significado e Alegações

O artigo alega fornecer um modelo de custo efetivo para o desenvolvimento sustentável de LLMs de grande escala. Seu significado principal reside em:

1. Reciclagem de Custos Irrecuperáveis: Demonstrar que checkpoints pré-treinados, frequentemente vistos como ativos intermediários descartáveis, podem ser sistematicamente aproveitados para criar modelos maiores, preservando assim o valor dos investimentos computacionais anteriores.
2. Estratégias de Crescimento Ótimas: Estabelecer que, para modelos bem convergidos, a cópia interposicional de camadas e a duplicação de especialistas com ruído mínimo são superiores ao empilhamento tradicional e à inicialização aleatória.
3. Escala Ortogonal: Provar que profundidade e largura podem ser expandidas independentemente e em qualquer ordem sem interferência, oferecendo flexibilidade em cenários de escala com recursos restritos.
4. Eficiência: Mostrar que o crescimento pode alcançar ganhos de desempenho que superam o treinamento do zero sob orçamentos computacionais extras idênticos, efetivamente reduzindo a barreira para o pré-treinamento de modelos na fronteira.

Os autores enfatizam que este trabalho não introduz novos riscos sociais além dos inerentes aos LLMs, mas visa reduzir a pegada ambiental do desenvolvimento de modelos ao melhorar a eficiência do treinamento.