Dense2MoE: Pushing the Pareto Frontier of On-Device LLMs via Unified Pruning and Upcycling

Resumo Técnico: Dense2MoE

Declaração do Problema
A implantação de Modelos de Linguagem de Grande Porte (LLMs) em dispositivos de borda com recursos limitados (por exemplo, direção autônoma, IA de borda) é severamente impedida pelo "muro de memória da inferência". De acordo com a teoria Roofline de hardware, LLMs densos são limitados por memória durante a decodificação autoregressiva, o que significa que simplesmente reduzir os FLOPs teóricos é insuficiente para melhorar a latência. As soluções existentes enfrentam um compromisso fundamental:

1. Poda Estruturada: Métodos como UIDL reduzem a latência descartando camadas redundantes, mas inevitavelmente sacrificam a precisão do modelo ao remover Perceptrons de Camada Múltipla (MLPs), que servem como os principais repositórios de conhecimento dos LLMs.
2. Reutilização para MoE: Converter modelos densos em arquiteturas de Mistura de Especialistas (MoE) aumenta a capacidade sem aumentar linearmente o cálculo dinâmico. No entanto, a reutilização padrão frequentemente sofre de "colapso de representação" devido à inicialização homogênea dos especialistas (duplicando MLPs nativos), exigindo pré-treinamento contínuo massivo para induzir especialização. Além disso, métodos recentes de reaproveitamento esparsos frequentemente focam em redundância intra-camada, falhando em abordar as sobreposições estruturais inter-camadas de nível macro que causam redundância na atenção e sobrecarga no carregamento do KV-cache.

Metodologia: Dense2MoE
Os autores propõem Dense2MoE, um framework unificado que integra poda estruturada de camadas com reutilização para MoE por meio de uma técnica inovadora chamada Reutilização por Fusão de Camadas (LF-UC). O método opera em três etapas:

1. Análise de Similaridade Inter-camada: O framework identifica redundância estrutural analisando camadas adjacentes quanto a alta similaridade cosseno tanto nas saídas do decodificador quanto nas entradas dos MLPs. Crucialmente, impõe uma restrição estrita de correspondência de norma de características para garantir que a reutilização de MLPs redundantes não introduza anomalias de ativação ou deslocamentos de escala.
2. Reutilização por Fusão de Camadas (LF-UC): Em vez de descartar completamente camadas redundantes, o Dense2MoE poda fisicamente apenas os módulos de atenção (MHA) dessas camadas para eliminar operações pesadas em largura de banda de memória. Os MLPs pré-treinados dessas camadas podadas são preservados e reaproveitados como especialistas heterogêneos dentro das camadas base retidas.
   • Isso cria um pool de especialistas onde "especialistas base" são MLPs nativos duplicados e "especialistas suplementares" são os MLPs preservados das camadas podadas.
   • Um roteador dinâmico em nível de token despacha seletivamente tokens para esses especialistas ou os contorna por meio de caminhos residuais.
   • Este design aproveita a heterogeneidade estrutural inter-camada para mitigar o colapso de roteamento "O Vencedor Leva Tudo" comum na reutilização homogênea.
3. Pré-treinamento Contínuo: A arquitetura modificada passa por uma fase de pré-treinamento contínuo leve (225 bilhões de tokens) para otimizar o mecanismo de roteamento e alinhar os especialistas fundidos. Este orçamento consome menos de 1% do cálculo necessário para treinar um modelo MoE do zero.

Para determinar a profundidade de poda ótima, os autores empregam uma métrica de recompensa consciente de hardware inspirada no MnasNet, mapeando restrições de poda para a fronteira de Pareto de latência versus precisão para hardware alvo específico (por exemplo, NVIDIA Jetson Thor-U).

Principais Contribuições

• Framework Unificado: O Dense2MoE integra perfeitamente a poda estruturada de camadas e a reutilização para MoE, resolvendo o compromisso entre eficiência de inferência e precisão do modelo.
• Arquitetura Inovadora (LF-UC): Ao podar gargalos de atenção enquanto recicla MLPs como especialistas heterogêneos, o método reduz o consumo de largura de banda de memória sem destruir a capacidade representacional.
• Optimalidade de Pareto de Última Geração: O framework avança a fronteira de Pareto para latência de inferência em dispositivo versus precisão do modelo, superando bases densas, métodos de compressão de última geração e abordagens padrão de reutilização.
• Escalabilidade Robusta: O método foi validado em diversas escalas de modelo (0,5B a 7B) e famílias arquitetônicas distintas (Qwen2.5 e LLaMA2).

Resultados Experimentais
Avaliado em 9 benchmarks diversos (incluindo MMLU, ARC, GSM8K e HumanEval) usando modelos Qwen2.5 e LLaMA2:

• Desempenho vs. Latência: A configuração ótima do Dense2MoE (retendo 19 camadas de uma semente Qwen2.5-0,5B de 24 camadas) alcançou uma latência física de inferência de 271,0 ms no SoC Thor-U (reduzido de 307,5 ms para o modelo semente) enquanto alcançava uma pontuação média de benchmark de 47,31, superando significativamente a base densa (37,53).
• Comparação com Bases: Diferentemente de métodos tradicionais de poda (por exemplo, UIDL, LLM-Streamline) e abordagens de reaproveitamento esparsas (por exemplo, ToMoE, LLaMA-MoE), que sofreram colapso severo de capacidade em tarefas de raciocínio e codificação, o Dense2MoE manteve uma retenção representacional robusta. Superou a base densa em mais de 9,78 pontos percentuais na pontuação média, ativando apenas 0,42B de parâmetros.
• Escalabilidade: Quando aplicado a modelos maiores (Qwen2.5-1,5B e LLaMA2-7B), o Dense2MoE superou consistentemente seus equivalentes densos em 3,04 e 3,54 pontos percentuais, respectivamente, após ajuste fino no mesmo conjunto de dados de 225 bilhões de tokens.

Significado
O artigo afirma que o Dense2MoE estabelece uma receita prática e padrão industrial para reaproveitar LLMs densos de código aberto em arquiteturas altamente eficientes para dispositivos de borda. Ao explorar sistematicamente a redundância estrutural inter-camada, permite a otimização conjunta da latência de inferência e da precisão do modelo. O trabalho demonstra que é possível romper o muro de memória da inferência e preservar capacidades complexas de raciocínio sem os custos proibitivos de treinar modelos MoE do zero ou a degradação de precisão associada à poda agressiva. Isso fornece um contínuo flexível de modelos ótimos para implantações em borda críticas para latência.