Resumo Técnico: LoRA Orçamentado

Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) alcançam alto desempenho por meio de escalabilidade, mas isso acarreta custos significativos de implantação. Embora a Distilação de Conhecimento (KD) ofereça uma solução ao transferir conhecimento de grandes professores para estudantes menores, os pipelines existentes frequentemente falham em produzir modelos que sejam estruturalmente eficientes no momento da inferência.

Uma limitação primária reside no uso generalizado de métodos de ajuste fino eficientes em parâmetros, como o LoRA (Adaptação de Baixo Rango). Embora o LoRA reduza o número de parâmetros treináveis ao congelar a espinha dorsal densa e adicionar atualizações de baixo rango, ele deixa a espinha dorsal densa inalterada. Consequentemente, a distilação baseada em LoRA produz estudantes mais baratos de treinar, mas que não reduzem o custo dominante de inferência, pois as multiplicações de matrizes densas permanecem intactas. As abordagens atuais frequentemente otimizam para a contagem de parâmetros ou eficiência de treinamento sem controlar explicitamente como a computação é alocada dentro da arquitetura do estudante durante a inferência.

Metodologia: LoRA Orçamentado

Os autores propõem o LoRA Orçamentado, um framework de distilação que reformula a compressão de modelos como um problema estruturado de alocação de computação. Em vez de fixar a arquitetura do estudante, o método introduz um orçamento global de computação que dita a fração final de computação densa a ser retida.

A abordagem opera em três etapas:

1. Módulos LoRA com Portão Sensíveis ao Orçamento:
   O modelo estudante substitui projeções lineares selecionadas (por exemplo, $q/k/v/o$ de atenção e $up/gate/down$ de MLP) por módulos com portão. Esses módulos calculam:
   $$y = d \cdot Wx + \frac{\alpha}{r_{max}} B((Ax) \odot g)$$
   Onde:

   • $W$ é o peso denso congelado.
   • $d \in [0, 1]$ é um coeficiente de retenção do caminho denso controlado por um agendador de orçamento.
   • $A, B$ são fatores LoRA, e $g$ representa portões aprendidos por rango.
   • Esta formulação permite que o modelo varie dinamicamente a capacidade ativa de ambos os caminhos densos e de baixo rango através de diferentes módulos.

2. Agendamento e Controlador de Orçamento:
   Um parâmetro de orçamento global $F \in [0, 1]$ especifica a fração alvo de computação densa a ser retida ao final do treinamento. Um controlador impõe este orçamento ao:

   • Definir uma trajetória alvo de custo denso $b(t)$ usando um agendamento cosseno.
   • Calcular o custo denso retido total $C(t)$ com base na contagem de MACs das projeções densas ativas.
   • Reduzir greedymente os coeficientes de retenção ao nível do módulo ($d$) para atender ao orçamento alvo, priorizando a remoção de módulos de menor custo primeiro.
   • Isso cria uma transferência controlada de computação dos caminhos densos para os caminhos de baixo rango durante o treinamento.

3. Compressão Pós-Treinamento:
   Após o treinamento, o modelo passa por uma etapa de consolidação para realizar economias reais de inferência. Com base no coeficiente de retenção final $d$, cada módulo é mapeado para uma das três formas de implantação:

   • Caso 1 (Remover Densidade): Se $d < \epsilon_{zero}$, o caminho denso é removido inteiramente, restando apenas o componente de baixo rango podado.
   • Caso 2 (Aproximação SVD): Se $\epsilon_{zero} \le d < \epsilon_{lr}$, a contribuição densa residual é aproximada via uma SVD de baixo rango, fundindo-a com os fatores LoRA.
   • Caso 3 (Mesclar Densidade): Se $d \ge \epsilon_{lr}$, o caminho denso é retido e mesclado com a atualização LoRA em uma camada linear densa padrão.

Este pipeline produz uma família de modelos estudantes parametrizados pelo único dial de orçamento $F$, permitindo uma troca direta entre eficiência e desempenho.

Contribuições Principais

• Nova Formulação: O artigo propõe tratar a distilação como um problema estruturado de alocação de computação sob uma restrição de orçamento global, em vez de uma otimização de arquitetura fixa.
• Framework LoRA Orçamentado: Introduz um método que combina retenção densa imposta por controlador com alocação de baixo rango aprendida, seguida por compressão pós-treinamento para criar estudantes eficientes para implantação.
• Troca Controlável: Demonstra que um único parâmetro de orçamento gera uma troca qualidade-eficiência controlável, formando uma fronteira de Pareto através dos custos de treinamento e inferência.
• Retenção de ICL: Mostra que o LoRA Orçamentado retém desempenho substancialmente superior à distilação baseada em LoRA padrão em sondas de Aprendizado em Contexto (ICL) no estilo de funções.

Resultados

Os autores avaliaram o LoRA Orçamentado usando um estudante de 6 camadas inicializado a partir de um modelo de 0,13B, distilado de professores de escalas variadas (0,13B, 0,5B, 1B).

• Troca Qualidade-Eficiência:

  • Com um orçamento moderado ($F=0,4$), o LoRA Orçamentado iguala a perplexidade do KD-LoRA padrão, enquanto alcança um aceleramento de 1,74× na inferência de módulo comprimido e reduz a computação de treinamento em 0,59×.
  • Com um orçamento agressivo ($F=0,0$), o método alcança um aceleramento de 4,05× na inferência de módulo comprimido e 0,38× de computação de treinamento, com apenas degradação moderada de perplexidade.
  • O método traça uma fronteira de Pareto clara, oferecendo pontos operacionais de menor custo com degradação de perplexidade controlada em comparação ao LoRA padrão.

• Retenção de Aprendizado em Contexto (ICL):

  • Em sondas de ICL no estilo de funções (adaptadas do conjunto Function Vectors), o LoRA Orçamentado supera significativamente o KD-LoRA padrão, especialmente ao distilar de professores maiores.
  • Por exemplo, com um professor de 1B, o KD-LoRA alcançou uma pontuação média macro de ICL de 2,7, enquanto o LoRA Orçamentado ($F=0,0$) atingiu 5,0 e ($F=0,4$) atingiu 5,6.
  • Isso sugere que preservar comportamentos de alto nível depende não apenas de igualar a perplexidade, mas de manter a estrutura computacional subjacente através da retenção densa orçamentada.

Significado e Alegações

O artigo alega que, sob distilação com restrição de computação, reter comportamento é menos sobre igualar a perplexidade ou simplesmente remover parâmetros; trata-se de controlar como a computação densa é transferida para caminhos de baixo rango.

Os autores argumentam que o LoRA padrão falha em entregar economias de inferência porque deixa a espinha dorsal densa intacta. Em contraste, o LoRA Orçamentado trata os caminhos densos e de baixo rango como recursos concorrentes. Os resultados sugerem que:

1. A Alocação Estruturada Importa: Uma redução gradual e orçamentada da capacidade densa permite uma transferência mais suave de capacidade para os caminhos de baixo rango em comparação com a remoção abrupta.
2. Além da Perplexidade: Comportamentos de alto nível (como ICL) são melhor preservados quando o processo de treinamento gerencia explicitamente a estrutura da computação, mesmo que a perplexidade final seja semelhante à do LoRA padrão.
3. Eficiência de Implantação: O método traduz com sucesso sinais de orçamentação no momento do treinamento em economias reais no momento da inferência através da compressão pós-treinamento, oferecendo um espectro de modelos desde densos até totalmente de baixo rango.

O trabalho se posiciona como um passo em direção a visualizar a eficiência não meramente como um problema de contagem de parâmetros, mas como um problema de alocação de computação durante o processo de treinamento e distilação.