Resumo Técnico: LoRDO – Otimização de Baixo Rank Distribuída com Comunicação Infrequente

1. Definição do Problema

O treinamento distribuído de modelos de fundação via Distributed Data Parallel (DDP) é cada vez mais limitado pela largura de banda de interconexão e pelo overhead de memória necessário para armazenar estados de otimizadores (ex: momentos do Adam). Embora estratégias que empregam comunicação infrequente (atualizações locais) reduzam a frequência de sincronização, elas permanecem limitadas pelos custos de memória e comunicação de estados de otimizadores de rank total.

Otimizadores adaptativos de baixo rank (ex: GaLore, LDAdam) oferecem um caminho para aliviar essas restrições ao projetar gradientes em um subespaço de baixo rank. No entanto, generalizar esses métodos para o regime de sincronização infrequente apresenta um desafio de otimização crítico:

1. Projeções Locais: Computar projeções de baixo rank localmente em cada worker usando gradientes estocásticos introduz ruído significativo devido aos pequenos tamanhos de batch efetivos por worker. Isso leva a subespaços de projeção subótimos e degradação de desempenho.
2. Projeções Globais: Usar uma matriz de projeção global derivada de pseudo-gradientes agregados (como em DDP) fornece uma base mais estável. No entanto, no regime de atualização local, essa abordagem restringe permanentemente a trajetória de otimização a um subespaço de baixo rank fixo. Sem sinais de rank total, a otimização estagna porque a base de projeção falha em "renovar-se" ou explorar novas direções fora do subespaço inicial.

A questão central abordada é: Como podemos projetar otimizadores de baixo rank de alto desempenho para treinamento com eficiência de comunicação que evitem a estagnação de subespaço, mantendo baixo overhead de comunicação e memória?

2. Metodologia: LoRDO

Os autores propõem o LoRDO (Low-Rank Distributed Optimization), um framework que unifica a otimização de baixo rank com a sincronização infrequente. A metodologia aborda a tensão entre estabilidade de projeção e exploração de subespaço através de três mecanismos principais:

A. Matrizes de Projeção Global

Em vez de cada worker determinar sua própria matriz de projeção $Q_m$ baseada em gradientes estocásticos locais (que sofre de alta variância), o LoRDO emprega uma matriz de projeção global $Q_t$. Esta matriz é computada na fronteira de sincronização usando o pseudo-gradiente agregado ($\Delta_t$), que representa a mudança total nos parâmetros do modelo ao longo de $K$ passos locais.

• Benefício: O pseudo-gradiente possui um tamanho de batch efetivo de $M \times B$ (onde $M$ é o número de workers), resultando em um sinal mais estável com variância reduzida em comparação aos gradientes locais. Ele também incorpora informações de curvatura contidas no pseudo-gradiente, levando a uma estimativa de base de maior qualidade.

B. Momento Quasi-Hiperbólico de Rank Total (QHM)

Para resolver o problema da "estagnação de subespaço" inerente às projeções globais, o LoRDO injeta um termo de momento quasi-hiperbólico de rank total na atualização local.

• Mecanismo: A regra de atualização combina o momento de baixo rank com um sinal de gradiente de rank total escalonado pelo segundo momento.
  $$\text{Update} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• Efeito: Essa injeção impede que o pseudo-gradiente agregado permaneça preso dentro de um subespaço fixo de rank-$r$. Ela permite que o procedimento de otimização agregue temporalmente subespaços de rank-$r$ de diferentes workers, permitindo a eventual exploração completa do subespaço. Isso restaura a capacidade de explorar o espaço de solução completo sem aumentar o overhead de comunicação.

C. Feedback de Erro e Rotação de Momento

O LoRDO incorpora feedback de erro localmente para contabilizar a informação perdida durante a compressão de baixo rank. Além disso, após o cálculo de uma nova matriz de projeção, os estados do otimizador (momentos) são rotacionados para garantir que permaneçam acumulados no mesmo subespaço, uma técnica adaptada do LDAdam.

3. Principais Contribuições

1. Análise de Modos de Falha de Projeção: Os autores demonstram que projeções locais prejudicam o desempenho devido à alta variância de pequenos batches de workers, enquanto projeções globais sem sinais de rank total induzem estagnação de subespaço.
2. Restauração da Exploração de Subespaço Completo: O LoRDO introduz um otimizador de baixo rank que injeta um sinal de gradiente de rank total (via QHM) nas atualizações locais, mantendo uma projeção global derivada de pseudo-gradientes agregados. Isso evita a estagnação sem aumentar o overhead de comunicação ou memória.
3. Paridade e Eficiência com DDP: O framework atinge quase a paridade com o DDP de baixo rank síncrono (gap de perplexidade < 1%), enquanto reduz o overhead de memória e comunicação em até $8\times$–$12\times$. Sob severas restrições de memória que exigem baixos ranks, o LoRDO supera o DDP em perplexidade.

4. Resultados Experimentais

Experimentos foram conduzidos em transformers decoder-only escalonados para 16M, 125M e 720M parâmetros usando a mistura de dados SmolLM2.

• Estagnação vs. Exploração: Em modelos de 16M, projeções globais sem QHM falharam em aprender (perplexidade 46.8) comparado ao DDP (36.6), confirmando a estagnação de subespaço. A adição do QHM de rank total restaurou o desempenho, permitindo que o LoRDO-Global superasse as variantes locais.
• Qualidade de Projeção: O LoRDO-Global demonstrou superior estabilidade e desempenho, particularmente em baixos ranks ($r=8, 16$), devido à base de projeção de maior qualidade derivada do pseudo-gradiente agregado. Mostrou-se menos sensível a mudanças no batch size local em comparação com métodos de projeção local.
• Escalabilidade para 720M: Na escala de 720M, o LoRDO-Global ($r=256$) alcançou quase a paridade com o DDP de baixo rank ($r=256$) com um gap de perplexidade inferior a 1% (10.41 vs. 10.34). Também igualou a precisão em tarefas downstream (ARC, HellaSwag, MMLU, PIQA) com o DDP.
• Ganhos de Eficiência:
  • Comunicação: O LoRDO reduz o overhead de comunicação em aproximadamente 10$\times$ em relação ao DDP de baixo rank e 25$\times$ em relação ao DDP de rank total.
  • Memória: Reduz o overhead de memória do estado do otimizador em 8$\times$ em relação aos equivalentes de rank total.
  • Tempo de Execução (Wall-Clock Time): Em ambientes de baixa largura de banda (modelado em 10 Gbps), o LoRDO supera significativamente o DDP em tempo de execução devido à redução na frequência de sincronização.

5. Significância e Alegações

O artigo afirma que o LoRDO resolve a tensão fundamental entre estabilidade de projeção e exploração de subespaço na otimização eficiente de comunicação de baixo rank. Ao utilizar pseudo-gradientes agregados para estabilidade de projeção e momento quasi-hiperbólico de rank total para exploração, o LoRDO possibilita:

• Treinamento Escalável: Desempenho competitivo em escala (até 720M de parâmetros) com custos significativamente reduzidos de comunicação e memória.
• Eficiência de Recursos: Viabilidade de pré-treinamento em hardware com recursos limitados (baixa memória/largura de banda), onde otimizadores de rank total são intratáveis.
• Aprendizado Descentralizado: Pavimentando o caminho para o aprendizado eficiente em ambientes descentralizados ao estender o treinamento de modelos além de centros de dados únicos.

Os autores observam limitações, incluindo a falta de uma prova formal de convergência e o foco atual em otimizadores adaptativos (Adam) em vez de otimizadores baseados em matrizes (ex: Muon), deixando estas questões para trabalhos futuros. No entanto, eles postulam que as tendências observadas em modelos menores se mantêm para escalas maiores.