Two-Stage Regularization-Based Structured Pruning for LLMs

Título: TRSP: Poda Estruturada Baseada em Regularização de Duas Etapas para Grandes Modelos de Linguagem (LLMs)

1. O Problema

A implantação de Grandes Modelos de Linguagem (LLMs) é frequentemente limitada pelo seu enorme número de parâmetros, o que exige alto custo computacional e de memória. A poda estruturada (remoção de camadas inteiras ou canais) surge como uma solução promissora para reduzir o tamanho do modelo e acelerar a inferência.

No entanto, os métodos de poda de camadas existentes apresentam limitações críticas:

• Perda de Conhecimento: Eles geralmente avaliam a importância das camadas, removem as "menos importantes" e, em seguida, tentam recuperar o desempenho através de um extenso fine-tuning (re-treinamento).
• Ineficiência: Mesmo camadas consideradas "menos importantes" podem conter conhecimento valioso. A remoção direta causa perda irreversível de informação, exigindo re-treinamento massivo para compensar a degradação de desempenho.
• Custo Computacional: O processo de re-treinamento após a poda consome recursos significativos, anulando parte dos benefícios da compressão.

2. Metodologia: TRSP

O TRSP propõe uma abordagem inovadora que realiza a poda sem re-treinamento (retraining-free), utilizando um processo de regularização de duas etapas antes da remoção física das camadas. O objetivo é redistribuir o conhecimento das camadas a serem removidas para as camadas preservadas antes da poda.

O processo segue quatro etapas principais:

1. Preparação de Dados: Seleciona-se uma pequena amostra de dados (ex: 128 instâncias) de benchmarks padrão (como WikiText-2) para guiar o processo de regularização.
2. Aprendizado de Pesos das Camadas (1ª Etapa de Regularização):
   • Um peso aprendível ($S[i]$) é atribuído a cada camada do Transformer.
   • A saída de cada camada é multiplicada por seu peso correspondente.
   • Um termo de regularização $\ell_1$ (norma L1) é adicionado à função de perda para incentivar que os pesos das camadas menos importantes se aproximem de zero.
   • Isso é feito iterativamente (estratégia gulosa): identifica-se a camada com o menor peso, mascara-se sua contribuição e repete-se o processo até identificar todas as camadas a serem podadas.
3. Regularização de Transferência de Conhecimento (2ª Etapa de Regularização):
   • Uma vez identificadas as camadas a serem removidas (conjunto $P$), aplica-se uma regularização adicional sobre a diferença entre a saída e a entrada dessas camadas específicas.
   • A função de perda inclui a norma ($\ell_1$ ou $\ell_2$) de $(X_{out}^i - X_{in}^i)$.
   • Objetivo: Forçar as camadas a serem removidas a se comportarem como funções identidade (onde a saída é quase igual à entrada). Isso obriga o modelo a transferir o conhecimento contido nessas camadas para as camadas restantes, minimizando a perda de informação quando elas forem removidas.
4. Poda: As camadas com os menores pesos são removidas diretamente do modelo. Como o conhecimento já foi transferido via regularização, o modelo mantido não requer re-treinamento.

3. Principais Contribuições

• Retenção de Conhecimento: Ao aplicar a regularização de duas etapas, o TRSP mitiga a perda de conhecimento ao redistribuir a informação das camadas a serem podadas para as camadas preservadas, evitando a degradação abrupta de desempenho.
• Eficiência e Custo Mínimo: O método é livre de re-treinamento (retraining-free). A quantidade de dados necessária para a regularização é mínima (apenas 128 amostras), reduzindo drasticamente o custo computacional e a dependência de dados de treinamento.
• Aceleração End-to-End: Como uma técnica de poda de camadas, oferece aceleração significativa na inferência, superando métodos de poda não estruturada que exigem hardware especializado para ganhos de velocidade.

4. Resultados Experimentais

O TRSP foi avaliado em diversos modelos (Phi-2, OPT, LLaMA2, LLaMA3) e tarefas (geração e zero-shot).

• Desempenho Superior: Em tarefas de geração (medida por Perplexidade no WikiText-2) e tarefas zero-shot (PIQA, WinoGrande, HellaSwag, ARC), o TRSP superou consistentemente métodos de poda de camadas de ponta (como SLEB, ShortGPT, LaCo).
  • Exemplo: No LLaMA2-7B com 25% de poda, o TRSP reduziu a perplexidade em 20% comparado ao ShortGPT e manteve uma queda de apenas 1% na precisão média.
• Aceleração:
  • No OPT-13B com 50% de poda, houve um aumento de 75% na taxa de transferência (throughput) e uma redução de 46% na latência.
  • No LLaMA2-13B, observou-se um aumento de 71% no throughput e redução de 45% na latência.
• Robustez: O método manteve desempenho estável em diferentes taxas de poda (até 60%) e foi robusto a diferentes conjuntos de dados de calibração (Alpaca, C4, PTB).
• Escalabilidade: Resultados adicionais em Llama-2-70B confirmaram que o método escala bem para modelos maiores, superando todas as baselines.

5. Significância

O TRSP representa um avanço significativo na eficiência de LLMs ao resolver o dilema entre compressão e perda de desempenho.

• Mudança de Paradigma: Em vez de "cortar e consertar" (poda seguida de re-treinamento), o TRSP adota uma abordagem de "preparar e cortar", garantindo que o modelo esteja pronto para a poda antes da remoção física.
• Viabilidade de Implantação: Ao eliminar a necessidade de re-treinamento custoso e manter a alta precisão com aceleração significativa, o TRSP torna a implantação de LLMs compactos em ambientes com recursos limitados (como dispositivos de borda ou servidores com restrições de memória) muito mais viável.
• Código Aberto: O código está disponível publicamente, facilitando a adoção e reprodução pelos pesquisadores.

Em resumo, o TRSP oferece uma solução elegante e eficiente para a poda estruturada de LLMs, equilibrando a redução de tamanho do modelo com a preservação máxima da capacidade de raciocínio e geração de linguagem.