Transformer em loop de 1B de parâmetros saiu 1,4 vez mais barato em FLOP

Os transformers em loop se tornaram um tema da moda este ano: quase todo dia saem artigos. A ideia é simples: em vez de adicionar novas camadas, o modelo passa várias vezes pelas mesmas. Os parâmetros são os mesmos, mas ele "pensa" mais fundo. Um novo trabalho de 149 páginas, assinado por pesquisadores entre os quais Mike Lewis e Artidoro Pagnoni, finalmente responde a uma pergunta prática: como exatamente fazer o loop para que valha a pena.

Qual é o problema do loop

Imagine uma linha de produção com 24 máquinas. Dá para comprar mais 24 e obter um processamento mais complexo, mas isso é caro em memória e parâmetros. Ou dá para mandar a peça por uma segunda volta nas mesmas máquinas. A memória necessária é a mesma, mas há mais computação. Essa abordagem tem dois pontos fracos. O primeiro: não está claro o que repetir, todas as camadas ou só parte delas. O segundo: ao gerar texto, cada novo token espera o término de todos os loops do anterior, e paralelizar essa decodificação é difícil.

O que os autores descobriram

A equipe pré-treinou decodificadores densos de até 1B de parâmetros e até 32 loops e comparou dois modos. O Full-Stack Looping repete todas as camadas. O Middle-Stack Looping repete apenas as intermediárias, e as primeiras e as últimas ficam como estão.

A conclusão depende do que lhe falta. Se o número de parâmetros é estritamente limitado, é melhor repetir todas as camadas. Se o limitado é o poder de computação (FLOP), as camadas intermediárias levam vantagem. Para um engenheiro, isso se transforma imediatamente em uma regra de escolha de arquitetura.

A segunda parte do trabalho trata da decodificação paralela. No esquema Parallel Loop Transformers (PLT), os estados dos loops são deslocados para os tokens vizinhos, de modo que diferentes loops são calculados simultaneamente. Os autores notaram que, com vários loops, a qualidade se degrada, porque o embedding do token atual e o estado do anterior são misturados em uma única soma. A correção deles, o Normalized Shifted PLT, normaliza essas duas entradas separadamente via RMSNorm antes da mistura. Além disso, eles adicionaram adaptadores compartilhados de baixo posto para os loops (Loop Matrix), em essência uma correção semelhante ao LoRA, para que cada passada seja um pouco diferente da anterior.

Os números são estes. No conjunto DCLM com 200M de parâmetros, a versão de quatro loops reduz os bits por byte para ~1,060, contra ~1,087 do PLT original e ~1,088 do modelo sem loops. Com 1B de parâmetros, o modelo de dois loops chega a ~0,86 contra ~0,90 com 80 tokens por parâmetro. A precisão em oito tarefas foi de ~47,5% contra ~46,1% com 160 tokens por parâmetro. Com treinamento longo, os modelos de dois loops exigem até 1,4 vez menos FLOP do que os modelos sem loops.

O que fazer com isso

Se você não treina modelos do zero, não há nada a fazer agora. Mas se planeja um modelo compacto próprio para um dispositivo ou servidor com limite rígido de memória, o loop se torna uma alavanca real. A memória necessária é a mesma, e a qualidade cresce. O esquema de escolha é simples: se o gargalo são os parâmetros, use o loop completo; se são os FLOP, repita o meio. Se a velocidade de geração importa, olhe para o deslocamento normalizado: é ele que faz a decodificação paralela funcionar.

Limitações

Há três ressalvas importantes. Primeiro, os experimentos pararam em 1B de parâmetros, e não se sabe se as tendências se mantêm em modelos de dezenas de bilhões. Segundo, os autores não mediram o tempo real de geração. Economizar FLOP ainda não significa que a resposta chegará mais rápido: na prática, memória e kernels decidem muito. Terceiro, o ganho de precisão é pequeno, cerca de 1,4 ponto percentual, então isso é mais uma otimização de engenharia do que uma revolução.

Em suma: o trabalho é útil como um mapa do terreno. Ele mostra que o loop não é uma única técnica, mas um conjunto de decisões com diferentes compromissos, e sugere qual escolher para o seu orçamento. Por enquanto, é material para quem treina modelos por conta própria, mas a tendência claramente aponta para extrair mais de cada parâmetro.