🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

O artigo apresenta o MoE-nD, uma abordagem de Mixture-of-Experts que otimiza a compressão do cache KV em LLMs ao rotear dinamicamente cada camada para uma configuração específica de evicção e quantização, alcançando uma compressão de 14x com perda mínima de precisão em tarefas de contexto longo.

Autores originais: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

Publicado 2026-04-21

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem um cérebro gigante (uma Inteligência Artificial) que precisa ler um livro inteiro de 100 páginas para responder a uma pergunta. O problema é que esse cérebro tem uma "mesa de trabalho" (a memória) muito pequena. Se ele tentar guardar todas as 100 páginas na mesa de uma só vez, a mesa vira, o livro cai e o cérebro trava.

Para resolver isso, os cientistas tentam duas coisas:

Jogar fora algumas páginas (Evicção): Manter apenas as partes mais importantes.
Escrever com uma caneta mais fina (Quantização): Guardar as informações com menos detalhes para economizar espaço.

O problema é que, até agora, todos tratavam todas as páginas do livro da mesma maneira. Se a mesa era pequena, eles jogavam fora 50% de todas as páginas, ou escreviam todas com a mesma caneta fina.

O Problema: Nem Toda Página é Igual

Os autores deste artigo descobriram algo fascinante: as páginas do livro não são iguais.

Algumas páginas são como o índice ou o sumário: se você jogá-las fora, o cérebro esquece tudo. Elas precisam ser guardadas com muito cuidado e em alta qualidade.
Outras páginas são como anotações de rascunho ou repetições: o cérebro consegue entender o resto mesmo se você jogar metade delas fora ou escrever com uma caneta bem grossa.

Tratar tudo igual é como tentar economizar dinheiro cortando o salário do CEO e do estagiário na mesma proporção. O CEO (as camadas importantes) vai parar de trabalhar, e o estagiário (as camadas menos importantes) nem vai sentir falta.

A Solução: MoE-nD (O "Gerente Inteligente")

O MoE-nD é como um gerente de escritório superinteligente que olha para cada página do livro individualmente antes de decidir o que fazer.

Ele usa um sistema de "especialistas" (daí o nome Mixture-of-Experts):

O Especialista em Jogar Fora: Decide quais páginas podem ser descartadas.
O Especialista em Economizar Espaço: Decide quais páginas podem ser escritas de forma mais simples.

Como funciona na prática?
Imagine que você tem um orçamento fixo para a mesa de trabalho.

O gerente olha para a Página 1 (muito importante) e diz: "Não podemos jogar nada fora, e vamos escrever com uma caneta de alta qualidade."
Ele olha para a Página 50 (menos importante) e diz: "Podemos jogar 75% desta página fora e escrever o resto com uma caneta bem barata."
Ele olha para a Página 100 e diz: "Vamos jogar metade fora, mas manter a qualidade."

Ele faz isso para cada uma das 28 camadas (páginas) do cérebro, criando um plano personalizado para cada uma, mas sempre respeitando o limite total de espaço na mesa.

O Resultado: O Truque Mágico

Os autores testaram isso em tarefas difíceis de raciocínio (como resolver problemas de matemática complexa ou responder perguntas sobre textos longos).

O Método Antigo: Tentava economizar espaço jogando fora as mesmas coisas em todas as páginas. O resultado? O cérebro perdia a noção e respondia errado (precisão baixa).
O MoE-nD: Com o mesmo espaço de memória (que era 14 vezes menor que o original!), o cérebro conseguiu responder tão bem quanto se tivesse a mesa gigante cheia.

É como se você conseguisse levar uma biblioteca inteira no bolso de uma calça jeans, mas em vez de amassar todos os livros, você apenas tirou as páginas vazias de alguns e reduziu a fonte de outros, mantendo o conteúdo essencial intacto.

Quando isso NÃO funciona?

O artigo é honesto e diz que essa técnica só brilha quando o livro é muito longo e o espaço é muito apertado.
Se o livro for curto (como uma mensagem de WhatsApp) e a mesa for grande o suficiente, o gerente inteligente percebe que não precisa fazer nada diferente. Ele diz: "Tudo bem, vamos guardar tudo do jeito normal". Nesse caso, a técnica não traz vantagem, mas também não atrapalha.

Resumo em uma frase

O MoE-nD é um sistema que ensina a IA a ser esperta na hora de economizar: em vez de cortar tudo igual, ela sabe exatamente onde pode cortar mais e onde precisa manter a qualidade, permitindo que modelos gigantes rodem em computadores comuns sem perder inteligência.

1. O Problema: Gargalo de Memória em Contextos Longos

A inferência de Grandes Modelos de Linguagem (LLMs) com contextos longos enfrenta um gargalo crítico: o KV Cache (Cache de Chave-Valor). Para sequências extensas, o consumo de memória pode atingir dezenas de GBs, dominando a latência e causando falhas de "out-of-memory" (OOM) em hardware comercial.

As técnicas de compressão existentes geralmente atuam em apenas um dos quatro eixos do tensor KV (4D) e aplicam a mesma estratégia uniformemente a todas as camadas do modelo:

Evicção (Eixo de Sequência): Remover tokens menos importantes.
Quantização (Eixo de Precisão): Reduzir o número de bits por elemento (ex: INT4, INT8).
Projeção de Baixo Rank: Reduzir a dimensão da cabeça.
Compartilhamento entre Camadas: Compartilhar KV entre camadas adjacentes.

A Limitação Atual: A maioria dos métodos aplica uma política uniforme (ex: mesma taxa de evicção ou mesma quantização para todas as camadas). O artigo demonstra que essa homogeneidade desperdiça potencial de precisão, pois diferentes camadas respondem de maneira drasticamente diferente a operações de evicção versus quantização. O mix ótimo de compressão varia significativamente de camada para camada.

2. Metodologia: MoE-nD

Os autores propõem o MoE-nD, um framework baseado em Mixture-of-Experts (MoE) que trata cada eixo de compressão como um "especialista" e utiliza um roteador para selecionar a combinação ideal para cada camada individualmente, sob um orçamento global de memória.

Componentes Principais:

Espaço de Design por Camada:
- Cada camada $\ell$ $ℓ$ possui três "knobs" (controles) independentes:
  - Taxa de Retenção (Keep Ratio): Fração de tokens mantidos após evicção (ex: 0.1, 0.25, 0.5, 0.75, 0.9, 1.0).
  - Precisão K (K-bits): Largura de bits para as chaves (ex: 16, 8, 4).
  - Precisão V (V-bits): Largura de bits para os valores (ex: 16, 8, 4).
- O objetivo é encontrar a tupla $(keep\_ratio, k\_bits, v\_bits)$ para cada camada que minimize a perda de qualidade dentro do orçamento de memória $M$ .
Calibração Offline (Tabela de Sensibilidade):
- Antes da inferência, o sistema mede a sensibilidade de cada camada a diferentes operações de compressão.
- Utiliza-se um proxy leve: o erro L2 relativo entre a saída de atenção de precisão total e a saída comprimida (aplicando a compressão apenas naquela camada específica).
- Isso gera uma tabela de sensibilidade $S$ que mapeia quanto cada operação degrada a precisão em cada camada.
Solver Guloso (Greedy Solver):
- Encontrar a configuração global ótima é um problema discreto complexo (NP-difícil).
- O MoE-nD utiliza um solver guloso que começa com a configuração mais barata em todas as camadas e iterativamente "melhora" a camada que oferece a maior redução de sensibilidade por unidade de memória adicionada ( $\Delta S / \Delta m$ ), até esgotar o orçamento.
- O tempo de solução é rápido (< 50ms em CPU).
Patch de Atenção Heterogênea:
- Para lidar com o fato de que cada camada pode ter um tamanho de cache diferente ( $T_\ell$ ) e diferentes larguras de bits, os autores implementam um patch de atenção personalizado.
- Este patch rastreia as posições originais dos tokens retidos em cada camada para reverter corretamente a rotação RoPE (Rotary Positional Embeddings) durante a geração, garantindo que a heterogeneidade não quebre a coerência temporal.

3. Contribuições Chave

Roteamento Multi-Eixo por Camada: É a primeira proposta a realizar roteamento conjunto de taxa de evicção e quantização (K e V) em granularidade de camada sob um orçamento global.
Descoberta de Heterogeneidade Crítica: O trabalho demonstra empiricamente que a sensibilidade à evicção agressiva varia em até 689x entre as camadas, enquanto a sensibilidade à quantização V é quase uniforme. Isso valida a necessidade de um roteador que veja ambos os eixos simultaneamente.
Isolamento de Contribuição: Através de ablações, o paper prova que o ganho de desempenho vem quase inteiramente do roteamento de evicção por camada, e não do roteamento de quantização (que, sozinho, traz ganhos marginais ou nulos).
Definição de Escopo: O trabalho identifica claramente quando o método funciona (contextos longos ou orçamentos apertados) e quando não funciona (contextos curtos onde o solver opta por não evictar nada).

4. Resultados Experimentais

Os experimentos foram conduzidos no modelo DeepSeek-R1-Distill-Qwen-7B em benchmarks como LongBench-v1 (subconjunto de 4 tarefas com 16k tokens) e AIME (raciocínio matemático).

Desempenho no LongBench:
- A variante heterogênea do MoE-nD (2d_hetero) alcançou 14x de compressão (reduzindo o cache de 1.9 GB para 136 MB) mantendo a mesma precisão do baseline sem compressão (11.5% vs 12.0% em média).
- Em comparação, outros métodos comprimidos (evicção uniforme, quantização uniforme, ou apenas quantização por camada) no mesmo ou menor orçamento de memória tiveram desempenho muito inferior (abaixo de 8/100 em média).
- Em tarefas individuais, métodos uniformes perderam até 99% de precisão, enquanto o MoE-nD manteve a estabilidade.
Desempenho no AIME (Raciocínio):
- O MoE-nD superou consistentemente o baseline de dois eixos não-heterogêneo em todas as configurações de orçamento.
- Em orçamentos muito apertados (ex: 64 tokens), o MoE-nD manteve desempenho (ex: 30% de acurácia), enquanto outros métodos caíram para 0%.
Resultados Nulos (MATH-500 e TREC):
- Em tarefas de contexto curto (MATH-500) ou com prompts curtos (TREC), o MoE-nD não trouxe vantagem sobre métodos uniformes.
- Causa: O solver guloso, diante de um orçamento frouxo e entradas curtas, escolhe manter 100% dos tokens (keep=1.0) na maioria das camadas. Sem evicção real, o roteamento heterogêneo degenera para uniforme, eliminando a vantagem. Isso valida que o método é específico para cenários onde a restrição de memória força escolhas difíceis.

5. Significado e Conclusão

O MoE-nD representa um avanço significativo na eficiência de LLMs para contextos longos. Ao abandonar a suposição de uniformidade entre camadas e adotar uma abordagem de "Mixture-of-Experts" para compressão, o método consegue:

Reduzir drasticamente o uso de memória (até 14x) sem perda de precisão em tarefas de raciocínio e contexto longo.
Demonstrar que a evicção inteligente por camada é o fator crítico, superando a simples quantização adaptativa.
Fornecer um template para futuras técnicas de compressão que combinam múltiplos eixos ortogonais sob um único roteador calibrado.

O trabalho destaca que a heterogeneidade das camadas dos transformadores não é apenas um detalhe, mas uma oportunidade fundamental para otimização, desde que o sistema de roteamento seja capaz de explorar as diferenças de sensibilidade entre evicção e quantização de forma conjunta.