Looping de camadas reduz em até 3,3 vezes o custo de treinamento de LLMs de difusão

Os modelos de linguagem de difusão com ruído uniforme (UDLM) há tempos são vistos como uma resposta cara, porém interessante, à autorregressão. Um novo trabalho mostra como reduzir esse custo de forma expressiva. Se algumas das primeiras camadas do transformer forem executadas de forma recursiva dentro de cada passo de denoising, a looped UDLM atinge o validation NLL de uma UDLM convencional treinada com cerca de 0,5 trilhão de tokens, com custo até 2,8 vezes menor no OpenWebText e até 3,3 vezes menor no FineWeb-Edu. O objetivo de treinamento e o sampler não foram alterados.

Por que treinar UDLMs custa caro e por que vale a pena

Os modelos de difusão mascarada (MDLM) trabalham com contexto limpo. Parte dos tokens é substituída por uma máscara, os demais permanecem corretos, e o modelo reconstrói as lacunas. Os modelos uniformes são mais rígidos. Qualquer token da sequência pode ser trocado por um token aleatório do vocabulário comum e só depois corrigido nos passos seguintes. Durante o treinamento, o modelo precisa remover o ruído do texto apoiando-se em um contexto que pode, ele próprio, não ser confiável.

A tarefa fica mais difícil, e por isso as UDLMs exigem mais FLOPs do que as MDLMs. Em troca, os autores apontam resultados mais fortes em tarefas downstream ligadas a raciocínio. Poder revisar um token já escrito dá ao modelo a chance de corrigir o próprio erro, em vez de carregá-lo até o fim da geração. A questão é o custo, e é justamente ele que o trabalho tenta reduzir.

Um loop no lugar de novos parâmetros

A ideia do selective looping já havia mostrado utilidade em MDLMs com o mesmo orçamento de FLOPs. Algumas das primeiras camadas do transformer são aplicadas de forma recursiva dentro de um único passo de denoising. Os pesos das camadas continuam os mesmos, mas a profundidade efetiva do cálculo aumenta. Os autores levaram a abordagem para as UDLMs e não mexeram em mais nada: nem na função de perda, nem no procedimento de amostragem.

Principais números do trabalho:

• Treinamento: a looped UDLM atinge o validation NLL da UDLM base treinada com ~0,5 trilhão de tokens, com custo até 2,8x menor no OpenWebText.
• No FineWeb-Edu, o ganho é maior: até 3,3x menos FLOPs.
• Geração: com o mesmo orçamento de NFE (número de avaliações da rede por amostra), a generation perplexity é até 9,8% menor do que a do modelo sem loops.
• A maior parte da melhoria vem de um número moderado de loops; aumentar além disso traz ganhos bem menores.
• Distribuição do número de loops no treinamento: deslocá-la para valores baixos reduz em 12% o número de loops na inferência, em relação à amostragem uniforme, sem queda de precisão nas tarefas downstream.

Vale destacar que o que conta é a combinação das métricas. O validation NLL mostra a qualidade do treinamento, enquanto a generation perplexity com NFE fixo mostra a qualidade da geração em si. O ganho aparece nas duas, e a verificação em tarefas downstream confirma que ele não se limita a métricas convenientes para o modelo.

Quantos loops usar e para quem isso serve

A conclusão prática é bem pé no chão. O retorno não cresce de forma linear com o número de loops, então não faz sentido perseguir a recursão máxima. Um valor moderado entrega a maior parte da melhoria; depois disso, o custo de cada passo cresce mais rápido que o benefício. Em modelos de difusão, nos quais a geração já exige muitas iterações de denoising, isso pesa ainda mais: cada loop extra é multiplicado pelo número de passos.

O segundo ponto prático envolve a distribuição do número de loops. A amostragem uniforme durante o treinamento parece a opção natural, mas valores baixos mais frequentes deixam o modelo robusto a um modo de operação reduzido. O resultado é que a inferência precisa de 12% menos loops com a mesma precisão. Para equipes que constroem LLMs de difusão, é uma receita pronta de ajuste. Ela não exige mudar nada no sampler, então pode ser incorporada a um pipeline existente.

O maior interesse do trabalho está em outro ponto. As UDLMs eram mais fortes em raciocínio, mas perdiam em custo de treinamento. Se o looping reduz o custo de forma consistente em 2,8–3,3 vezes, a diferença de preço para as MDLMs diminui e o argumento a favor do ruído uniforme fica mais convincente. A questão da escala além do OpenWebText e do FineWeb-Edu continua em aberto. Mas a direção é clara: a qualidade dos modelos de linguagem de difusão pode ser obtida com arquitetura, sem depender só de mais tokens e horas de GPU.