MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

Resumo Técnico: MaskPro – Aprendizado Probabilístico em Espaço Linear para Esparsidade Estrita (N:M) em LLMs

Declaração do Problema
O escalonamento rápido dos Grandes Modelos de Linguagem (LLMs) criou um gargalo significativo na eficiência de inferência, tornando necessárias técnicas de compressão eficazes. A esparsidade semi-estruturada, especificamente o padrão (N:M) onde $N$ elementos são retidos a cada $M$ pesos consecutivos, oferece uma solução amigável ao hardware para aceleração e redução de memória. No entanto, identificar a máscara (N:M) ótima é um problema formidável de otimização combinatória. As abordagens existentes dividem-se em duas categorias com limitações distintas:

1. Métodos baseados em regras: Estes empregam buscas ganancias, camada por camada, utilizando métricas criadas manualmente (por exemplo, minimização do erro de reconstrução por camada). Embora computacionalmente baratos, sofrem de vieses inerentes e apresentam lacunas consideráveis em relação ao desempenho final em tarefas.
2. Métodos baseados em aprendizado (por exemplo, MaskLLM): Estes otimizam máscaras diretamente via descida de gradiente em uma perda de geração. Embora eficazes, incorrem em custos de treinamento proibitivos. Especificamente, o MaskLLM requer o armazenamento de logits para todas as máscaras possíveis dentro de um grupo de $M$ pesos, levando a uma complexidade de memória de $O(\binom{M}{N} \frac{d}{M})$. Para $M$ grande, isso escala exponencialmente, tornando-o inviável para modelos grandes. Além disso, o vasto espaço combinatório induz alta variância nos gradientes de política, causando instabilidade no treinamento.

Metodologia: MaskPro
Os autores propõem o MaskPro, um novo framework probabilístico em espaço linear projetado para aprender esparsidade (N:M) de forma eficiente e estável.

1. Relaxação Probabilística em Espaço Linear:
   Em vez de definir uma distribuição categórica sobre todas as $\binom{M}{N}$ máscaras possíveis para um grupo de $M$ pesos, o MaskPro reformula o problema. Ele estabelece uma distribuição categórica sobre os $M$ elementos individuais dentro de cada grupo. A máscara (N:M) é então gerada realizando uma amostragem N-ária sem reposição a partir dessa distribuição.

   • Base Teórica: Os autores provam que qualquer máscara (N:M) pode ser representada como uma soma probabilística de $N$ vetores de base distintos.
   • Eficiência de Memória: Essa reformulação reduz o número de parâmetros aprendíveis (logits) de $O(\binom{M}{N} \frac{d}{M})$ para um linear $O(d)$, onde $d$ é o número total de pesos. Isso elimina a explosão de memória associada aos métodos anteriores baseados em aprendizado.

2. Estimador de Gradiente de Política Refinado (PGE):
   Para otimizar os logits, o MaskPro utiliza gradientes de política. Os autores identificam que gradientes de política padrão usando valores de perda brutos são instáveis em LLMs devido à alta variância causada por diferenças em mini-lotes (onde a variação de perda proveniente da amostragem de dados frequentemente supera a variação de perda proveniente da amostragem de máscaras).

   • Resíduos de Perda: Em vez de usar a perda bruta $f(m \odot w, \xi)$, o MaskPro atualiza os logits com base no resíduo de perda relativo a uma máscara de linha de base inicial $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. Isso isola o impacto da mudança na máscara.
   • Rastreador de Média Móvel: Para estabilizar ainda mais o treinamento e prevenir explosões numéricas, um rastreador de média móvel $\delta$ é introduzido para suavizar a linha de base. A regra de atualização torna-se:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • Garantias Teóricas: O artigo fornece provas mostrando que esses estimadores refinados permanecem como estimadores não viesados do verdadeiro gradiente de política. Além disso, sob a condição de que a perda da máscara amostrada seja menor que o dobro da perda da linha de base (uma condição facilmente satisfeita na prática), a variância do estimador refinado é estritamente menor que a do gradiente de política padrão.

Contribuições Principais

1. Framework em Espaço Linear: O MaskPro reduz o requisito de memória para aprender esparsidade (N:M) de exponencial $O(\binom{M}{N} \frac{d}{M})$ para linear $O(d)$, modelando o processo de seleção como amostragem N-ária sem reposição a partir de uma distribuição categórica sobre $M$ elementos.
2. Gradiente de Política Aprimorado: Os autores introduzem um mecanismo de atualização com variância reduzida usando resíduos de perda e um rastreador de média móvel, abordando a instabilidade de treinamento causada pela alta variância dos gradientes de política em grandes espaços combinatórios.
3. Análise Teórica: Provas abrangentes são fornecidas quanto à não viés dos estimadores e às propriedades de redução de variância das regras de atualização propostas.

Resultados Experimentais
Extensos experimentos foram conduzidos em múltiplos LLMs (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) em várias tarefas downstream (por exemplo, HellaSwag, RACE, PIQA) e padrões de esparsidade ((2:4), (4:8), (8:16)).

• Desempenho: O MaskPro supera consistentemente os métodos baseados em regras mais avançados (SparseGPT, Wanda, Pruner-Zero) e alcança desempenho comparável ao MaskLLM, que consome muita memória, frequentemente com uma melhoria média superior a 2% em precisão.
• Escalabilidade de Memória: Enquanto o MaskLLM encontra erros de Memória Insuficiente (OOM) em 8x GPUs A100 para esparsidade (4:8) em modelos de 7B, o MaskPro mantém um uso de memória estável (~36-37 GB) independentemente do padrão de esparsidade.
• Eficiência de Dados: O MaskPro demonstra robustez notável ao volume de dados. Ele alcança desempenho estável mesmo quando treinado com apenas 1 amostra de dados, enquanto o MaskLLM requer centenas de milhares de amostras para convergir.
• Eficiência de Treinamento: O método requer significativamente menos amostras de treinamento e menos tempo de treinamento em comparação com o MaskLLM, mantendo ao mesmo tempo um uso de memória comparável aos métodos baseados em regras.

Significado e Alegações
O artigo afirma que o MaskPro fornece um caminho viável para a implantação local prática e escalável de LLMs, resolvendo os gargalos de memória e estabilidade do aprendizado de esparsidade semi-estruturada baseada em aprendizado. Os autores enfatizam que sua abordagem preenche a lacuna entre o alto desempenho dos métodos baseados em aprendizado e a eficiência dos métodos baseados em regras. Ao reduzir a pegada de memória para espaço linear e estabilizar o processo de treinamento por meio de gradientes de política refinados, o MaskPro permite o aprendizado de máscaras (N:M) ótimas em modelos de grande escala sem a sobrecarga computacional proibitiva anteriormente associada a tais tarefas. O método é apresentado como um framework generalizável para aprendizado de máscaras probabilísticas que é robusto à escassez de dados e à escala do modelo.