Duplicação de camadas no treinamento elevou o GSM8K de 12% para 19% em modelo de 3B

Os pesquisadores Ferdinand Kapl, Kaitlin Maile, Stefan Bauer e Johannes von Oswald compararam três formas de construir um modelo de linguagem de até 3 bilhões de parâmetros. A primeira é a rede densa convencional. A segunda é uma rede que fica gradualmente mais profunda durante o treinamento, por meio da cópia de camadas intermediárias (método LIDAS).

A terceira é uma rede "em loop", em que as mesmas camadas são executadas várias vezes em um único passe. O foco foi o in-context learning (ICL): a capacidade de o modelo entender a tarefa a partir de alguns exemplos no prompt, sem ajuste fino. O resultado: os modelos com aumento de profundidade e com looping aprendem ICL de forma nitidamente melhor que o modelo denso.

Na escala de 3B, o LIDAS atingiu cerca de 19% de correspondência exata no GSM8K com 8 exemplos no prompt, contra cerca de 12% do modelo denso.

Duas formas de reaproveitar profundidade

Em um modelo convencional, o número de camadas é definido de antemão e cada camada é treinada do zero. No LIDAS é diferente: o treinamento começa com uma rede rasa, depois as camadas intermediárias são duplicadas e a rede cresce ao longo do processo. No início, as cópias são quase idênticas; depois, elas divergem. A fase inicial do treinamento acontece em um modelo pequeno e, por isso, sai mais barata. Segundo os autores, no 3B o LIDAS gasta 26% menos FLOPs (operações de ponto flutuante) no treinamento e, ainda assim, supera o modelo denso.

O looping é mais simples. Pega-se um conjunto pequeno de camadas, e os dados passam por ele várias vezes em um único forward pass. O resultado é um número menor de parâmetros únicos e uma profundidade efetiva maior.

Os experimentos usaram modelos de 360M a 3B de parâmetros. O código foi construído sobre a biblioteca nanotron, da Hugging Face. Os autores não publicam código próprio nem pesos.

Onde os modelos se saem melhor

O GSM8K é um conjunto de problemas escolares de matemática em texto. Oito exemplos resolvidos entram no prompt, e verifica-se se a resposta do modelo coincide com o gabarito. Resultados dos modelos de 3B:

• LIDAS: cerca de 19% de correspondências exatas;
• modelo denso: cerca de 12%.

O segundo teste é o RegBench: linguagens formais novas, que o modelo nunca viu. A partir de alguns exemplos, ele precisa captar as regras e prever o próximo símbolo. A qualidade é medida pela distância de variação total, que mostra a diferença entre as probabilidades previstas e as reais. Quanto menor, melhor:

• aumento de profundidade: ~0.119;
• looping: ~0.125;
• modelo denso: ~0.137.

Aqui é ICL em estado puro: no teste não há nada para recuperar dos pesos, só resta ler o contexto.

Induction heads multiplicadas e o preço do ganho

Os autores explicam o resultado pelas induction heads (cabeças de indução). São cabeças de atenção que localizam no contexto uma ocorrência anterior do token atual e sugerem o que veio depois dela. Elas são consideradas um dos principais mecanismos do ICL.

No modelo de 3B com LIDAS, 10 das 12 induction heads mais fortes descendem de dois ancestrais da rede inicial de quatro camadas. O circuito encontrado no modelo pequeno é copiado por toda a profundidade conforme a rede cresce. Parece que ele não é substituído depois pelo conhecimento que a rede memoriza nos pesos.

Outra verificação foi a ablação. Os autores desligaram os 20% principais dessas cabeças e observaram quanto cai o F1 médio em 19 tarefas com poucos exemplos:

• modelo denso: queda de cerca de 0.08;
• LIDAS: cerca de 0.18;
• modelo em loop: cerca de 0.28.

Ou seja, os modelos expandidos e em loop dependem mais desses circuitos, e toda a capacidade de aprender com exemplos fica concentrada neles.

O que os engenheiros devem considerar

Para quem faz pré-treinamento de modelos, a conclusão é prática. Se o objetivo é economizar computação e ter bom ICL, o aumento gradual de profundidade parece uma opção razoável. No 3B, ele trouxe melhores resultados e 26% de economia de FLOPs.

O looping tem um custo maior. Esse modelo tem menos parâmetros únicos e se sai pior ao extrair conhecimento armazenado nos pesos. O ganho em ICL tem um preço: o modelo lê melhor o contexto, mas lembra pior os fatos.

Há também uma ressalva sobre a confiabilidade. Cada configuração, de 360M a 3B, foi treinada com apenas um seed, ou seja, um único estado inicial do gerador de números aleatórios. A variação entre execuções não foi medida. Sete pontos percentuais no GSM8K parecem convincentes, mas, sem repetições, não dá para dizer o quanto são estáveis. Além disso, 3B é um tamanho pequeno para os padrões dos modelos atuais, e a extrapolação das conclusões para dezenas de bilhões de parâmetros precisa ser verificada à parte.

Ainda assim, o resultado é inesperado. O modelo denso, treinado da forma habitual, ficou atrás justamente onde era tido como forte: na capacidade de aprender com exemplos do contexto. A explicação dos autores é a cópia precoce dos circuitos de atenção, e ela pode ser testada em outros tamanhos de modelo.