Scalable LLM Reasoning Acceleration with Low-rank Distillation

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) modernos, especialmente aqueles focados em raciocínio (como modelos de "pensamento" ou thinking models), dependem de gerações longas e complexas, frequentemente utilizando Cadeias de Pensamento (Chain-of-Thought - CoT). Isso demanda recursos computacionais significativos e aumenta a latência.

Para mitigar isso, foram desenvolvidos métodos de inferência eficiente baseados em esparsidade (como GRIFFIN e CATS) que comprimem os blocos Feedforward (FF) das redes neurais. Embora esses métodos preservem bem o desempenho em tarefas de linguagem padrão (como resumo ou compreensão de leitura), eles degradam severamente o desempenho em tarefas de raciocínio matemático.

• Causa: Os erros de aproximação introduzidos pela compressão acumulam-se ao longo das gerações longas típicas do raciocínio, desviando a trajetória de geração e levando a respostas incorretas.
• Desafio: Existe uma necessidade urgente de algoritmos de inferência eficiente que mantenham simultaneamente a capacidade de raciocínio e a eficiência computacional, sem penalizar tarefas de linguagem.

2. Metodologia: Caprese

O artigo propõe o Caprese (CAPability REcovery with Scalable Efficiency), um método de distilação de recursos que recupera as capacidades de raciocínio perdidas ao usar inferência eficiente, focando especificamente nos blocos Feedforward (FF).

Princípio Central

A motivação surge da observação de que, embora a compressão FF (como top-k ou thresholding) reduza a precisão, o resíduo do erro (a diferença entre a saída do modelo original e a saída comprimida) possui uma estrutura de baixo rango. Isso significa que o erro pode ser aproximado eficientemente com poucas dimensões.

Arquitetura e Processo

1. Compressão Base: O método utiliza algoritmos de compressão existentes (como GRIFFIN ou CATS) para reduzir o número de parâmetros ativos e a latência.
2. Camadas de Baixo Rango (Low-Rank): Em vez de treinar o modelo inteiro, o Caprese adiciona pequenas camadas lineares de baixo rango ($L \in \mathbb{R}^{D \times r}$ e $R \in \mathbb{R}^{r \times D}$) aos blocos FF comprimidos.
   • O objetivo é aprender a aproximar o resíduo: $\min_{L,R} \|FF(x) - \hat{FF}(x) - xLR\|^2$.
   • O tamanho interno $r$ é muito pequeno (ex: 256) comparado à dimensão intermediária do FF (ex: 14.336), adicionando apenas ~1% de parâmetros totais.
3. Distilação:
   • Distilação por Camada (Layer-wise): Treina as camadas de baixo rango independentemente para cada bloco FF, usando o modelo original como referência.
   • Distilação End-to-End (E2E): Após a inicialização, todo o conjunto de camadas é ajustado conjuntamente para minimizar o erro na saída final do embedding (antes da cabeça linear), congelando os pesos originais do modelo.
4. Eficiência de Inferência: As operações das camadas de baixo rango ($xLR$) podem ser executadas em paralelo com as operações do bloco FF comprimido. Além disso, as matrizes $L$ e $R$ podem ser concatenadas às matrizes de projeção existentes, evitando sobrecarga de latência significativa.

3. Contribuições Principais

1. Recuperação de Desempenho: O Caprese recupera grande parte (se não toda) a capacidade de raciocínio perdida pela compressão eficiente, permitindo que modelos comprimidos superem ou igualem o desempenho do modelo original em tarefas matemáticas.
2. Eficiência e Baixo Custo:
   • Adiciona apenas ~0.8% a 1% de parâmetros adicionais.
   • Requer apenas 20.000 amostras sintéticas de treinamento matemático para a distilação.
   • Reduz o número de parâmetros ativos em ~2B (para modelos de 8B/9B parâmetros) ao combinar compressão + distilação.
3. Redução de Latência e Tokens:
   • Reduz o tempo até o próximo token (time-to-next-token) em mais de 16%.
   • Incentiva implicitamente respostas mais curtas (redução de até 8.5% no número de tokens), sem necessidade de regularização explícita de comprimento.
4. Versatilidade: Funciona bem em modelos de "pensamento" (DeepSeek-R1) e modelos de instrução (Llama, Gemma), preservando a qualidade em tarefas de linguagem enquanto melhora o raciocínio.

4. Resultados Experimentais

Os experimentos foram realizados em diversos modelos (Llama 3.1/3.2, Gemma 2, DeepSeek-R1-Distill-Qwen) em tarefas como GSM8K, MATH-500, AIME 2024 e AMC 2023.

• Recuperação de Precisão:
  • No modelo DeepSeek-R1-Distill-Qwen 7B, o método CATS (compressão pura) reduziu a precisão no AMC 2023 de 75.00% para 62.50%. O Caprese (E2E) recuperou esse desempenho, elevando-o para 78.13% (superando o modelo original).
  • Em modelos menores (1.5B), o Caprese recuperou a maior parte da queda drástica causada pelo GRIFFIN (ex: de 42% para 74.8% no MATH-500).
• Escalabilidade (Test-Time Scaling):
  • Ao escalar a geração (aumentando o número de tentativas ou o comprimento do CoT), o Caprese mantém a curva de desempenho superior à do modelo comprimido puro, aproximando-se do modelo completo.
  • No Llama 3.2 3B, o Caprese melhorou o Pass@100 em 7.0% em relação ao modelo original, usando 15.8% menos computação ativa.
• Latência:
  • Para o Qwen 2.5 14B, o Caprese reduziu a latência total em >16% em comparação ao modelo completo, com uma sobrecarga mínima em relação ao método de compressão puro (GRIFFIN).
  • Em comparação com LoRA (Low-Rank Adaptation) aplicado à compressão, o Caprese é superior em latência, pois LoRA exige operações sequenciais que anulam as economias de tempo da compressão.
• Brevidade: O Caprese tende a gerar respostas mais curtas do que o modelo completo ou o método CATS, economizando tokens e memória sem sacrificar a precisão.

5. Significância e Conclusão

O trabalho Caprese representa um avanço crucial para a viabilidade de inferência de LLMs de raciocínio em ambientes com restrições de recursos. Ele resolve o dilema entre eficiência (compressão) e precisão (raciocínio complexo).

• Impacto Prático: Permite a implantação de modelos de raciocínio avançados em hardware mais acessível, reduzindo custos de inferência e tempo de resposta.
• Inovação Técnica: Demonstra que a distilação de resíduos de erro em camadas de baixo rango é uma estratégia superior à simples compressão ou ao ajuste fino tradicional (LoRA) para cenários de inferência esparsa.
• Futuro: O método abre caminho para inferência de LLMs sem degradação de desempenho, permitindo que técnicas de test-time scaling (como geração de múltiplas respostas ou CoTs longos) sejam aplicadas de forma eficiente em larga escala.

Em resumo, o Caprese oferece uma solução "ganha-ganha": mais rápido, mais barato e tão inteligente quanto o modelo original para tarefas de raciocínio.