TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Resumo Técnico: TAPIOCA – Por que a Poda Consciente da Tarefa Melhora a Capacidade de Modelos OOD

Declaração do Problema
Trabalhos recentes demonstraram que remover camadas de Modelos de Linguagem de Grande Escala (LLMs) no momento da inferência — especificamente por meio de métodos de poda consciente da tarefa como TALE (Eliminação de Camada Consciente da Tarefa) — pode melhorar o desempenho em tarefas específicas sem retreinamento. No entanto, os mecanismos que impulsionam esses ganhos permanecem obscuros. Uma questão central é saber se essas melhorias decorrem de benefícios gerais de compressão, efeitos de regularização ou interações específicas entre as representações internas do modelo e a distribuição de entrada. Além disso, é desconhecido se esses ganhos são universais ou específicos de certos tipos de deslocamento de dados.

Metodologia
Os autores investigam esse fenômeno usando uma combinação de experimentos sintéticos controlados e avaliações em larga escala em LLMs ajustados finamente.

1. Configuração de Regressão Controlada: O estudo começa com um transformer de 12 camadas treinado do zero em regressão linear em contexto. A distribuição de treinamento é explicitamente definida (coeficientes amostrados de $U(-1, 1)$). Os autores avaliam o modelo tanto em dados dentro da distribuição (ID) (mesma distribuição) quanto em dados fora da distribuição (OOD) (coeficientes amostrados de intervalos mais amplos $U(-\sigma, \sigma)$ onde $\sigma > 1$). Eles aplicam o TALE, um algoritmo de poda ganancioso que remove camadas para maximizar a precisão de validação em uma distribuição específica.
2. Avaliação em Larga Escala de LLMs: As descobertas são estendidas para LLMs ajustados finamente (Llama 3.1 8B e GPT-OSS 120B) especializados em raciocínio matemático e geração de código. Aqui, a distribuição "adaptada" é a tarefa de ajuste fino (por exemplo, MATH500), e as tarefas OOD incluem benchmarks gerais como MMLU e BoolQ.
3. Análise Geométrica e Causal: Para explicar por que a poda funciona, os autores analisam a geometria dos estados ocultos. Eles medem as normas representacionais por camada e as distâncias entre pares de tokens. Eles também empregam substitutos lineares locais para aproximar transformações de camada e realizam intervenções causais, como a inserção de mapas substitutos inversos e o redimensionamento de atualizações residuais, para isolar o efeito da distorção geométrica.

Principais Descobertas

• Ganhos Específicos de OOD: A principal descoberta empírica é que a poda consciente da tarefa produz ganhos de desempenho apenas sob deslocamento de distribuição. Quando avaliado na distribuição à qual o modelo foi adaptado (ID), o TALE falha em melhorar o desempenho ou degrada-o ativamente ao remover camadas necessárias para precisão ótima. Por outro lado, em dados OOD, o TALE identifica e remove consistentemente camadas que melhoram a precisão.
• Incompatibilidade da Geometria de Representação: Os autores demonstram que entradas OOD induzem uma geometria de representação distorcida em comparação com entradas ID. Especificamente, entradas OOD causam que as normas de estados ocultos por camada e as distâncias entre pares se desviem sistematicamente dos perfis observados em dados ID. No cenário de regressão, entradas OOD causam a inflação de distâncias através das camadas; em LLMs, trajetórias OOD divergem de trajetórias ID nas camadas do meio até o final.
• Mecanismo de Amplificação: O estudo identifica que certas camadas atuam como "amplificadores dependentes da distribuição". Em entradas ID, essas camadas realizam refinamentos úteis. No entanto, em entradas OOD, as mesmas camadas exibem amplificação de alto ganho e baixo rank, empurrando representações ainda mais para longe da geometria adaptada à tarefa. O TALE melhora o desempenho ao remover essas camadas específicas, prevenindo assim a amplificação da distorção geométrica.
• Alinhamento, Não Apenas Contração: A poda não suprime simplesmente as normas de ativação. Em experimentos onde a direção do deslocamento de distribuição é invertida (treinamento em distribuições amplas, teste em distribuições estreitas), a poda expande as distâncias de representação para alinhar com o perfil ID. O mecanismo central é o realinhamento das representações OOD em direção à geometria aprendida durante a adaptação, independentemente de isso exigir contração ou expansão.
• Evidência Causal: Intervenções causais confirmam que a distorção geométrica é o motor da perda de desempenho. A inserção de um mapa substituto inverso após uma camada podada (para desfazer o efeito geométrico da camada) restaura o desempenho OOD quase ao nível do modelo podado. Da mesma forma, o redimensionamento da atualização residual de uma camada selecionada melhora continuamente a precisão à medida que a contribuição da camada é atenuada, descartando artefatos topológicos ou simples redundância como a única causa.

Significado e Alegações
O artigo alega fornecer uma explicação geométrica novel para o sucesso da poda no momento da inferência. Ele argumenta que a importância da camada não é fixa, mas condicional tanto à tarefa quanto à distribuição de entrada. Os autores postulam que modelos pré-treinados codificam múltiplas funções especializadas em tarefas, cada uma associada a uma geometria de representação local. Entradas ID se encaixam nessa geometria, enquanto entradas OOD a violam.

O significado deste trabalho reside em:

1. Clarificando os Limites da Poda: Estabelece que a poda consciente da tarefa não é um impulsionador universal de desempenho, mas um mecanismo corretivo para deslocamento de distribuição.
2. Interpretação Geométrica: Desloca a compreensão da redundância de camadas de contagens estáticas de parâmetros para distorções geométricas dinâmicas e dependentes da distribuição.
3. Insight Mecanístico: Fornece evidência causal de que os ganhos de desempenho surgem da correção da trajetória das representações internas para corresponder à geometria adaptada do modelo, em vez de aumentar a capacidade do modelo ou reduzir ruído.

Os autores concluem que a poda melhora o desempenho OOD ao atenuar camadas que amplificam a incompatibilidade entre a distribuição de entrada e a geometria de tarefa aprendida pelo modelo, realinhando efetivamente o estado interno do modelo com as condições sob as quais ele foi otimizado.