Resumo Técnico: Teoria de Cronogramas Ótimos de Taxa de Aprendizado e Leis de Escalonamento para um Modelo de Características Aleatórias

Declaração do Problema

O treinamento de modelos de aprendizado profundo requer a seleção de hiperparâmetros, como taxa de aprendizado (LR), tamanho de lote e horizonte de treinamento. Embora o "tentativa e erro" empírico seja comum, há uma falta de compreensão teórica sobre como os hiperparâmetros ótimos escalam com o horizonte de treinamento $T$. Estratégias existentes para transferir hiperparâmetros entre tamanhos de modelo frequentemente falham ao serem transferidas entre diferentes horizontes de treinamento. Além disso, embora vários cronogramas de LR (por exemplo, decaimento linear, aquecimento-estabilidade-decaimento, annealing cosseno) sejam utilizados na prática, as condições teóricas que ditam qual cronograma é ótimo para uma dada tarefa permanecem obscuras. Este artigo aborda a necessidade de um framework principiado para derivar cronogramas de LR ótimos que dependam do horizonte de treinamento e da estrutura específica da tarefa de aprendizado.

Metodologia

Os autores empregam um modelo de características aleatórias de lei de potência treinado com Descida de Gradiente Estocástica (SGD) para derivar cronogramas ótimos de hiperparâmetros. O modelo assume dados de entrada $x$ com uma função alvo $y(x) = w^* \cdot \psi(x) + \sigma_0 z$, onde as características $\psi(x)$ exibem decaimento espectral de lei de potência ($\lambda_k \sim k^{-b}$) e os pesos do professor escalam como $(w^*_k)^2 \lambda_k \sim k^{-a}$.

A metodologia central envolve teoria de controle ótimo:

1. Formulação da Dinâmica: A perda de teste esperada $L_t$ é derivada como uma função do cronograma de taxa de aprendizado $\{\eta_t\}$ e do cronograma de tamanho de lote $\{m_t\}$. A dinâmica da perda é governada por uma relação de recorrência linear para erros de modo $c_{t,k}$.
2. Objetivo de Otimização: O objetivo é minimizar a perda de teste final $L_T$ em um horizonte fixo $T$. Os autores formulam três problemas de otimização:
   • LR constante ótima.
   • Cronograma de LR ótimo (variável no tempo $\eta_t$).
   • Cronograma conjunto ótimo de LR e tamanho de lote.
3. Abordagens de Solução:
   • Numérica: Utilizando um método de disparo único (implementado no CasADi) para resolver o problema de controle ótimo de tempo discreto para $T$ finito.
   • Analítica: Derivando aproximações de tempo contínuo para $T$ grande usando o teorema de Noether e análise assintótica do funcional de perda. Isso envolve minimizar um funcional da taxa de aprendizado integrada $\chi(t) = \int_t^T \eta(s) ds$.

Principais Contribuições e Resultados

1. Identificação de Duas Fases Distintas

A análise revela dois regimes distintos para cronogramas ótimos de taxa de aprendizado, determinados pela relação entre os expoentes espectrais $a$ (dificuldade da tarefa) e $b$ (complexidade da característica):

• Fase Fácil ($a > b$):

  • O cronograma ótimo segue uma forma de decaimento polinomial: $\eta^*_T(t) \sim T^{b/a - 1} f(t/T)$, onde $f(z) \propto (1-z)^{2b-1}$.
  • A taxa de aprendizado decai gradualmente para equilibrar a redução de viés e variância ao longo da trajetória de treinamento.
  • O expoente de escalonamento da perda é $\zeta_T = (a-1)/a$, correspondendo à taxa minimax-ótima para tamanhos de lote fixos.

• Fase Difícil ($a < b$):

  • O cronograma ótimo assemelha-se ao Aquecimento-Estabilidade-Decaimento (WSD). Mantém uma taxa de aprendizado máxima constante ($\eta_{max}$) para a maioria das etapas de treinamento, seguida por uma fase de annealing rápida sobre uma fração vanescente do tempo total ($1 - t_s/T \sim T^{-(b-a)/(2b-1)}$).
  • Esta estratégia prioriza a supressão de viés via uma LR constante máxima, comprimindo a redução de variância nas etapas finais.
  • O expoente de escalonamento da perda melhora para $\zeta_T = (a-1)/b$, superando estratégias de LR constante.

2. Superioridade sobre Benchmarks

Os cronogramas ótimos derivados superam estritamente os benchmarks padrão:

• LR Constante: O cronograma ótimo alcança expoentes de escalonamento melhores em ambas as fases.
• Cronogramas de Lei de Potência: Mesmo ao otimizar os expoentes de cronogramas de lei de potência ($\eta \sim T^{-\xi} t^{-\delta}$), eles recuperam assintoticamente o mesmo comportamento de escalonamento da linha de base de LR constante ótima, falhando em corresponder ao desempenho dos cronogramas ótimos dependentes do horizonte derivados via teoria de controle.

3. Otimização Conjunta de Taxa de Aprendizado e Tamanho de Lote

Os autores estendem a análise para otimizar conjuntamente $\eta(t)$ e $m(t)$ sob um orçamento de dados fixo $B_{tot}$:

• Fase Fácil: Um "ramp de lote" (aumento do tamanho de lote ao longo do tempo) combinado com uma taxa de aprendizado máxima constante reduz o tempo de relógio. A perda escala como $T^{-(a-1)/b}$, melhorando o expoente de relógio comparado a tamanhos de lote fixos, enquanto o expoente de orçamento de amostra permanece $(a-1)/a$.
• Fase Difícil: A otimização conjunta produz os mesmos expoentes de escalonamento que o ótimo de lote fixo, com o tamanho de lote permanecendo limitado.

4. Cronogramas Ótimos de Momento

Para tarefas com momento, os autores otimizam o parâmetro de momento $\beta(t)$ juntamente com $\eta(t)$.

• Na Fase Fácil, o cronograma de momento ótimo produz leis de escalonamento similares ao SGD ótimo.
• Na Fase Difícil, otimizar o cronograma de momento proporciona uma melhoria no expoente de escalonamento além do alcançável com SGD sozinho. A dinâmica de momento ótima mostra uma diminuição transitória inicial seguida por um aumento à medida que $t/T \to 1$.

5. Validação Empírica em Redes Profundas

As fases teóricas são validadas em benchmarks de aprendizado profundo:

• ResNet no CIFAR-5M: Exibe comportamento de Fase Difícil. A taxa de aprendizado base ótima para cronogramas de annealing é aproximadamente constante através dos horizontes, consistente com o perfil WSD.
• Transformadores Estilo GPT-2 no C4: Exibem comportamento de Fase Fácil. A taxa de aprendizado base ótima desloca-se (diminui) com o horizonte $T$, mesmo sob annealing, consistente com o perfil de decaimento polinomial.

Significado e Alegações

O artigo alega fornecer uma base teórica para entender por que cronogramas de taxa de aprendizado variam entre diferentes tarefas e horizontes de treinamento. Ao utilizar a teoria de controle ótimo em um modelo de características aleatórias solúvel, os autores demonstram que:

1. Dependência do Horizonte: Taxas de aprendizado ótimas não são automaticamente transferíveis entre horizontes de treinamento; o cronograma deve ser explicitamente projetado com base em $T$.
2. Estrutura da Tarefa: A "dificuldade" de uma tarefa (caracterizada pelos expoentes espectrais $a$ e $b$) dita se uma tarefa cai na fase fácil ou difícil, determinando a forma do cronograma ótimo (decaimento polinomial vs. WSD).
3. Eficiência: Cronogramas dependentes do horizonte e otimização conjunta do tamanho de lote podem melhorar significativamente as leis de escalonamento e reduzir o tempo de treinamento de relógio comparado a heurísticas padrão.

Os autores reconhecem limitações, observando que sua teoria atualmente se aplica a modelos lineares no regime preguiçoso/kernel e ainda não deriva analiticamente o expoente específico melhorado para momento na fase difícil ou considera otimizadores sofisticados como Adam ou métodos pré-condicionados. No entanto, o trabalho oferece uma explicação principiada para a eficácia de cronogramas como WSD e sugere que o ajuste empírico de taxas de aprendizado base através de horizontes deve levar em conta as propriedades espectrais subjacentes do modelo e da tarefa.