Resumo Técnico: GEFEN: OTIMIZADOR ESTOCÁSTICO OTIMIZADO

Definição do Problema

O aprendizado profundo moderno depende fortemente do otimizador AdamW devido ao seu desempenho robusto em diversas arquiteturas. No entanto, o AdamW mantém médias móveis dos primeiros e segundos momentos para cada parâmetro, tipicamente em FP32. Isso resulta em uma pegada de memória do estado do otimizador de aproximadamente o dobro do tamanho dos parâmetros do modelo (cerca de 8 bytes por parâmetro). Para modelos de grande escala (ex: 13B parâmetros), esse estado sozinho pode consumir ~97 GiB de memória, restringindo severamente o tamanho dos modelos treináveis, o tamanho máximo dos lotes (batch sizes) e o rendimento (throughput) do treinamento em hardware com restrição de memória.

Alternativas existentes de eficiência de memória enfrentam compromissos significativos:

• Otimizadores Quantizados (ex: Adam8bit, Adam4bit): Reduzem a memória, mas frequentemente dependem de codebooks manuais e subótimos, tamanhos de bloco fixos (introduzindo hiperparâmetros) ou podem não ser universalmente aplicáveis (ex: evitando a quantização de embeddings).
• Otimizadores Estruturados (ex: Adam-mini): Compartilham estatísticas do segundo momento baseadas em regras manuais que mimetizam estruturas de bloco-diagonal da Hessiana. No entanto, estes requerem metadados específicos da arquitetura (ex: número de cabeças de atenção) e carecem de uma justificativa teórica forte sobre por que o agrupamento alinhado à Hessiana é benéfico.

O objetivo é projetar um otimizador que reduza significativamente a pegada de memória (além do Adam-mini) mantendo o desempenho de nível AdamW, sem exigir configuração manual, hiperparâmetros adicionais ou reajuste da taxa de aprendizado.

Metodologia

O método proposto, Gefen, alcança a redução de memória através de dois mecanismos primários: compartilhamento automático do segundo momento e quantização aprendida do primeiro momento.

1. Motivação Teórica

A abordagem é fundamentada em um resultado teórico (Teorema 1.1) demonstrando que, para parâmetros com grandes entradas mistas da Hessiana, a razão de seus quadrados de gradientes é restrita em direção a um. Isso implica que parâmetros com alta afinidade de Hessiana possuem naturalmente magnitudes de gradiente ao quadrado semelhantes, tornando-os candidatos ideais para compartilhar estimativas do segundo momento. Embora computar a Hessiana diretamente seja intratável em escala, os autores observam que a estrutura dos gradientes ao quadrado no passo inicial de treinamento correlaciona-se com essa estrutura da Hessiana.

2. Particionamento de Bloco Automático

Em vez de depender de regras manuais ou computação de Hessiana, o Gefen infere a estrutura de bloco ideal a partir dos gradientes ao quadrado do primeiro passo de otimização ($g_0$).

• Algoritmo: Para um tensor de parâmetros achatado (flattened), o algoritmo testa todos os divisores próprios $p$ do tamanho do tensor $n$. Ele remodela os gradientes ao quadrado em blocos de tamanho $p$ e calcula a heterogeneidade intra-bloco $E(p)$ (a variância dos gradientes ao quadrado dentro dos blocos).
• Estratégia de Seleção: O algoritmo seleciona o período $p$ que produz a primeira queda pronunciada em $E(p)$ conforme o tamanho do bloco aumenta, favorecendo blocos maiores para eficiência de memória enquanto garante a similaridade do gradiente. Se nenhuma queda significativa for encontrada ou se o $p$ ideal for $< 8$, ele assume por padrão $p=1$ (sem compartilhamento).
• Uso: Este esquema de particionamento único é usado tanto para o compartilhamento do segundo momento quanto para a quantização do primeiro momento.

3. Quantização por Programação Dinâmica Exata

Para reduzir a memória do primeiro momento (momentum), o Gefen emprega um codebook de quantização aprendível e exato, em vez de esquemas fixos ou heurísticos.

• Aprendizado do Codebook: Usando os gradientes iniciais, o método constró-lo um histograma de valores de gradiente normalizados. Em seguida, resolve o problema de quantização escalar ótimo usando Programação Dinâmica (DP) (inspirado em Wang & Song, 2011) para encontrar o codebook exato para a distribuição e largura de bits dadas.
• Estabilidade: Ao contrário do Lloyd-Max, que pode sofrer de "contração" (encolhimento do intervalo de quantização ao longo do tempo) e sensibilidade à inicialização, o Gefen força os endpoints do codebook a serem fixos em $-1$e$1$. Isso garante que o intervalo de quantização seja preservado durante todo o treinamento.
• Eficiência: O codebook é aprendido apenas uma vez no início do treinamento (usando CUDA para o histograma e CPU/Numba para a DP) e é reutilizado para toda a execução do treinamento.

4. Fluxo do Algoritmo

O Gefen opera como um substituto direto (drop-in replacement) para o AdamW:

1. Inicialização: Particiona os parâmetros com base em $g_0$ e aprende o codebook de quantização.
2. Passo de Atualização:
   • Computa os gradientes.
   • Atualiza o primeiro momento (momentum) desquantizando os valores do codebook armazenado, aplicando o momentum, requantizando e atualizando o máximo absoluto por bloco.
   • Atualiza o segundo momento computando a média dos gradientes ao quadrado dentro de cada bloco e aplicando a média móvel exponencial.
   • Aplica a correção de viés (bias correction) e atualiza os parâmetros usando o segundo momento compartilhado e o primeiro momento quantizado.

Principais Contribuições

1. Eficiência de Memória: O Gefen reduz a pegada de memória do otimizador em aproximadamente 8× comparado ao AdamW (reduzindo de ~2.0× o tamanho do parâmetro para ~0.25×), correspondendo a uma redução de ~6.5 GiB por bilhão de parâmetros.
2. Adaptação Automática: O método não requer metadados específicos da arquitetura, regras de agrupamento manual ou hiperparâmetros adicionais além dos padrões do AdamW. Ele infere automaticamente as estruturas de bloco a partir dos dados.
3. Fundamentação Teórica: Fornece um elo teórico entre as entradas da Hessiana e as razões de gradiente ao quadrado, justificando o compartilhamento de estatísticas do segundo momento para parâmetros alinhados à Hessiana.
4. Quantização Ótima: Introduz uma abordagem de programação dinâmica baseada em histograma para quantizar o momentum, evitando a subotimalidade e instabilidade de métodos iterativos como o Lloyd-Max.
5. Paridade de Desempenho: Demonstra que essas reduções não trazem prejuízos à qualidade do modelo, igualando o desempenho do AdamW em várias tarefas.

Resultados Experimentais

Os autores avaliaram o Gefen contra AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor e SM3 em modelos CNN, ResNet, GPT-2 e Llama 3.

• Pegada de Memória de Pico: O Gefen consistentemente alcançou o menor uso de memória de pico entre os métodos comparados. Para um modelo de 1.5B de parâmetros, o Gefen reduziu o estado do otimizador de 11.2 GiB (AdamW) para 1.5 GiB.
• Desempenho (Qualidade): Em todas as configurações (OpenWebText, C4, ImageNet, MNIST), o Gefen alcançou perdas de validação e acurácias comparáveis ao AdamW. Notavelmente, superou o Adam4bit e o Adam-mini em várias métricas mantendo o mesmo cronograma de taxa de aprendizado.
• Rendimento (Throughput): A redução na pegada de memória permitiu lotes micro (microbatches) maiores em treinamento distribuído:
  • FSDP (Llama 3-1.5B): O Gefen permitiu um tamanho de microbatch de 2 (vs. 1 para outros), resultando em um aumento de 56% no throughput (tokens/seg).
  • DDP (GPT-2-1B): O Gefen permitiu o treinamento com um tamanho de microbatch de 2, enquanto o AdamW não conseguia comportar sequer um microbatch de 1. Isso resultou em uma melhoria de 21% no throughput sobre o Adam-mini.
• Estabilidade: Estudos de ablação confirmaram que o codebook aprendido permanece estável durante todo o treinamento, validando a decisão de aprendê-lo apenas uma vez na inicialização.

Significância e Alegações

O artigo posiciona o Gefen como um substituto prático e "drop-in" para o AdamW que aborda o gargalo de memória no treinamento em larga escala sem sacrificar o desempenho ou exigir ajuste manual.

• Impacto Prático: Ao reduzir o estado do otimizador, o Gefen permite treinar modelos maiores ou usar lotes maiores em hardware fixo, melhorando diretamente o rendimento do treinamento e a eficiência de amostragem.
• Generalidade: Ao contrário de métodos anteriores que dependem de heurísticas manuais ou nomes de tensores específicos, o Gefen é agnóstico à arquitetura e se adapta automaticamente aos dados.
• Complementaridade: Os autores observam que a compressão do estado do otimizador do Gefen é ortogonal à compressão de comunicação de gradiente (ex: 1-bit Adam) e às estratégias de treinamento distribuído (FSDP/DDP), o que significa que pode ser combinado com essas técnicas para ganhos adicionais.

Os autores fornecem uma implementação Python completa com kernels CUDA fundidos, enfatizando a reprodutibilidade e a facilidade de integração em pipelines de treinamento existentes.