STAPLE adiciona bilhões de parâmetros ao OLMo-2 e decodifica duas vezes mais rápido que MoE

Uma equipe de pesquisadores publicou no arXiv o STAPLE, um método que aumenta o número de parâmetros de um modelo de linguagem já treinado sem exigir um pré-treinamento esparso caro, feito do zero. No OLMo-2 1B, ele entrega +4,27 pontos no HumanEval, +2,80 no MBPP, +2,27 no GSM8K e +1,65 no DROP em relação à simples continuação do treinamento dense. Segundo os autores, a geração fica no máximo 1,5% mais lenta e o treinamento, no máximo 3%, mesmo com bilhões de parâmetros incluídos tardiamente no modelo.

Por que não MoE

A computação esparsa é vista há tempos como uma forma de aumentar a capacidade do modelo sem inflar o custo computacional ativo por token. O problema é o preço: o pré-treinamento de arquiteturas esparsas exige muitos dados e muito dinheiro. Os métodos de upcycling existentes, que convertem um checkpoint dense pronto em um modelo esparso, miram principalmente Mixture-of-Experts. Os autores apontam três gargalos nessa abordagem: homogeneidade dos especialistas, baixa sample efficiency e overhead de sistema com roteamento e memória.

A memória esparsa estática funciona de outra forma. O acesso a ela é previsível: os parâmetros são escolhidos pelo identificador do token, e não pela decisão de um roteador. As atualizações também ficam atreladas a cada token. Mas, como observam os autores, inserir essa memória de forma ingênua em um modelo pré-treinado gera um retorno fraco.

Como o STAPLE funciona

O método se baseia no STEM, que substitui as projeções up dos blocos FFN por embeddings indexados por token. Em treinamento do zero isso funciona, mas, em um checkpoint pronto, a substituição quebra a computação já acumulada. O STAPLE faz diferente: as projeções up continuam onde estão, e os embeddings entram como um caminho residual. A parte pré-treinada segue operando como antes, enquanto a nova memória aprende aquilo que o modelo conhece pior sobre os tokens: informação adicional específica de cada token.

Há um detalhe importante de engenharia: a implementação faz offload dos embeddings para o host. Apenas cerca de 5% dos embeddings mais frequentes ficam na GPU, e o restante permanece na memória do host. Por causa da distribuição dos tokens, isso basta para que o throughput quase não caia.

Os números

Principais resultados do preprint:

• OLMo-2 1B, ganho sobre o continued dense training: HumanEval +4,27, MBPP +2,80, GSM8K +2,27, DROP +1,65.
• A média cresce junto com a capacidade estática, e as melhorias se mantêm no OLMo-2 7B.
• O STAPLE supera as baselines esparsas estáticas: STEM e um modelo no estilo Engram.
• Contra o upcycling para MoE: média mais alta, decodificação 2 vezes mais rápida, menos da metade da memória de GPU para o modelo e menos energia por token.
• Overhead: geração dentro de 1,5% da base dense, treinamento dentro de 3%.

A ressalva é importante: a comparação com MoE em velocidade e energia foi feita em uma avaliação focada em throughput de decodificação. Em cenários com requisitos rígidos de latência para uma única requisição, o resultado pode ser diferente. Além disso, trata-se ainda de um preprint com modelos de 1B e 7B, e o escalonamento para dezenas e centenas de bilhões de parâmetros não foi demonstrado.

O que isso significa na prática

Para equipes que fazem fine-tuning de modelos abertos, o STAPLE oferece um caminho barato para ganhar capacidade: pega-se um checkpoint existente, adiciona-se a memória estática, e o modelo continua treinando com uma pequena perda de velocidade. O ganho é mais visível em código e matemática, áreas em que o conhecimento preciso de tokens e padrões é especialmente útil. Na inferência, o perfil de memória chama atenção: a maior parte dos parâmetros fica no host, e a GPU precisa de apenas um cache pequeno. Isso reduz a exigência de VRAM em comparação com um MoE de qualidade equivalente.

Antes de levar a abordagem para produção, vale checar três pontos: o comportamento em modelos maiores que 7B, a robustez a mudanças de domínio, em que a distribuição dos tokens frequentes é outra, e a latência com batches pequenos. Se esses testes forem bem-sucedidos, o upcycling via embeddings estáticos passará a ser uma alternativa séria ao MoE para expandir modelos já treinados.