Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

Resumo Técnico: Aceleração da Otimização Espectral de Ordem Zero com Ortogonalização Parcial a partir de Iteração de Potência

Declaração do Problema
A otimização de ordem zero (ZO) emergiu como uma técnica crítica para o ajuste fino de Modelos de Linguagem de Grande Escala (LLMs) em dispositivos de borda com restrições de memória, pois evita a retropropagação intensiva em memória exigida por métodos de primeira ordem (FO). No entanto, os métodos ZO sofrem de alta variância nas estimativas de gradiente devido a perturbações de Monte Carlo, levando a uma convergência significativamente mais lenta em comparação com os métodos FO.

Tentativas recentes de acelerar a otimização ZO focaram na busca em subespaço de baixa dimensão (por exemplo, LOZO, SubZero) e em otimizadores espectrais como o Muon, que exploram a geometria da matriz por meio de ortogonalização. Embora o Muon supere o AdamW em configurações FO ao ortogonalizar gradientes para explorar direções espectrais fracas, os autores observam que aplicar o Muon diretamente em configurações ZO produz ganhos limitados. A questão central identificada é que as estimativas de gradiente ZO são inerentemente ruidosas e pouco confiáveis. Diferentemente dos gradientes FO, que estão concentrados em direções singulares dominantes, as estimativas ZO exibem uma "cauda espectral mais pesada e plana" (Figura 1). Consequentemente, a ortogonalização Newton-Schulz (NS) padrão do Muon, que equaliza a energia em todo o espectro ativo, amplifica inadvertidamente essas direções ruidosas e imprecisas, transformando erros de estimativa em componentes de atualização prejudiciais.

Metodologia: ZO-MOPI
Os autores propõem o ZO-MOPI (Momento de Ordem Zero com Ortogonalização Parcial), um método projetado para acelerar a otimização ZO extraindo seletivamente direções espectrais dominantes enquanto suprime o ruído. A abordagem integra três componentes-chave:

1. Ortogonalização Parcial via Iteração de Potência em Fluxo (SPI):
   Em vez de realizar ortogonalização de espectro completo (como iterações NS), o ZO-MOPI emprega uma estratégia de "ortogonalização parcial". Foca exclusivamente nas top-$k$ direções singulares dominantes, descartando a cauda espectral ruidosa.

   • Mecanismo: O método substitui o procedimento NS por uma Iteração de Potência em Fluxo (SPI) mais rápida e concentrada. A SPI rastreia o subespaço de dimensão $k$ líder ao longo das etapas de treinamento.
   • Eficiência: Ao armazenar em cache o subespaço singular direito ($V_{t-1}$) e usá-lo como inicialização quente para a próxima iteração, a SPI evita SVD completo ou multiplicações de matriz caras. O custo computacional é reduzido para $O(rnk + nk^2)$, onde $k \ll r \ll \min(m, n)$, comparado ao custo mais alto dos métodos de espectro completo.
   • Estabilidade: O método baseia-se na observação de que as direções espectrais dominantes evoluem suavemente sob suavização de momento, tornando o subespaço anterior uma inicialização confiável para a etapa atual.

2. Estimativa ZO de Subespaço:
   Para reduzir ainda mais a variância, a estimativa de gradiente é restringida a um subespaço de baixo posto. A matriz de perturbação aleatória $Z$ é decomposta em $AB$, onde $A$ é uma matriz de subespaço fixa e $B$ é uma perturbação aleatória dentro desse subespaço. Isso reduz a dimensão de busca efetiva de $mn$ para $rn$, diminuindo significativamente a variância do estimador de gradiente.

3. Momento Projetado:
   O momento é mantido dentro do sistema de coordenadas reduzido $r \times n$, em vez do espaço de parâmetros completo, para economizar memória. Um componente crítico do método é o projeto do momento quando o subespaço $A$ é atualizado periodicamente (a cada $\nu$ iterações). O momento $M_{old}$ é projetado no novo subespaço $A_{new}$ via $M_{new} \leftarrow A_{new}^\top A_{old} M_{old}$.

   • Necessidade: Os autores demonstram que o momento não serve meramente para redução de variância, mas é essencial para a estabilidade da SPI. Sem momento, o subespaço dominante flutua drasticamente devido às estimativas ZO ruidosas, causando instabilidade na iteração de potência.

Contribuições Principais

• Identificação de Mismatch Espectral: O artigo destaca que a ortogonalização de espectro completo (Newton-Schulz) é prejudicial em configurações ZO porque amplifica o ruído inerente às estimativas de gradiente ZO, enquanto os gradientes FO estão naturalmente concentrados.
• Ortogonalização Parcial: A introdução de uma estratégia de ortogonalização seletiva que preserva apenas os componentes espectrais mais confiáveis (top-$k$), ignorando a cauda ruidosa.
• Iteração de Potência em Fluxo (SPI): Um algoritmo leve e eficiente para rastrear direções singulares dominantes que aproveita a continuidade temporal e a inicialização quente para evitar computações caras.
• Integração de Momento Projetado: Um mecanismo para manter a estabilidade do momento através de atualizações de subespaço, demonstrado como crítico para a eficácia do rastreamento espectral em ambientes ruidosos.

Resultados Experimentais
Os autores avaliam o ZO-MOPI em tarefas de ajuste fino de LLMs usando os modelos OPT-13B, LLaMA3-8B e Gemma2-2B em benchmarks SuperGLUE (SST-2, RTE, BoolQ, SQuAD).

• Velocidade de Convergência: O ZO-MOPI alcança convergência 1,5x a 4x mais rápida em comparação com o ZO-Muon, estado da arte atual. Por exemplo, nas tarefas RTE e SQuAD, o ZO-Muon requer aproximadamente 3x o tempo de relógio para atingir a mesma precisão que o ZO-MOPI.
• Precisão Final: O método alcança precisões finais competitivas ou superiores em comparação com bases fortes como MeZO, LOZO e ZO-Muon. Em vários casos (por exemplo, OPT-13B em SST-2 e RTE), o ZO-MOPI estabelece novos recordes de precisão ZO.
• Eficiência: O método mantém a eficiência de memória do ajuste fino ZO, com apenas sobrecarga de memória negligenciável para armazenar o subespaço e os estados da iteração de potência.
• Estudos de Ablação: Experimentos confirmam que remover o momento ou usar um intervalo de atualização de subespaço muito curto degrada significativamente o desempenho, validando a necessidade do projeto de momento proposto e das estratégias de amostragem de subespaço preguiçosa.

Significado e Alegações
O artigo alega que o ZO-MOPI conecta com sucesso a lacuna de desempenho entre a otimização ZO e FO para ajuste fino de LLMs, abordando as características espectrais específicas de estimativas de gradiente ruidosas. Os autores argumentam que a chave para acelerar a otimização ZO não é apenas reduzir a variância por meio de projeção de subespaço, mas também adotar um alinhamento espectral seletivo que evite amplificar o ruído. Ao substituir a ortogonalização agressiva de espectro completo por ortogonalização parcial via iteração de potência, o método permite que otimizadores espectrais funcionem efetivamente no regime ZO de alta variância. O trabalho sugere que a estrutura espectral pode ser melhor explorada mesmo sob estimativas ruidosas, oferecendo um caminho para o ajuste fino eficiente e com restrições de memória de grandes modelos sem depender da retropropagação.