Test-time training surge como novo eixo de escala para modelos

A promessa de um novo eixo de escala

O autor parte de uma observação sobre o desenvolvimento de modelos: cada nova forma de escalar trouxe ganhos grandes de eficácia. Numa visão simplificada, os GPTs 1 a 3 escalaram dados, arquiteturas MoE escalaram parâmetros separadamente dos FLOPs, e tudo desde o o1 também escala compute no momento do teste, por meio de raciocínio.

Nesse contexto, o test-time training (TTT) parece atraente como um novo eixo de escala. Ele também dialoga com o tema do aprendizado contínuo: os modelos atuais são muito inteligentes, mas para evitar que repitam o mesmo erro é preciso incluir avisos como "IMPORTANT: DO NOT IGNORE" nos prompts. Se eles pudessem aprender, isso não seria necessário.

O que os artigos mostram

Ao comentar o artigo "Test-Time Training with KV Binding Is Secretly Linear Attention", o autor nota que ele mostra que uma ampla classe de arquiteturas TTT pode ser expressa como uma forma de operador de atenção linear aprendido - o que pode parecer uma reformulação de algo já existente. Ele considera o artigo interessante, com conclusões úteis para quem constrói certos tipos de modelos, mas o que mais lhe chamou atenção foi a confusão terminológica em torno de "test time training".

O artigo que define o termo, "Test Time Training under Distribution Shift", mostra o esperado: atualizar os pesos do modelo para lidar com mudanças de distribuição vistas no momento do teste. São atualizações de slow weights - pesos atualizados pelo otimizador, uma vez por minibatch, com checkpoints, ou seja, o que normalmente se entende por "pesos".

Já os fast weights são produzidos no forward pass, como a matriz de estado de uma camada de atenção linear. O treino normal não otimiza esses pesos diretamente; aprende uma regra que os cria.

KV binding como fast weights

No caso analisado, o TTT treina fast weights: uma camada de KV binding que adiciona uma matriz de tamanho fixo, que não cresce com o comprimento da sequência, e roda um passo de gradiente atualizando essa matriz a cada elemento da sequência. O autor observa que isso lembra uma camada de atenção linear ou uma RNN.

Isso resolve o aprendizado contínuo?

Não, segundo o texto: os fast weights são reiniciados a cada nova sequência, então o mecanismo não proporciona aprendizado persistente entre sequências.