From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Resumo Técnico: De Camadas para Submódulos (SUBFIT)

Definição do Problema
A compressão pós-treinamento de Grandes Modelos de Linguagem (LLMs) visa reduzir os custos de implantação através da remoção de componentes arquiteturais. Embora métodos que deletam camadas inteiras ou blocos contíguos sejam eficazes, eles interrompem o fluxo residual, frequentemente causando uma degradação acentuada na qualidade. Abordagens recentes baseadas em "substituição" tratam disso substituindo os componentes removidos por módulos leves ajustados. No entanto, os métodos existentes compartilham duas restrições de design rigorosas:

1. Granularidade de camada completa: Eles tratam toda a camada Transformer como a unidade atômica de remoção.
2. Seleção contígua: Eles removem componentes em intervalos de profundidade consecutivos.

Os autores argumentam que essas restrições são excessivamente limitantes. A redundância em transformers pré-treinados não se limita a regiões contíguas, nem é distribuída uniformemente entre as saídas de Atenção e FeedForward (FFN). Consequentemente, diferentes estratégias são necessárias para aproximar diferentes tipos de submódulos, e os componentes removíveis não precisam se agrupar em intervalos de profundidade contíguos.

Metodologia: SUBFIT
O artigo apresenta o SUBFIT (Submodule-level Fitted residual replacement), um método de compressão pós-treinamento que opera na granularidade de submódulos individuais de Atenção e FFN. O processo envolve duas etapas principais:

1. Seleção Não Contígua:

   • O SUBFIT seleciona um número alvo de submódulos de Atenção e FFN de forma independente ao longo da profundidade do modelo com base em um orçamento de esparsidade $s$.
   • A seleção é não contígua; os módulos são escolhidos independentemente por pontuação, sem formar um bloco consecutivo.
   • Estratégia de Pontuação:
     • Atenção: Os módulos são classificados por uma "pontuação de impacto" calculada em ativações de calibração, medindo o quanto um submódulo rotaciona o fluxo residual e sua magnitude relativa.
     • FFN: Para modelos base, utiliza-se uma pontuação sensível à substituição que combina efeitos de remoção e erro substituto. Para modelos ajustados por instrução (instruction-tuned), utiliza-se uma pontuação apenas de cosseno para evitar ruído de magnitudes residuais mal calibradas.
   • Ordenação: Os submódulos de Atenção são selecionados e substituídos primeiro. Os submódulos de FFN são então pontuados e selecionados no modelo parcialmente comprimido para contabilizar o deslocamento do fluxo residual induzido pela substituição da Atenção.

2. Substituição de Resíduo Ajustada:

   • Cada submódulo selecionado é substituído por um substituto paramétrico leve $\hat{F}_\ell$ que aproxima a contribuição residual original ($y = F_\ell(x)$) em vez de seu cálculo interno.
   • Estrutura do Substituto: O substituto utiliza uma correção afim (portão elemento a elemento, viés e centralização) mais um termo de interação estruturado definido por um operador de baixo posto $U_\ell = V_\ell^\top W_\ell$.
   • Parametrização Assimétrica:
     • Atenção: Cada submódulo de Atenção removido recebe sua própria base de baixo posto ($V_\ell$) e parâmetros independentes, refletindo papéis heterogêneos entre as profundidades.
     • FFN: Uma base de entrada compartilhada $V$ é usada em todas as camadas de FFN selecionadas. Apenas os fatores de saída ($W_\ell$) e os termos afins permanecem específicos por camada. Esse compartilhamento de parâmetros entre camadas reduz significativamente o custo de implantação.
   • Ajuste: Os parâmetros são recuperados analiticamente via ajuste de forma fechada sobre estatísticas de calibração (mínimos quadrados regularizados por ridge e decomposição de autovalores), não exigindo otimização baseada em gradiente ou ajuste fino de todo o modelo.

Principais Contribuições

1. Identificação de Restrições: Os autores identificam a granularidade de camada completa e a seleção contígua como fatores limitantes na compressão baseada em substituição existente.
2. Framework SUBFIT: Um novo método que realiza seleção e substituição conjunta ao nível de submódulo (Atenção e FFN) com seleção não contígua.
3. Design Assimétrico: Uma arquitetura específica onde os bypasses de Atenção usam bases independentes, enquanto os bypasses de FFN compartilham uma base de entrada comum para minimizar o overhead de parâmetros.
4. Avaliação Abrangente: Testes extensivos em 10 LLMs (5 base, 5 ajustados por instrução) e 5 níveis de esparsidade (12,5% a 37,5%), demonstrando trade-offs melhorados em comparação com quatro baselines (Streamline, ReplaceMe).

Resultados Experimentais

• Perplexidade (Modelos Base): O SUBFIT alcança a menor degradação de perplexidade em 9 de 10 configurações de modelo-esparsidade. A lacuna em relação ao baseline mais forte aumenta significativamente sob compressão agressiva (ex: em 25% de esparsidade, o SUBFIT mostra uma degradação de perplexidade de 2,42× vs. 4,34× para o melhor baseline).
• Acurácia em Tarefas de Avaliação (Ajustados por Instrução): Com 25% de esparsidade, o SUBFIT retém, em média, 84,6% da acurácia densa de tarefas de avaliação, superando o baseline mais forte (81,6%). É o único método a reter acima de 80% de acurácia com 25% de esparsidade e acima de 73% com 37,5%.
• Estabilidade: O SUBFIT exibe uma variância significativamente menor na degradação de perplexidade comparado aos baselines, que mostram instabilidade (ex: desvio padrão crescendo ordens de magnitude) em alta esparsidade.
• Eficiência: O método gera ganhos mensuráveis de velocidade de inferência (1,18×–1,40× no tempo até o primeiro token) e economia de cache KV (proporcional à esparsidade) sem exigir kernels esparsos ou hardware customizado.
• Estudos de Ablação:
  • Compensação de Resíduo: Substituir FFNs é o fator dominante para a estabilidade; substituir tanto Atenção quanto FFN produz os melhores resultados.
  • Base Compartilhada: A base de FFN compartilhada alcança desempenho comparável a bases por camada, mas com significativamente menos parâmetros.
  • Sensibilidade de Posto (Rank): A qualidade do FFN melhora até o posto 4096, enquanto a da Atenção satura no posto 256, justificando as configurações de posto assimétricas.

Significância e Alegações
O artigo afirma que a granularidade é um eixo subexplorado na compressão pós-treinamento. Ao mover-se da substituição ao nível de camada para o nível de submódulo, o SUBFIT recupera mais o comportamento do modelo denso do que a substituição contígua de camada completa sob o mesmo orçamento de parâmetros. Os resultados sugerem que a recuperação não contígua e específica de submódulos fornece um trade-off qualidade-eficiência agregada mais forte, particularmente sob compressão agressiva. O método foi projetado para ser barato de implantar (ajuste de forma fechada) e compatível com stacks de inferência existentes. Os autores sugerem modestamente que esta abordagem pode se generalizar para outras unidades de remoção, como cabeças de atenção ou neurônios individuais, mas não reivindicam generalização imediata para essas unidades específicas neste trabalho.