SPARKLING: Balancing Signal Preservation and Symmetry Breaking for Width-Progressive Learning

Resumo Técnico: SPARKLING

Declaração do Problema

O Aprendizado Progressivo (PL - Progressive Learning) visa reduzir o custo computacional do treinamento de Grandes Modelos de Linguagem (LLMs) ao expandir gradualmente a escala do modelo durante o treinamento, em vez de treinar a escala total do zero. Embora a pesquisa existente em PL tenha explorado extensivamente a expansão de profundidade (ex: empilhamento de camadas), a expansão de largura permanece significativamente subestudada. Métodos anteriores de expansão de largura são amplamente limitados aos estágios iniciais do treinamento (ex: <10–30% dos tokens), oferecendo vantagens computacionais negligenciáveis em relação ao treinamento do modelo de largura alvo do zero.

Expandir a largura durante o estágio intermediário do pré-treinamento é essencial para maximizar a economia computacional, mas apresenta desafios formidáveis. O artigo identifica dois problemas centrais que causam instabilidade no treinamento e degradação de desempenho neste regime:

1. Falha na Preservação do Sinal: A inicialização ingênua no ponto de expansão interrompe a distribuição estatística das ativações intermediárias (especificamente a escala da raiz do valor quadrático médio ou RMS). Esse descompasso altera as magnitudes do sinal camada a camada, desestabilizando a otimização mesmo que a perda imediata não apresente picos.
2. Bloqueio de Simetria (Symmetry Locking): A inicialização baseada em cópia, embora preserve a função de propagação (forward function), induz uma simetria de retropropagação (backward symmetry). Canais duplicados recebem gradientes idênticos e, se os estados do otimizador também forem simétricos, evoluem de forma idêntica. Isso impede que a nova capacidade se diversifique em características distintas, tornando os parâmetros expandidos funcionalmente redundantes. Este problema persiste através de diferentes famílias de otimizadores, incluindo otimizadores de elemento único (AdamW) e otimizadores de estilo espectral (Muon).

Metodologia: SPARKLING

Os autores propõem o SPARKLING (equilibrando a Preservação de Sinal e a manutenção da Simetria para aprendizado de largura progressiva - *balancing Signal Preservation And Symmetry Keeping for width-progressive LearnIng*), um framework projetado para equilibrar a preservação de sinal e a quebra de simetria durante a expansão de largura em estágio intermediário. O framework consiste em quatro mecanismos concretos:

1. Preservação de Sinal via Consistência de Escala RMS

Os autores argumentam que a verdadeira preservação de sinal requer a manutenção da distribuição estatística das representações ocultas, especificamente a escala RMS, em vez de apenas a preservação da função (continuidade da perda).

• Mecanismo: Eles derivam fatores de reescalonamento para os pesos durante a expansão para garantir que a escala RMS de saída permaneça invariante.
  • Expansão de Fan-out: Novos canais de saída são inicializados por cópia ou amostragem aleatória da mesma distribuição dos pesos existentes.
  • Expansão de Fan-in: Novos canais de entrada requerem um reescalonamento específico de pesos. Se as novas coordenadas forem aleatórias ou copiadas de um lado só, os pesos são escalonados por $\sqrt{d_{in}/d'_{in}}$. Se ambos os lados forem copiados (duplicando as dimensões existentes), os pesos são escalonados por $1/\sqrt{1+3c}$ (onde $c$ é a razão de cópia) para compensar a contribuição quadrática da covariância.
  • Expansão de RMSNorm: Novos parâmetros de RMSNorm são inicializados por cópia ou amostragem aleatória sem reescalonamento adicional, pois a contagem de parâmetros é negligenciável em relação às camadas lineares.

2. Quebra de Simetria via Intervenções Assimétricas

Para resolver o "bloqueio de simetria" onde os componentes copiados falham em se diversificar, o SPARKLING introduz intervenções direcionadas que atuam apenas nas dinâmicas de retropropagação (backward), enquanto preservam a função de propagação (forward).

• Reset Assimétrico do Estado do Otimizador: Em vez de copiar ou resetar globalmente os estados do otimizador (momento), o framework retém os estados do otimizador para os parâmetros originais ($W$), mas reseta os estados para os novos parâmetros introduzidos ($\tilde{W}$) para zero. Isso quebra o acoplamento de gradiente imediatamente, permitindo que os novos parâmetros divirjam.
• Reaquecimento Assimétrico da Taxa de Aprendizado (Asymmetric Learning Rate Re-warmup): No ponto de expansão, os parâmetros originais continuam no cronograma de taxa de aprendizado base. Os novos parâmetros passam por um reaquecimento começando da taxa de aprendizado atual ($\eta_e$) para um pico mais alto ($\hat{\eta}_{max} = \rho \cdot \eta_e$) ao longo de uma janela curta ($\tau_w$), antes de decair de volta para o mínimo compartilhado. Isso injeta uma assimetria de otimização controlada para estimular a diversificação de características sem perturbar globalmente o modelo pré-treinado bem adaptado.

Principais Contribuições

1. Identificação de Desafios de Estágio Intermediário: O artigo estabelece que o sucesso da expansão de largura em estágio intermediário depende de dois princípios complementares: Preservação de Sinal (estabilizando as estatísticas de ativação via consistência RMS) e Quebra de Simetria (resolvendo o acoplamento de gradiente na inicialização baseada em cópia).
2. O Framework SPARKLING: Uma implementação prática combinando escalonamento de preservação de RMS, inicialização baseada em cópia, reset assimétrico do estado do otimizador e reaquecimento assimétrico da taxa de aprendizado.
3. Validação Empírica: Experimentos extensos em arquiteturas densas e de Mistura de Especialistas (MoE), múltiplos eixos de largura (dimensão oculta e dimensão intermediária do especialista MoE) e famílias de otimizadores (AdamW e Muon).

Resultados

• Desempenho: Sob um orçamento fixo de tokens, o SPARKLING supera consistentemente o treinamento do modelo de escala total do zero em avaliações de downstream (ex: ARC, MMLU, HellaSwag) e iguala ou excede o desempenho de baselines ingênuos de preservação de função.
• Eficiência: O SPARKLING reduz significamente os custos de treinamento. Sob uma expansão de largura de 2×, ele alcança até 35% de redução nos FLOPs de treinamento e até 1.49× de aceleração de tempo de execução (wall-clock) em comparação com o treinamento do modelo expandido do zero.
• Convergência: O escalonamento de preservação de RMS melhora a convergência de estágio tardio em comparação com a inicialização não escalonada ingênua. As intervenções assimétricas (reset de estado e reaquecimento) são críticas para quebrar o bloqueio de simetria, sendo que a inicialização cópia-cópia combinada com essas intervenções atinge a melhor perda final entre todas as variantes.
• Generalidade: O método se generaliza para modelos densos, diferentes otimizadores (incluindo Muon) e expansões iterativas de múltiplos estágios (ex: 256→512→1024).

Significância e Alegações

O artigo afirma que o SPARKLING torna a expansão de largura em estágio intermediário praticamente viável, abordando a principal motivação do Aprendizado Progressivo: reduzir os custos de treinamento mantendo ou superando o desempenho dos modelos de largura alvo.

Os autores enfatizam que sua abordagem resolve a tensão entre continuidade de propagação (mantendo a função pré-treinada) e diversidade de retropropagação (permitindo que novos parâmetros aprendam características distintas). Ao desacoplar esses dois aspectos — preservando as estatísticas de sinal enquanto quebra a simetria nas dinâmicas de otimização — o SPARKLING permite a escala eficiente de LLMs. O trabalho sugere que o framework oferece um comportamento de escalonamento de computação-perda mais favorável, alcançando menor perda em orçamentos de computação equivalentes em comparação ao treinamento do zero.

O artigo conclui observando que, embora os resultados sejam promissores, trabalhos futuros são necessários para estabelecer um princípio unificado para a expansão simultânea de largura e profundidade e para investigar se a estratégia de preservação de RMS satisfaz a condição $\mu$P (maximal update parameterization) para transferibilidade de hiperparâmetros.