Pesquisadores encontram em Pythia uma janela de 0,4–1,1% do pré-treino em que o readout é reorganizado

Costuma-se considerar que o modelo aprendeu uma habilidade no momento em que começa a gerar o token correto. Um novo estudo, destacado no Hugging Face Daily Papers, mostra que essa lógica introduz um desvio. Os estados ocultos podem carregar informação sobre a tarefa antes de a saída do modelo expressá-la. Usando a Pythia como exemplo, os autores encontraram uma janela estreita, de 0,4–1,1% do caminho pelo pré-treino, em que a matriz de unembedding (readout) é reorganizada e as categorias do WordNet aparecem em suas linhas.

A saída também aprende

A previsão de token passa pela matriz de unembedding, que converte o estado oculto em logits sobre o vocabulário. Essa matriz é treinada ao mesmo tempo que todo o resto. Por isso, o gráfico comportamental, em que a acurácia sobe a partir do passo N, mistura dois processos: o que o modelo já codificou internamente e o que o seu readout aprendeu a expressar externamente.

Os autores testaram isso diretamente. Probes lineares nos estados ocultos de checkpoints iniciais já decodificam informação sobre a tarefa, embora as saídas do próprio modelo ainda não a mostrem. O resultado é uma defasagem entre o conhecimento e a sua manifestação, e é ela que distorce a datação de muitas "transições" no desenvolvimento do modelo.

Parameter-trajectory crosscoding

Para acompanhar o readout, os pesquisadores propuseram o parameter-trajectory crosscoding. O método treina um único dicionário esparso sobre as linhas do readout de cada token em cada checkpoint. Os crosscoders clássicos trabalham com ativações e exigem passar dados pelo modelo. Aqui, bastam os pesos dos checkpoints, o que simplifica bastante a análise de modelos abertos com histórico de treinamento preservado.

Principais observações:

• janela de reorganização do readout na Pythia: 0,4–1,1% do pré-treino;
• nessa janela, as categorias do WordNet passam a ser decodificáveis a partir das linhas do readout;
• ao final dela, os probes nos estados ocultos já leem informação que não está nas saídas;
• a análise exige apenas os pesos dos checkpoints, sem ativações;
• com learning rate mais alta no readout, o loss final quase não muda.

Experimento de controle com a learning rate

Correlação não basta, então os autores fizeram uma intervenção. Eles treinaram do zero modelos pequenos que diferiam apenas na learning rate do readout. Se a defasagem é determinada pela velocidade de aprendizado da camada de saída, ela deve se mover junto com essa velocidade.

Foi o que aconteceu. Com learning rate maior no readout, ele se forma mais cedo e a distância entre o estado oculto e a saída se fecha mais rápido. O loss final quase não é prejudicado. Ou seja, a velocidade de aprendizado do readout influencia quando a capacidade se torna visível por fora, mas não aquilo a que o modelo chega no fim do treinamento.

Como usar o resultado na prática

A conclusão diz respeito à fase inicial do treinamento, em que muitas transições "súbitas" são datadas. Se você constrói curvas de surgimento de capacidades, uma métrica baseada só nas saídas não é suficiente. Vale colocar ao lado probes lineares nos estados ocultos e comparar as datas: se o probe se antecipa à saída, você provavelmente está medindo a maturação do readout, e não o surgimento da habilidade.

Para quem treina os próprios modelos, é também um motivo para testar uma learning rate separada para o unembedding. O trabalho foi verificado em modelos pequenos e na Pythia, então é preciso cautela ao estender as conclusões para arquiteturas modernas de grande porte. Mas, para checkpoints abertos, a ferramenta já é aplicável, porque os pesos estão disponíveis para todos.