🤖 aprendizado de máquina

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

Este artigo apresenta o MKOR, um otimizador de segunda ordem com suporte a momentum que utiliza atualizações de rank-1 e inversão baseada em Sherman-Morrison para reduzir a complexidade computacional e de comunicação de cúbica para quadrática e linear, respectivamente, acelerando significamente o treinamento e a convergência de grandes modelos de linguagem em comparação com os métodos de primeira e segunda ordem de última geração.

Autores originais: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Publicado 2024-01-30✓ Author reviewed ⓘ

📖 7 min de leitura🧠 Leitura aprofundada

Autores originais: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

No vasto cenário da inteligência artificial moderna, os sistemas mais poderosos são construídos sobre redes neurais profundas, estruturas matemáticas intrincadas que aprendem a reconhecer padrões, traduzir idiomas e gerar texto. Para ensinar esses sistemas, os pesquisadores devem resolver um problema de otimização massivo: encontrar o conjunto perfeito de números internos, ou pesos, que permitam à rede fazer previsões precisas. Por décadas, a abordagem padrão tem sido dar um leve empurrão nesses pesos na direção que mais rapidamente reduz o erro, um método semelhante a um caminhante tateando o caminho para baixo de uma montanha com neblina, verificando a inclinação sob seus pés. Este método de primeira ordem é confiável e amplamente utilizado, mas pode ser lento, muitas vezes levando milhares de passos para alcançar o fundo do vale.

Uma estratégia mais ambiciosa, conhecida como otimização de segunda ordem, tenta dar ao caminhante um mapa de todo o terreno, incluindo a curvatura do solo, permitindo descidas muito mais rápidas. Ao compreender como a inclinação muda, esses métodos podem, teoricamente, alcançar a solução ideal em muito menos passos. No entanto, historicamente, essa abordagem tem sido computacionalmente cara demais para ser prática para os modelos massivos usados hoje. Calcular a informação de curvatura necessária para uma rede grande exige um poder de processamento e memória imensos, muitas vezes tornando o método mais lento do que a abordagem mais simples que ele visa substituir. O desafio tem sido capturar a velocidade do caminhante que lê o mapa sem o pesado fardo de carregar o mapa.

Uma equipe de pesquisadores da Universidade de Toronto, da Universidade Rutgers e do Texas Advanced Computing Center desenvolveu uma nova ferramenta chamada MKOR que resolve esse dilema. Eles criaram um otimizador com suporte a momentum que utiliza um atalho matemático específico para aproximar a forma do terreno sem calcular cada detalhe. Em vez de tentar computar a curvatura total e complexa do cenário de perda, o MKOR utiliza uma técnica chamada atualizações de rank-1. Esta abordagem simplifica o cálculo ao focar na direção de mudança mais significativa, reduzindo efetivamente o custo computacional de uma relação cúbica para uma quadrática. Em termos práticos, isso significa que o método escala muito melhor à medida que os modelos crescem, evitando os gargalos que anteriormente impediam que métodos de segunda ordem fossem usados em grandes modelos de linguagem.

Os pesquisadores testaram este novo otimizador em um grande modelo de linguagem conhecido como BERT-Large-Uncased, realizando experimentos em 64 processadores gráficos de alto desempenho. Eles compararam o MKOR contra os melhores métodos de primeira ordem existentes e as implementações de segunda ordem mais avançadas disponíveis. Os resultados foram impressionantes. Em uma tarefa envolvendo resposta a perguntas, o método de primeira ordem padrão exigiu 1.536 passos para atingir um nível específico de precisão. O novo método MKOR alcançou a mesma ou melhor precisão em apenas 600 passos. Ao medir o tempo total necessário para completar o treinamento, o MKOR terminou a tarefa em pouco mais de três horas, enquanto o método padrão levou quase oito horas. Isso representa um aumento de velocidade de 2,57 vezes, cortando efetivamente o tempo de treinamento em mais da metade.

Além da velocidade bruta, o estudo abordou uma questão crítica de estabilidade que frequentemente assombra os métodos de segunda ordem. Como esses métodos ajustam o processo de aprendizado com base em fatores matemáticos complexos, eles são propensos a "gradientes explosivos", um fenômeno onde os ajustes tornam-se tão grandes que o processo de treinamento falha ou diverge. Soluções anteriores frequentemente exigiam a redução significativa da taxa de aprendizado para manter a segurança, o que anulava os benefícios de velocidade. O MKOR inclui um estabilizador integrado que monitora o tamanho desses ajustes. Se os valores crescerem demais, o sistema os reescala automaticamente, evitando que o treinamento exploda, enquanto mantém o ritmo agressivo da abordagem de segunda ordem. Isso permite que o otimizador se mova rapidamente sem o risco de instabilidade.

Os pesquisadores também observaram que a vantagem dos métodos de segunda ordem é mais pronunciada nos estágios iniciais do treinamento, quando o modelo está aprendendo os contornos amplos dos dados. À medida que o modelo se aproxima da solução ideal, os benefícios dos cálculos de curvatura complexos diminuem. Para resolver isso, eles introduziram uma versão híbrida do otimizador, chamada MKOR-H. Esta versão começa com o poderoso método de segunda ordem para acelerar a fase inicial de aprendizado, mas muda automaticamente para um método de primeira ordem mais simples assim que a perda de treinamento para de diminuir rapidamente. Esta abordagem híbrida permitiu que o sistema convergisse em ainda menos passos, alcançando a mesma precisão final que o método padrão em menos da metade do tempo.

O estudo também examinou por que este atalho matemático específico funciona tão bem. Os pesquisadores analisaram as estruturas de dados dentro da rede neural e descobriram que as matrizes de covariância, que descrevem como diferentes partes da rede interagem, são frequentemente de baixo rank. Isso significa que, apesar da alta dimensionalidade dos dados, a informação está concentrada em algumas direções fundamentais. Ao aproximar essas matrizes usando atualizações de rank-1, o MKOR captura a informação essencial necessária para a otimização sem o overhead dos cálculos de matriz completa. A equipe verificou, através de testes extensivos, que essa aproximação introduz muito pouco erro, confirmando que a simplificação não ocorre às custas da precisão.

Além de velocidade e estabilidade, o novo método reduz significamente a quantidade de dados que precisam ser transferidos entre diferentes processadores durante o treinamento. Na computação distribuída, onde milhares de processadores trabalham juntos, a comunicação pode se tornar um grande gargalo. O MKOR reduz essa complexidade de comunicação de uma relação quadrática para uma linear, o que significa que, à medida que mais processadores são adicionados, o sistema escala eficientemente sem ser sobrecarregado por atrasos na transferência de dados. Essa eficiência permite que o método lide com os tamanhos de lote massivos exigidos pelos modernos grandes modelos de linguagem, o que era uma limitação para abordagens anteriores de segunda ordem.

As implicações deste trabalho estendem-se para além do treinamento de apenas um modelo específico. Ao tornar a otimização de segunda ordem prática para sistemas de grande escala, o MKOR oferece um caminho para o treinamento mais eficiente da inteligência artificial. Os pesquisadores observaram que seu método pode ajudar a economizar quantidades significativas de energia e tempo para profissionais de aprendizado de máquina. Embora a tecnologia em si seja uma ferramenta neutra, sua aplicação em campos que vão desde o diagnóstico médico ao modelagem climática pode acelerar a descoberta científica. No entanto, os autores também reconheceram que algoritmos de otimização eficientes podem ser aplicados a qualquer modelo, incluindo aqueles com impactos sociais negativos, enfatizando que o uso responsável de tal tecnologia depende das escolhas das pessoas que a utilizam.

Em última análise, o artigo apresenta uma solução concreta para um problema de longa data no aprendizado profundo. Ao combinar momentum, aproximações de rank-1 e um mecanismo de alternância híbrido, os pesquisadores criaram um otimizador que é mais rápido, mais estável e mais escalável que seus predecessores. Os experimentos demonstram que é possível aproveitar o poder dos métodos de segunda ordem sem os custos computacionais proibitivos que os mantiveram à margem. À medida que o campo da inteligência artificial continua a avançar em direção a modelos maiores e mais complexos, ferramentas como o MKOR fornecem uma maneira de treinar esses sistemas de forma mais eficiente, transformando o que antes era uma vantagem teórica em uma realidade prática.

Resumo Técnico: MKOR – Otimizador Baseado em Fator de Kronecker com Suporte a Momentum Usando Atualizações de Rank-1

1. Definição do Problema

Métodos de otimização de segunda ordem, como o Gradiente Natural (NGD) e a Curvatura Aproximada Fatorada por Kronecker (KFAC), oferecem taxas de convergência superiores em comparação com métodos de primeira ordem (ex: SGD, Adam, LAMB) ao utilizar informações de curvatura (aproximações da matriz Hessiana ou da Matriz de Informação de Fisher). No entanto, sua adoção no treinamento de Redes Neurais Profundas (DNNs) de larga escala, particularmente Transformers e Grandes Modelos de Linguagem (LLMs), é dificultada por três gargalos principais:

Complexidade Computacional: Métodos tradicionais de KFAC exigem a inversão de matrizes de covariância de tamanho $d \times d$ (onde $d$ é a dimensão da camada), resultando em uma complexidade de $O(d^3)$ . Da mesma forma, métodos baseados em Sherman-Morrison-Woodbury (SNGD) frequentemente escalam com o tamanho do lote $b$ ( $O(b^3)$ ). Em Transformers, onde os tamanhos de lote escalam com o comprimento da sequência, essas complexidades tornam-se proibitivas.
Sobrecarga de Comunicação: O treinamento distribuído exige a sincronização de informações de segunda ordem (matrizes de covariância e suas inversas) entre os trabalhadores. O KFAC exige a sincronização de $O(d^2)$ dados, enquanto o SNGD exige $O(bd + b^2)$ , criando gargalos de comunicação significativos que limitam a escalabilidade.
Frequência de Atualização e Estabilidade: Devido aos altos custos de inversão, implementações de última geração (como o KAISA) atualizam as informações de segunda ordem infrequencialmente (ex: a cada 100–1000 iterações) usando fatores "obsoletos" (stale). Isso reduz a eficácia da informação de curvatura. Além disso, métodos de segunda ordem são propensos à instabilidade numérica (gradientes explosivos, matrizes singulares) e frequentemente exigem fatores de amortecimento ou aritmética de alta precisão, o que aumenta ainda mais os custos computacionais.

2. Metodologia: O Algoritmo MKOR

Os autores propõem o MKOR (Otimizador Baseado em Fator de Kronecker com Suporte a Momentum Usando Atualizações de Rank-1), que aborda as questões acima através de quatro mecanismos centares:

Aproximação de Rank-1 & Inversão Baseada em SM:
Em vez de computar e inverter matrizes de covariância completas, o MKOR aproxima a covariância de ativações e gradientes usando matrizes de rank-1 derivadas da média das amostras ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). Ao aplicar a identidade de Sherman-Morrison a essas atualizações de rank-1, o MKOR atualiza diretamente a inversa dos fatores de Kronecker. Isso reduz a complexidade computacional da inversão de fatores de $O(d^3)$ para $O(d^2)$ .
- Regra de Atualização: Os fatores inversos são atualizados usando momentum ( $\gamma$ ) e a fórmula de Sherman-Morrison, permitindo atualizações frequentes sem o custo da inversão de matriz completa.
Complexidade de Comunicação Linear:
Em ambientes distribuídos, o MKOR sincroniza apenas os dois vetores de aproximação de rank-1 (tamanho $2d$ ) entre os trabalhadores, em vez de matrizes de covariância completas ( $4d^2$ ). Isso reduz a complexidade de comunicação de quadrática $O(d^2)$ para linear $O(d)$ . Adicionalmente, o MKOR utiliza meia-precisão (FP16) para esses vetores, reduzindo ainda mais os custos de comunicação, um feito difícil para outros métodos de segunda ordem devido à sensibilidade numérica.
Mecanismos de Estabilização:
Para prevenir a instabilidade numérica e gradientes explosivos comuns em métodos de segunda ordem:
- Estabilizador Baseado em Norma: Monitora a norma infinito dos fatores inversos. Se excederem um limite, os fatores são misturados com a matriz identidade ( $\hat{L} = \zeta L + (1-\zeta)I$ ) para evitar divergência.
- Redimensionamento de Gradiente: Gradientes pré-condicionados são redimensionados para corresponder à norma dos gradientes originais, garantindo que os agendadores de taxa de aprendizado permaneçam eficazes.
- Estabilidade Teórica: Os autores provam (Lema 3.1) que os fatores permanecem definidos positivos sem a necessidade de fatores de amortecimento artificiais ( $\mu I$ ), e que a divisão escalar única na regra de atualização é garantida como não-zero.
MKOR Híbrido (MKOR-H):
Reconhecendo que as vantagens de segunda ordem frequentemente diminuem conforme o treinamento converge (onde o cenário de perda se aproxima da identidade), o MKOR-H introduz um mecanismo de alternância. Ele monitora a taxa de redução da perda e retorna automaticamente para um otimizador de primeira ordem se as atualizações de segunda ordem não fornecerem mais aceleração, otimizando o equilíbrio entre custo e velocidade de convergência.

3. Principais Contribuições

Redução de Complexidade: O MKOR alcança complexidade computacional de $O(d^2)$ e complexidade de comunicação de $O(d)$ para atualizações de segunda ordem, aliviando os gargalos de $O(d^3)$ e $O(d^2)$ encontrados no KFAC e $O(b^3)$ no SNGD.
Atualizações de Alta Frequência: Ao reduzir os custos de inversão, o MKOR permite atualizações de segunda ordem a cada iteração (ou muito frequentemente), eliminando a necessidade de fatores "obsoletos" e explorando totalmente a informação de curvatura.
Robustez: A introdução da estabilização baseada em norma e do redimensionamento de gradiente permite que o MKOR opere de forma estável sem fatores de amortecimento ou aritmética de dupla precisão, permitindo o uso eficiente de hardware de meia-precisão.
Estratégia Híbrida: A variante MKOR-H alterna dinamicamente entre métodos de segunda e primeira ordem para manter a eficiência durante todo o ciclo de vida do treinamento.

4. Resultados Experimentais

Os autores avaliaram o MKOR no BERT-Large-Uncased utilizando 64 GPUs A100, comparando-o com o otimizador de primeira ordem de última geração LAMB e as implementações de segunda ordem KAISA (KFAC) e HyLo (SNGD), além de Eva.

Aceleração (Speedup): O MKOR alcançou um aceleração de 2.57× sobre o LAMB e uma aceleração de 1.85× sobre o KAISA na tarefa de pré-treinamento. Especificamente, reduziu o tempo de treinamento do BERT-Large-Uncased de 8 horas para 3 horas.
Convergência: Na tarefa de questionamento-resposta SQuAD v1.1, o MKOR-H convergiu em 600 passos (vs. 1.536 para o LAMB) enquanto alcançou pontuações F1 comparáveis ou superiores (90.64 vs. 90.44).
Benchmark GLUE: O MKOR alcançou novas métricas de última geração no dataset GLUE com 1.500 passos. Notavelmente, enquanto o KAISA falhou em convergir para a média de precisão de referência mesmo após 1.563 passos, o MKOR convergiu com sucesso.
Escalabilidade: Ao contrário do HyLo, que divergiu ou não coube na memória da GPU devido à alta sobrecarga de memória ( $O(bd + b^2)$ ), o MKOR treinou com sucesso com grandes tamanhos de lote (até 64k de tamanho de lote efetivo via acumulação) sem problemas de memória.
Validade da Aproximação: Experimentos no ResNet-50 e BERT mostraram que aproximações de rank-1 de matrizes de covariância incorrem em erro desprezível, validando o pressuposto central de aproximação.

5. Significância e Alegações

O artigo afirma que o MKOR consegue unir com sucesso o benefício teórico da otimização de segunda ordem com as restrições práticas do treinamento de DNNs modernas de larga escala. Ao reduzir a complexidade das atualizações de segunda ordem para um nível comparável aos métodos de primeira ordem (em termos de escalonamento com o tamanho do modelo) enquanto mantém as vantagens de convergência, o MKOR permite:

Treinamento Mais Rápido: Redução significativa no tempo de treinamento real (wall-clock) para LLMs.
Melhor Convergência: A capacidade de convergir para métricas de precisão mais altas onde outros métodos de segunda ordem (como o KAISA) falham ou estagnam.
Escalabilidade: A capacidade de escalar para grandes tamanhos de lote e dimensões de modelo sem os gargalos de comunicação e memória que anteriormente limitavam os métodos de segunda ordem.

Os autores enfatizam que seu trabalho torna a otimização de segunda ordem viável para a era dos "modelos grandes", abordando especificamente os problemas de escalabilidade introduzidos pelos tamanhos de lote dependentes do comprimento da sequência em arquiteturas Transformer. Eles concluem que o MKOR representa um novo estado da arte para o treinamento eficiente e eficaz de redes neurais profundas.