🤖 machine learning

Riemannian Gradient Descent for Low-Rank Architectures

Este artigo investiga o gradiente descendente riemanniano em dez designs algorítmicos para parâmetros de matrizes fatoradas por posto em aprendizado profundo, mas descobre que, apesar do ajuste, esses métodos não superam conclusivamente um baseline AdamW quando aplicados à atenção de múltiplas cabeças em modelos de linguagem pequenos.

Autores originais: Nicholas Knight

Publicado 2026-06-02

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Nicholas Knight

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

A Visão Geral: Navegando em uma Montanha vs. Dirigindo um Carro

Imagine que você está tentando encontrar o ponto mais baixo de um vasto vale enevoado (isso representa treinar um modelo de aprendizado de máquina para cometer menos erros). Normalmente, os cientistas da computação usam um método padrão chamado AdamW, que é como dirigir um carro com um GPS muito inteligente. O GPS diz qual caminho é "para baixo", e você dá um passo. Se você atingir um solavanco, a suspensão do carro (momentum) ajuda você a continuar se movendo. Isso funciona muito bem, mas trata cada parte do motor do carro como uma parte separada e independente.

Este artigo faz uma pergunta diferente: E se tratarmos as partes do motor como um sistema único e conectado?

Na IA moderna, muitas partes do cérebro (o modelo) são construídas usando matrizes de "baixo posto" (low-rank). Pense nelas não como um bloco gigante e sólido de dados, mas como uma folha de papel dobrada. Você pode descrever toda a forma do papel apenas conhecendo as coordenadas das linhas de dobra (duas matrizes menores, $A$ e $B$ ).

O autor argumenta que os métodos padrão tratam as linhas de dobra como se fossem independentes, o que é como tentar alisar um papel amassado puxando pelos cantos sem perceber que o papel está conectado. Em vez disso, o autor propõe o uso de Gradiente Descendente Riemanniano.

A Analogia:

Método Padrão (Euclidiano): Imagine caminhar em um chão plano. Você dá um passo na direção que o GPS indica. Se você estiver em um papel amassado, pode acidentalmente dar um passo fora do papel, sendo forçado a pular de volta para dentro.
Método Riemanniano: Imagine que você é um surfista andando em uma onda. Você está restrito a permanecer sobre a superfície da onda. Você não apenas caminha; você desliza ao longo da curva da onda. Este método garante que você nunca saia da "forma" da matriz de baixo posto, mantendo a matemática limpa e teoricamente sólida.

O Experimento: Testando Novas Pranchas de Surf

O autor construiu dez diferentes "pranchas de surf" (algoritmos) baseadas nesta ideia de surfar ondas. Ele as testou em um pequeno modelo de linguagem (um cérebro que aprende a prever a próxima palavra em uma frase).

Ele testou dois tipos principais de ondas:

Ondas de Posto Fixo (Fixed-Rank): Apenas mantendo o papel dobrado com uma certa rigidez.
Ondas de Isometria Parcial (Partial Isometry): Uma regra mais estrita onde o papel deve ser perfeitamente ortogonal (como um ângulo reto perfeito), o que é como uma onda mais rígida e estruturada.

Ele também testou uma versão de "Grade" (Grid), onde múltiplas partes do modelo compartilham as mesmas linhas de dobra (como um grupo de pessoas segurando a mesma corda).

Os Resultados: Um Resultado "Bom", mas Não "Ótimo"

Aqui está a conclusão honesta do artigo, livre de hype:

Funciona, Mas Não é uma Bala de Prata: Os novos métodos treinaram o modelo com sucesso. Eles não travaram e aprenderam a tarefa. Isso prova que a matemática é sólida e o código funciona.
Sem uma Vitória Clara: Após ajustar cuidadosamente a "velocidade" (taxa de aprendizado) das novas pranchas de surf, elas não venceram consistentemente o carro AdamW padrão.
- Em alguns casos, foram ligeiramente melhores.
- Em outros casos, foram ligeiramente piores.
- No geral, a diferença foi tão pequena que poderia ser atribuída ao ruído aleatório (como uma leve brisa mudando o caminho de uma folha).
O Custo: Os novos métodos são computacionalmente mais caros. É como dirigir um carro elétrico de alta tecnologia que requer uma estação de carregamento especial, enquanto o método padrão é um carro a gasolina confiável que funciona em qualquer lugar. Como o novo carro não o levou ao destino significativamente mais rápido ou melhor, o autor conclui que ele ainda não está pronto para substituir o método padrão para uso cotidiano.

Pontos Principais para o Público Geral

A Teoria é Bela: A ideia de tratar os parâmetros de uma matriz como uma superfície curva (uma variedade ou manifold) em vez de uma grade plana é matematicamente elegante e logicamente consistente. Ela respeita as relações ocultas entre os números.
A Prática é Complicada: Embora a teoria diga "você deve chegar lá mais rápido", a realidade do aprendizado profundo (com seu ruído, dados massivos e arquiteturas complexas) significa que os métodos mais simples e padrão (AdamW) ainda são os campeões por enquanto.
Potencial Futuro: O autor permanece otimista. Só porque um novo motor não vence uma corrida em uma pista pequena não significa que ele não vencerá em uma rodovia enorme. O autor sugere que, à medida que os modelos crescerem (escalabilidade), esses métodos geométricos podem finalmente mostrar seu verdadeiro poder.

O Que o Artigo Não Alega

Ele não alega que este método curará doenças ou resolverá as mudanças climáticas.
Ele não alega que este é o futuro da IA agora.
Ele não alega que o método é mais rápido ou barato (na verdade, admite que é provavelmente mais lento e complexo).

Em resumo: O autor construiu uma forma geométrica sofisticada de treinar modelos de IA que respeita a estrutura interna dos dados. Funciona perfeitamente bem, mas, por enquanto, o método antigo e simples (AdamW) ainda é a melhor escolha para realizar o trabalho de forma eficiente. O novo método é uma "aposta de longo prazo" promissora que pode valer a pena quando construirmos modelos muito maiores no futuro.

Resumo Técnico: Gradiente Descendente Riemanniano para Arquiteturas de Baixo Posto

Enunciado do Problema
No aprendizado profundo moderno, os parâmetros do modelo são frequentemente matrizes restritas a um baixo posto para reduzir o custo computacional e a contagem de parâmetros (por exemplo, em mecanismos de atenção multi-cabeça onde $W = AB^T$ ). As atualizações padrão de treinamento utilizam gradientes euclidianos para os fatores $A$ e $B$ de forma independente, ignorando a estrutura geométrica do conjunto de matrizes de posto- $r$ . Este artigo investiga se tratar a matriz de posto restrito $W$ como um ponto em uma subvariedade suave e aplicar otimização Riemanniana pode resultar em melhor convergência ou melhores soluções. O desafio central é que a fatoração $W=AB^T$ não é única (invariante sob $A \to AS, B \to BS^{-T}$ ) e o conjunto de matrizes de posto- $r$ é uma variedade, não um espaço vetorial. Os autores visam restringir as modificações apenas ao otimizador, evitando mudanças na arquitetura do modelo, mantendo uma complexidade por iteração comparável aos métodos padrão ( $O(mr^2 + nr^2)$ ).

Metodologia
Os autores propõem um framework de gradiente descendente riemanniano adaptado para parâmetros fatorados em posto. A iteração geral envolve:

Recuperar o gradiente euclidiano ambiente $G_W$ a partir dos gradientes dos fatores $G_A$ e $G_B$ .
Projetar este gradiente no espaço tangente da variedade específica para obter o gradiente riemanniano $G_W^R$ .
Atualizar o momento e normalizar o tamanho do passo.
Mover-se ao longo da geodésica (usando uma retração) para a próxima iteração.
Transportar paralelamente o vetor de momento para o novo espaço tangente.

O artigo instancia este template em dez designs algorítmicos específicos categorizados por geometria e restrições:

Geometrias para Matrizes de Posto Fixo ( $R_{r}^{m \times n}$ ):
- Geometria Embebida: Utiliza a métrica de Frobenius ambiente. A retração é realizada via SVD truncada de posto- $r$ de $W + \Xi$ . O transporte paralelo utiliza projeção métrica.
- Geometria de Quociente: Trata a variedade como um quociente do espaço produto dos fatores $(A, B)$ . A métrica é definida "acima" nos fatores, induzindo uma métrica no quociente. A retração é aditiva nos fatores; o transporte envolve projeção horizontal.
Geometrias para Isometrias Parciais de Posto- $r$ ( $P_{r}^{m \times n}$ ):
- Aqui, $W$ é restrita de tal forma que seus valores singulares não nulos são todos 1 (equivalente a $A$ e $B$ serem ortonormais em suas colunas).
- Geometria Embebida: Restringe a métrica de Frobenius. A retração usa SVD truncada com restrições de ortogonalidade específicas.
- Geometria de Quociente: Utiliza o produto das variedades de Stiefel ( $St_{m,r} \times St_{n,r}$ ) como o espaço "acima". As retrações usam decomposições QR ou Polares.
- Geometria Canônica: Utiliza a métrica canônica nos fatores de Stiefel, induzindo uma métrica de produto no quociente.
Compartilhamento de Pesos (Variedades de Grade):
- Estende os itens acima para arquiteturas onde os fatores são compartilhados entre blocos de linhas e blocos de colunas (ex: Group-Query Attention).
- Para Posto Fixo, os autores mostram que empilhar os fatores em matrizes maiores permite o uso de algoritmos de instância única, desde que a preservação do posto seja verificada por bloco.
- Para Isometrias Parciais, o empilhamento simples falha devido às restrições de ortogonalidade. Os autores desenvolvem algoritmos mais complexos, incluindo uma abordagem "híbrida" para o caso embebido, onde elevam para a geometria de quociente para realizar retrações e transporte, e então projetam de volta para baixo.

Principais Contribuições

Exploração Sistemática: O artigo fornece um levantamento abrangente e uma implementação de dez variantes distintas de otimização riemanniana para parâmetros de matriz de baixo posto, cobrindo tanto restrições de posto fixo quanto de isometria parcial, bem como cenários de compartilhamento de peso.
Detalhes de Implementação: Detalha as retrações específicas (SVD, QR, Polar) e os transportes de vetores necessários para cada geometria, abordando a não unicidade das fatorações e a preservação das restrições da variedade (posto e ortogonalidade).
Avaliação Empírica: Os métodos são aplicados às matrizes Q, K, V, O em modelos de linguagem de pequena escala (6 camadas, dimensão de embedding 512) usando o dataset FineWeb.

Resultados Experimentais

Desempenho: Após o ajuste das taxas de aprendizado, os métodos riemannianos não superaram conclusivamente o baseline AdamW. Nas configurações de Multi-Head Attention (MHA) e Group-Query Attention (GQA), as diferenças de perda estavam frequentemente dentro da variância das sementes aleatórias.
Hiperparâmetros: A normalização da métrica proporcionou melhorias leves (0.01–0.05 nats). O momento pesado ( $\nu=7/8$ ) geralmente prejudicou o desempenho ou proporcionou ganhos negligenciáveis em comparação ao baseline.
Estabilidade: Algumas execuções exibiram picos de perda (ex: por volta do passo 6660) que se recuperaram, frequentemente correlacionados com dados estruturados no conjunto de treinamento.
Complexidade: Os autores reconhecem que, embora os métodos sejam matematicamente sólidos, sua complexidade computacional por iteração é maior que a do AdamW, tornando-os menos atraentes para os experimentos de pequena escala conduzidos.

Significância e Alegações
O artigo mantém uma postura modesta em relação à sua novidade e impacto imediato:

Redescoberta: Os autores afirmam explicitamente que qualquer aparente novidade é "quase certamente uma redescoberta", baseando-se fortemente em livros didáticos estabelecidos (Absil et al., Boumal) e literatura prévia.
Motivação Heurística: O trabalho é motivado pela heurística de que considerar melhor as relações entre parâmetros (via geometria de variedade) poderia levar a uma convergência ou soluções melhores, embora isso não tenha sido definitivamente provado nos experimentos.
Potencial Futuro: Os autores argumentam que a falta de superação em modelos de pequena escala não exclui o mérito em escalas maiores. Eles sugerem que a abordagem é matematicamente sólida e que o "tradeoff custo-precisão" pode mudar favoravelmente com modelos maiores ou arquiteturas diferentes (ex: modelos baseados em tensores como MLA, MoLAE).
Limitações: O estudo admite limitações significativas, incluindo o tamanho pequeno do modelo, a falta de decaimento da taxa de aprendizado (annealing) e o uso de um único baseline (AdamW) sem o ajuste de seus hiperparâmetros.

Em conclusão, o artigo serve como uma demonstração técnica rigorosa da aplicação de otimização riemanniana a componentes de baixo posto de aprendizado profundo, fornecendo uma biblioteca de algoritmos e um baseline para pesquisas futuras, ao mesmo tempo em que reconhece que os resultados empíricos atuais ainda não justificam a substituição de otimizadores padrão como o AdamW.