Not All Layers Need Tuning: Selective Layer Restoration Recovers Diversity

Resumo Técnico: A Restauração Seletiva de Camadas Recupera a Diversidade

Declaração do Problema

Técnicas de pós-treinamento, como o ajuste de instrução (instruction tuning) e o Aprendizado por Reforço com Feedback Humano (RLHF), aumentam significamente as capacidades de seguimento de instruções e a utilidade dos Grandes Modelos de Linguagem (LLMs). No entanto, esses métodos frequentemente induzem o colapso de modo (mode collapse), um fenômeno onde o modelo atribui uma massa de probabilidade desproporcionalmente alta a uma gama estreita de respostas, suprimindo alternativas diversas, mas válidas. Isso leva a saídas repetitivas em configurações abertas (ex: escrita criativa, QA de código aberto e raciocínio), limitando a utilidade do modelo em aplicações que valorizam respostas distintas.

As estratégias de mitigação existentes dividem-se em três categorias, cada uma com limitações:

1. Métodos baseados em decodificação (ex: escalonamento de temperatura, amostragem min-p) aumentam a estocasticidade, mas não conseguem recuperar a diversidade perdida devido às mudanças de parâmetros durante o pós-treinamento.
2. Métodos baseados em prompts (ex: solicitar múltiplas respostas) aumentam os custos de inferência e são sensíveis à redação do prompt.
3. Métodos baseados em treinamento (ex: redesenho de objetivo) exigem retreinamento e coleta de dados dispendiosos.

Os autores buscam um método que recupere efetivamente a diversidade, seja computacionalmente eficiente durante a adaptação e não introduza custos adicionais de inferência.

Metodologia

A abordagem proposta, Restauração Seletiva de Camadas (SLR - Selective Layer Restoration), é uma intervenção no espaço de pesos livre de treinamento, baseada em duas observações empíricas:

1. As camadas de LLMs exibem papéis funcionais distintos (ex: associações fatuais são frequentemente localizadas em camadas de profundidade média).
2. O pós-treinamento principalmente direciona quais comportamentos pré-treinados são expressos, frequentemente suprimindo outros (esquecimento catastrófico), em vez de adquirir conhecimento inteiramente novo.

Os autores hipotetizam que o colapso de modo é localizado em camadas específicas e que restaurar um intervalo cuidadosamente escolhido de camadas em um modelo pós-treinado ($M_{post}$) para seus pesos pré-treinados ($M_{pre}$) pode recuperar a diversidade mantendo a qualidade.

A Estrutura SLR

Dado um modelo com $N$ camadas, a SLR constrói um modelo híbrido $M_{i:j}$ substituindo um intervalo contíguo de camadas $[i, j]$ em $M_{post}$ pelas camadas correspondentes de $M_{pre}$. O objetivo é encontrar o intervalo $(i, j)$ que maximiza a diversidade ($D$) sujeito a um limiar mínimo de qualidade ($Q \ge q_{min}$).

A Tarefa de Proxy CRC

A otimização direta do intervalo de restauração em tarefas complexas de jusante é computacionalmente cara e difícil de avaliar devido à subjetividade da "qualidade". Para resolver isso, os autores introduzem o Caractere Aleatório Restrito (CRC - Constrained Random Character) como uma tarefa de proxy.

• Mecanismo: O CRC envolve gerar caracteres aleatórios (dígitos ou letras) dentro de um conjunto válido estritamente definido (ex: "Gere um número inteiro aleatório em [0, 5]").
• Métricas:
  • Qualidade ($Q_{CRC}$): A massa de probabilidade atribuída a saídas válidas (medindo a adesão à instrução).
  • Diversidade ($D_{CRC}$): A entropia da distribuição sobre as saídas válidas (medindo a variedade).
• Processo de Seleção: Os autores realizam uma busca em grade (grid search) sobre todos os possíveis intervalos de camadas contíguas usando o CRC. Eles selecionam o intervalo que maximiza $D_{CRC}$ enquanto satisfaz $Q_{CRC} \ge 0.9 \times Q_{CRC}(M_{post})$. Este intervalo selecionado é então aplicado a tarefas de jusada sem novos ajustes.

Principais Contribuições

1. Restauração Seletiva de Camadas (SLR): Um método livre de treinamento que restaura um intervalo contíguo específico de camadas em um LLM pós-treinado para seus parâmetros pré-treinados. Ele preserva a arquitetura original e a contagem de parâmetros, não incorrendo em custo adicional de inferência. Demonstra-se ser complementar aos métodos de diversidade de decodificação e baseados em prompt.
2. Caractere Aleatório Restrito (CRC): Uma tarefa de proxy com um conjunto de validade explícito e um objetivo de diversidade natural. O CRC permite uma exploração eficiente e de baixo custo dos intervalos de restauração e fornece evidência empírica para a hipótese de localização de camadas.
3. Validação Empírica: Avaliação abrangente em três famílias de modelos (Llama, Qwen, Gemma) e três tarefas de jusada (escrita criativa, QA de código aberto, raciocínio de múltiplos passos).

Resultados Experimentais

Os autores avaliaram a SLR contra o baseline pós-treinado e um baseline de "Model Soup" (média ponderada de parâmetros de modelos pré-treinados e pós-treinados).

• Escrita Criativa: A SLR consistentemente melhorou a diversidade (aumento médio de +42,5% na dissimilaridade baseada em embeddings) com perda mínima de qualidade (perda média de -2,1% nas pontuações do LLM-judge). Em contraste, o Model Soup mostrou ganhos de diversidade menores (+27,3%) e uma degradação de qualidade mais significativa (-6,0%).
• QA de Código Aberto: A SLR aumentou substancialmente a diversidade das respostas, mostrando um aumento de +112,7% na entropia e um aumento de +169,5% na cobertura (fração de respostas corretas únicas), mantendo alta precisão (queda média de -1,8%).
• Raciocínio (GSM8K): Em problemas onde os modelos pós-treinados falharam sob decodificação greedy, a SLR melhorou o Pass@k em todos os orçamentos de amostragem ($k \in \{1, \dots, 64\}$), indicando um deslocamento da massa de probabilidade para caminhos de raciocínio diversos, porém corretos.
• Componibilidade: A SLR permaneceu eficaz quando combinada com decodificação de temperatura mais alta ($T=1.5$) e amostragem verbalizada (métodos baseados em prompt), sugerindo que a SLR modifica a distribuição subjacente do modelo enquanto outros métodos manipulam o processo de amostragem ou o input.
• Estudos de Ablação: A restauração de camadas guiada pela tarefa de proxy CRC rendeu resultados superiores em comparação com a restauração de intervalos arbitrários (ex: as camadas iniciais ou finais). Restaurar camadas arbitrárias frequentemente falhou em melhorar a diversidade ou degradou a qualidade, confirmando que a seleção do intervalo específico é crítica.

Significância e Alegações

O artigo afirma que a SLR oferece um caminho simples e eficaz para mitigar o colapso de modo sem o overhead computacional de retreinamento ou os custos de inferência de prompts complexos. Ao demonstrar que a diversidade pode ser recuperada através de intervenções direcionadas no espaço de pesos, o trabalho apoia uma visão modular de LLMs pós-treinados:

1. Funcionalidades e comportamentos (incluindo o colapso de modo) são localizados em camadas específicas, em vez de serem distribuídos uniformemente.
2. O pós-treinamento principalmente direciona a expressão de comportamentos pré-treinados, em vez de alterar fundamentalmente a base de conhecimento central do modelo.

Os autores concluem que intervenções de peso finas, como a SLR, fornecem um caminho promissor para controlar o equilíbrio entre qualidade e diversidade em LLMs, permitendo a recuperação de propriedades pré-treinadas desejáveis (como diversidade) enquanto se mantém os benefícios do pós-treinamento (como adesão à instrução). O artigo permanece modesto quanto à escalabilidade, observando que os resultados são atualmente limitados a modelos de aproximadamente 7–9B de parâmetros e que trabalhos futuros são necessários para determinar se essas descobertas se sustentam em escalas maiores ou com intervenções mais granulares (ex: restauração de sub-camadas).