Tuning Language Models by Mixture-of-Depths Ensemble

Resumo Técnico: Ajuste de Modelos de Linguagem por Ensamble de Mistura de Profundidade (MoDE)

Definição do Problema
Os Modelos de Linguagem de Grande Escala (LLMs) tradicionais baseados em Transformer dependem exclusivamente das representações das camadas finais tanto para a predição quanto para o cálculo da perda durante o ajuste fino (fine-tuning). Embora ferramentas de interpretabilidade como o "logit lens" tenham demonstrado que os estados ocultos de camadas tardias já contêm predições formadas e relevantes para a tarefa, essas representações intermediárias são tipicamente tratadas como subprodutos implícitos, em vez de sinais de treinamento explícitos. Consequentemente, o potencial preditivo das camadas tardias é subutilizado, e o ajuste fino padrão otimiza apenas a camada final, deixando as representações intermediárias otimizadas apenas indiretamente. Além disso, métodos existentes para aproveitar camadas intermediárias (ex: saída antecipada para eficiência ou direcionamento de representação) frequentemente não tratam os logits específicos de camada como um componente treinável direto da predição final.

Metodologia: Ensamble de Mistura de Profundidade (MoDE)
Os autores propõem o Mixture-of-Depths Ensemble (MoDE), um framework de ajuste que trata as saídas das últimas $k$ camadas de um LLM como um ensemble de "especialistas" dentro de um único modelo. Diferente do Mixture-of-Experts (MoE), que roteia diferentes tokens para diferentes modelos, o MoDE roteia a predição final para uma combinação ponderada de logits de diferentes profundidades dentro do mesmo modelo.

O framework consiste em três componentes principais:

1. Extração de Logits de Camada Tardia: Para um modelo com $n$ camadas, as últimas $k$ camadas ($K = \{L_{n-k+1}, \dots, L_n\}$) são selecionadas. A cabeça de vocabulário $\phi(\cdot)$ é aplicada aos estados ocultos de cada uma dessas $k$ camadas para gerar os logits $\ell_i(x)$.
2. Rede de Roteamento Aprendida: Um roteador leve $G(x)$, implementado como uma única camada linear seguida de softmax, calcula os pesos de roteamento baseados no estado oculto imediatamente antes das últimas $k$ camadas ($x = H_{n-k}$). Os logits finais são computados como uma soma ponderada:
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   O roteador pode ser tornado esparso (Top-K) para permitir a saída antecipada (early exit) durante a inferência.
3. Adaptação via Normalização e Destilação: Aplicar diretamente a cabeça pré-treinada às camadas intermediárias frequentemente resulta em uma perda inicial alta. Para resolver isso sem adicionar parâmetros pesados (ao contrário de matrizes afins por camada), os autores introduzem:
   • Normalização por Camada ($N_k$): Módulos de normalização treináveis (correspondentes à arquitetura pré-treinada, ex: LayerNorm) são adicionados a cada camada tardia para adaptar as representações ao espaço de entrada da cabeça.
   • Perda de Destilação ($L_{distill}$): A camada final atua como um "professor". Uma perda de divergência KL é aplicada para puxar as distribuições de saída das camadas tardias anteriores em direção à distribuição da camada final, estabilizando o treinamento e alinhando o ensemble com o objetivo da tarefa.

A perda total é $L_{MoDE} = L_{task} + \lambda L_{distill}$. O MoDE foi projetado para ser um mecanismo plug-in que pode ser aplicado sobre qualquer método de ajuste fino eficiente em parâmetros (PEFT) (ex: LoRA, DoRA) ou usado independentemente.

Principais Contribuições

• Sinais de Camada Tardia Treináveis: O artigo demonstra que os logits de camadas tardias, quando otimizados individualmente com destilação, convergem para perdas comparáveis à camada final e resolvem um conjunto complementar de problemas.
• Framework MoDE: Um mecanismo de ajuste simples e de baixo overhead que faz o ensemble de logits de camadas tardias via um roteador aprendido. Ele adiciona parâmetros treináveis desprezíveis (aproximadamente +0,04% do tamanho do modelo), enquanto permite que o modelo substitua módulos treináveis substancialmente maiores (ex: LoRA de camada total) com desempenho comparável.
• Eficiência e Flexibilidade: O framework suporta roteamento esparso (Top-K), permitindo saída antecipada e acelerações de geração de até 1,6$\times$ com perda mínima de precisão. Ele se compõe com métodos de PEFT existentes, aumentando seu desempenho.

Resultos Experimentais
Experimentos foram conduzidos no LLaMA2-7B e LLaMA3-8B em tarefas de raciocínio aritmético, raciocínio de senso comum e seguimento de instruções.

• Raciocínio Aritmético: O MoDE melhorou consistentemente o desempenho dos métodos base (LoRA, DoRA). A variante leve (LoRA aplicado a todas as camadas exceto as últimas $k$, mais MoDE) alcançou desempenho dentro de 0,3–0,5 pontos do LoRA de camada total, utilizando significativamente menos parâmetros.
• Raciocínio de Senso Comum: O MoDE proporcionou ganhos aditivos sobre métodos baseados em pesos (LoRA, DoRA). Embora métodos baseados em representação (ReFT) tenham performado fortemente por conta própria, o MoDE ofereceu melhorias confiáveis quando combinado com abordagens baseadas em pesos.
• Seguimento de Instruções: O MoDE gerou ganhos pequenos, mas consistentes, no MT-Bench, embora as melhorias tenham sido mais modestas do que nas tarefas de raciocínio, sugerindo que o seguimento de instruções depende fortemente de formatação de camada tardia que o ensemble captura apenas parcialmente.
• Ablação e Análise:
  • O roteador naturalmente concentra o peso nas últimas uma ou duas camadas, mas os componentes de destilação e normalização permitem que as camadas tardias anteriores contribuam efetivamente.
  • Um tamanho de ensemble de $k \approx 3\text{--}4$ é ideal; valores de $k$ maiores ajudam em datasets fáceis, mas prejudicam tarefas de raciocínio complexo.
  • O roteamento esparso Top-K mantém a maior parte da precisão enquanto permite acelerações significativas.

Significância e Alegações
Os autores enquadram o MoDE como um estudo de mecanismo, demonstrando que os logits de camadas tardias podem ser tornados diretamente úteis para o ajuste fino. O artigo afirma que:

1. As representações de camadas tardias possuem um potencial preditivo significativo que pode ser desbloqueado através de uma simples adaptação (normalização e destilação).
2. Esses sinais intermediários podem substituir módulos treináveis substancialmente maiores (como o LoRA de camada total) com desempenho comparável, oferecendo uma alternativa eficiente em parâmetros.
3. O MoDE fornece uma direção complementar para transformar representações de camadas intermediárias em sinais de treinamento acionáveis, distintos de abordagens de sondagem em tempo de inferência ou edição de pesos.

O artigo permanece modesto quanto ao seu escopo, observando limitações como a dependência de um $k$ fixo (em vez de seleção dinâmica), ganhos menores em seguimento de instruções em comparação ao raciocínio, e avaliação limitada a LLMs do tipo decoder-only. Trabalhos futuros são sugeridos para seleção dinâmica de $k$, famílias de modelos mais amplas e modelos maiores.