🤖 machine learning

Token Geometry

O artigo apresenta o Ember, um otimizador leve que explora a geometria de gradiente única das matrizes de embedding e da camada LM-head para reduzir significativamente o uso de VRAM e melhorar a eficiência de treinamento em diversas tarefas, ao mesmo tempo em que desafia a visão convencional dos horizontes de otimização de redes neurais.

Autores originais: Kathan Shah

Publicado 2026-07-03

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Kathan Shah

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

Imagine que você está ensinando um robô gigante e superinteligente a falar a linguagem humana. Para fazer isso, você dá ao robô um dicionário massivo (a tabela de embeddings) e um guia de tradução (o LM-head) que conecta seus pensamentos internos às palavras reais.

Por muito tempo, a maneira padrão de ensinar esse robô foi usando um método chamado Adam. Pense no Adam como um professor muito minucioso, porém pesado e autoritário. Para cada única palavra no dicionário, o Adam mantém um caderno de "notas de estudo" (estado do otimizador) massivo para lembrar como o robô reagiu a essa palavra no passado.

O problema? À medida que o dicionário cresce para incluir milhões de palavras, esses cadernos tornam-se tão grandes que não cabem na memória do robô (VRAM). Isso força os pesquisadores a dividir o trabalho entre muitos computadores, o que é lento, caro e complicado.

Surge o Ember, um novo professor mais leve apresentado neste artigo.

A Grande Ideia: Um Toque Mais Leve

Os autores do artigo descobriram que a parte do "dicionário" do robô aprende de forma diferente do resto do seu cérebro. Enquanto o resto do cérebro precisa de notas complexas e pesadas, o dicionário só precisa de uma abordagem simples e simplificada.

A Analogia: A Verificação do "Z-Score"
Imagine que você está corrigindo a prova de um aluno.

Adam olha para cada resposta individual, lembra-se de cada vez que o aluno acertou ou errou e mantém um arquivo detalhado sobre cada erro específico. É como manter uma biografia de 100 páginas para cada única palavra.
Ember faz uma pergunta mais simples: "Com que frequência este aluno acerta esta palavra e qual é o nível de confiança dele?" Ele essencialmente converte o desempenho do aluno em um simples escore padrão (um "z-score"). Ele remove o peso excessivo e foca apenas no sinal central: O robô está melhorando nesta palavra ou não?

Como o Ember Economiza Espaço

O artigo afirma que o Ember é incrivelmente eficiente porque para de manter essas biografias massivas de 100 páginas.

Memória do Adam: Se você tiver um dicionário com 100.000 palavras, o Adam precisa de um caderno para cada palavra que é enorme. O artigo diz que isso ocupa gigabytes de espaço (como uma biblioteca inteira).
Memória do Ember: O Ember percebe que só precisa de uma lista minúscula de "com que frequência" e "quão confiante" para cada palavra. Ele reduz essa biblioteca para um único post-it (kilobytes).

O artigo mostra que, embora o Ember seja muito mais leve, ele ensina o robô tão bem quanto o Adam, que é pesado e autoritário. Na verdade, em algumas situações complicadas (como quando o robô está aprendendo com lotes de dados muito pequenos), o Ember na verdade aprende mais rápido e de forma mais estável.

A Descoberta Surpreendente: Uma Linha Reta

Uma das descobertas mais fascinantes do artigo é sobre como o robô aprende.

Crença Antiga: Os cientistas pensavam que o caminho de aprendizado do robô era como uma caminhada caótica por uma cordilheira nebulosa, zigzagueando para cima e para baixo, ficando preso em pequenos vales e seguando um caminho muito longo e sinuoso até o topo.
A Realidade do Ember: Os autores descobriram que, ao usar o Ember, o caminho de aprendizado do robô é, na verdade, uma rodovia reta e suave. Se você plotasse o progresso do robô em um gráfico, pareceria uma linha simples e reta movendo-se de "iniciante" para "especialista".

Isso sugere que o "cenário" de aprendizado para essas partes específicas do dicionário não é tão bagunçado quanto pensávamos. É um caminho direto, e o Ember é o veículo que dirige direto por ele sem se perder.

Por Que Isso Importa (De Acordo com o Artigo)

É Barato: Você pode treinar esses modelos em um único computador em vez de precisar de um enorme cluster de supercomputadores apenas para suportar a memória.
É Rápido: Como a memória é tão pequena, os pesquisadores podem testar novas ideias muito mais rápido.
Funciona em Todo Lugar: O artigo testou o Ember em diferentes tamanhos de robôs (de pequenos a muito grandes) e diferentes tarefas (aprender a falar, aprender a raciocinar e até gerar imagens). Em quase todos os casos, o Ember igualou ou superou o antigo padrão (Adam) enquanto usava uma fração da memória.

Resumo

O artigo apresenta o Ember, uma nova maneira de treinar a parte do "vocabulário" dos modelos de IA. Ele substitui o método pesado e guloso de memória "Adam" por uma abordagem inteligente e leve que trata o processo de aprendizado como um problema matemático simples (um z-score) em vez de uma biografia complexa. O resultado? Você obtém a mesma (ou melhor) inteligência, mas não precisa de um supercomputador para guardar as notas. Além disso, acontece que o robô aprende em uma linha reta, não em um zig-zag caótico.

Resumo Técnico: Geometria de Tokens e o Otimizador Ember

Declaração do Problema

Os atuais modelos de linguagem de grande escala (LLMs) otimizam as tabelas de embedding e as matrizes do LM-head (coletivamente denominadas "interface de tokens") utilizando o otimizador Adam. Embora o Adam apresente um bom desempenho, ele impõe um fardo de engenharia significativo no treinamento distribuído. Especificamente, o Adam mantém buffers de primeiro e segundo momentos para cada parâmetro, resultando em uma complexidade de estado do otimizador de $O(2VD)$ para um tamanho de vocabulário $V$ e dimensão de embedding $D$ . Para modelos grandes, esse estado consome gigabytes de VRAM, necessitando de estratégias de particionamento complexas (ex: ZeRO/FSDP) mesmo para parâmetros que não são o gargalo primário de computação do modelo.

Além disso, o artigo argumenta que a aplicação padrão do Adam às interfaces de tokens ignora a geometria específica desses parâmetros. Diferente das camadas lineares densas, as interfaces de tokens são atualizadas via gradientes esparsos (apenas a linha do token ativo é atualizada por passo) e o panorama de perda subjacente para a entropia cruzada pode possuir propriedades distintas que os otimizadores padrão falham em aproveitar eficientemente.

Metodologia: O Otimizador Ember

Os autores introduzem o Ember, um otimizador leve projetado especificamente para matrizes de embedding e LM-head. A metodologia central fundamenta-se na Geometria da Informação, especificamente na métrica de Informação de Fisher.

Fundamentação Teórica

O artigo postula que, para a perda de entropia cruzada, a métrica de Fisher (o produto externo esperado dos gradientes) é o termo de segunda ordem exato da divergência de Kullback–Leibler. Próximo ao ótimo, a métrica de Fisher aproxima o Hessiano.

Raiz Quadrada de Fisher: Em vez de aplicar um passo de Newton completo (invertendo o Hessiano), o que é computacionalmente caro e instável, os autores argumentam que condicionar os gradientes pela raiz quadrada da métrica de Fisher é suficiente. Isso efetivamente normaliza o gradiente pelo seu desvio padrão (uma transformação de "z-score").
Esparsidade de Gradiente: Nas interfaces de tokens, o gradiente para um token $i$ é não-nulo apenas quando esse token é selecionado. O artigo observa que a energia do gradiente por linha é um estimador não enviesado da probabilidade de participação do token ( $p_i$ ).

Design Algorítmico

O Ember substitui a pegada de memória $O(2VD)$ do Adam por uma abordagem $O(V + D)$ , fatorizando a estimativa do segundo momento:

Remoção do Primeiro Momento: Os autores descobriram que remover o primeiro momento (momentum) do EMA do Adam favorece inteiramente o desempenho e economiza memória.
Segundo Momento Fatorizado: Em vez de manter um buffer de segundo momento completo de $V \times D$ $V \times D$ , o Ember mantém dois vetores 1D:
- $r_t$ : Uma estimativa de segundo momento por linha (rastreando a frequência/energia do token).
- $c_t$ : Uma estimativa de segundo momento por coluna (rastreando a energia por característica).
Pré-condicionamento: A regra de atualização forma um pré-condicionador de produto externo leve. O denominador efetivo para um parâmetro $(i, j)$ é derivado da média geométrica das estimativas de linha e coluna ( $\sqrt{\hat{r}_i \hat{c}_j}$ ).
Correção de Viés: A correção de viés padrão é aplicada às estimativas de linha e coluna antes de formar o pré-condicionador.

A regra de atualização resultante efetivamente divide o gradiente pela estimativa da raiz quadrada da métrica de Fisher, normalizada para ser independente da frequência do token, enquanto reduz a variância para tokens frequentes.

Principais Contribuições

Eficiência de Memória: O Ember reduz a pegada de estado do otimizador para interfaces de tokens de $O(2VD)$ para $O(V + D)$ . Isso elimina a necessidade de particionar os estados do otimizador para tabelas de embedding, permitindo que todo o estado caiba em uma única GPU mesmo para modelos massivos.
Insight Geométrico: O artigo fornece uma visão fundamentada de que otimizadores competitivos para Transformers (incluindo Adam, Lion e Muon) podem ser unificados sob o framework de condicionar gradientes pela raiz quadrada da métrica de Fisher (ou seus análogos de matriz).
Análise de Trajetória de Otimização: Através da Decomposição em Valores Singulares (SVD) das trajetórias de tokens, os autores demonstram que os tokens de embedding otimizam ao longo de um raio 1D surpreendentemente simples (base monomial). Isso desafia a visão prevalecente de que a otimização de redes neurais navega em um panorama altamente não-convexo e caótico, sugerindo, em vez disso, um caminho de convergência linear e suave para as interfaces de tokens.
Implementação Open-Source: Uma implementação distribuída compatível com configurações ZeRO/FSDP existentes é lançada, facilitando a adoção imediata.

Resultados

Os autores fornecem evidências empíricas em múltiplos regimes:

Ajuste Fino Supervisionado (SFT): Em modelos GPT-2 e Pythia, o Ember iguala ou melhora ligeiramente a perda de validação do AdamW, enquanto reduz o estado do otimizador de gigabytes para kilobytes (ex: 2,06 GB $\to$ 413 KB para Pythia-2.8B).
Aprendizado por Reforço (RL): No pós-treinamento do Qwen2.5-3B com GRPO, o Ember alcança paridade com o AdamW em vários tamanhos de lote de rollout, utilizando 4.000 vezes menos estado do otimizador.
Pré-treinamento: No FineWeb, o Ember iguala o AdamW e supera o Adafactor, particularmente em regimes de baixo lote, onde a decomposição de gradiente de rank-1 oferece estabilidade contra gradientes esparsos/ruidosos.
Leis de Escala: O Ember escala efetivamente com o tamanho do modelo e o tamanho do lote. A lacuna de desempenho entre o Ember e o AdamW desaparece à medida que o tamanho do modelo aumenta, sugerindo que o buffer completo de segundo momento do Adam é bem aproximado pela fatoração de rank-1 do Ember em escala.
Estabilidade: O Ember não exibe picos de perda no início do treinamento com lotes pequenos, um modo de falha observado no AdamW.

Significância e Alegações

O artigo afirma que o Ember oferece uma melhoria na "fronteira de Pareto", alcançando desempenho de estado da arte enquanto reduz drasticamente a sobrecarga de memória. A significância é dupla:

Engenharia Prática: Ao remover a necessidade de particionar os estados do otimizador de embeddings, o Ember simplifica configurações de treinamento distribuído e permite iterações mais rápidas em dispositivos únicos para modelos de até 7 bilhões de parâmetros (e além), onde o estado do Adam era anteriormente uma restrição limitante.
Compreensão Teórica: A descoberta de que as trajetórias de otimização de tokens se decompõem em uma base monomial 1D limpa sugere que a natureza "não-convexa" do panorama de perda pode ser um artefato do otimizador ou da parametrização, e não uma propriedade intrínseca da interface de tokens. Isso implica que o espaço de busca para otimizadores ideais para Transformers é mais estreito do que se pensava anteriormente, centrado em torno da raiz quadrada da métrica de Fisher.

Os autores concluem que, embora o Ember seja uma ferramenta prática para treinamento com restrição de memória, seu enquadramento teórico fornece uma nova lente para entender como os LLMs aprendem, destacando especificamente a simplicidade do caminho de otimização para representações de tokens.