COREY: A Prototype Study of Entropy-Guided Operator Fusion with Hadamard Reparameterization for Selective State Space Models

Título: COREY: Um Estudo de Protótipo de Fusão de Operadores Guiada por Entropia com Reparametrização de Hadamard para Modelos de Espaço de Estado Seletivos (SSMs)

1. O Problema

Os Modelos de Espaço de Estado Seletivos (SSMs), representados pela família Mamba, oferecem modelagem de sequências em tempo linear ($O(L)$), tornando-se atraentes para inferência em contextos longos. No entanto, a implantação prática enfrenta dois gargalos principais:

1. Limitação de Largura de Banda de Memória: As atualizações de estado seletivo são frequentemente decompostas em múltiplos kernels fragmentados, resultando em materialização repetida de tensores intermediários e tráfego excessivo de memória global (DRAM). Isso transforma o caminho crítico em um trabalho limitado por memória, não por computação.
2. Instabilidade em Baixa Precisão (Outliers): As ativações internas dos SSMs, especialmente nos canais de saída, apresentam estatísticas de cauda pesada (heavy-tailed). Esses "outliers" desestabilizam a quantização de baixa precisão (ex: W4A8, W8A8) e limitam a profundidade segura da fusão de operadores devido à sensibilidade numérica amplificada.

2. Metodologia: O Framework COREY

O COREY é um framework de protótipo que unifica a otimização de sistemas (fusão de operadores) com o ajuste de distribuição de ativações (reparametrização).

• Fusão de Operadores Guiada por Entropia:

  • Em vez de usar limites de profundidade estáticos ou heurísticas fixas, o COREY utiliza a entropia de ativação (estimada via histogramas de largura fixa) como um sinal de agendamento em tempo de execução.
  • Uma pontuação de fusão $S(R)$ é calculada combinando entropia normalizada, intensidade aritmética e estimativas de tráfego de memória.
  • Regiões com alta entropia (distribuição mais uniforme) permitem tiles (blocos) maiores e fusão mais agressiva, enquanto regiões de baixa entropia (com outliers) exigem tiles menores para evitar sobrecarga de registradores e instabilidade.

• Reparametrização com Hadamard Fundido:

  • Para suavizar as distribuições de ativação de cauda pesada, o método insere uma transformação de Hadamard normalizada ($H$) nas projeções lineares: $y = W H^\top H x$.
  • Isso preserva a equivalência funcional, mas redistribui a energia dos outliers através dos canais, reduzindo a concentração de coordenadas de pico.
  • A transformação de Hadamard, a projeção linear e os operadores elementares adjacentes são fundidos em um único kernel para evitar passagens extras de memória.

• Co-design Entropia-Tiling:

  • O sistema ajusta dinamicamente o tamanho dos tiles com base na entropia local. Alta entropia $\rightarrow$ tiles maiores (melhor reutilização de dados); baixa entropia $\rightarrow$ tiles menores (proteção contra instabilidade).

3. Principais Contribuições

1. Critério de Fusão Dinâmica: Introdução de um critério baseado em entropia para selecionar dinamicamente os limites de fusão em kernels de SSM, conectando a geometria da ativação à viabilidade da fusão.
2. Camada de Hadamard Fundida: Design de uma camada que incorpora transformações diretas e inversas de Hadamard nos blocos de projeção, preservando a equivalência funcional enquanto reduz a concentração de picos.
3. Conexão Teórica e Prática: Estabelecimento de um vínculo explícito entre o ganho de entropia (pós-rotação de Hadamard) e a supressão de outliers, quantificando como isso expande a região de operação segura para fusão agressiva e execução de baixa precisão.
4. Validação de Protótipo: Validação do design através de um estudo de protótipo reproduzível e verificações de saneamento (sanity checks) em checkpoints reais, separando claramente os resultados de pesquisa de reivindicações de implantação em nível de hardware.

4. Resultados

O estudo foi conduzido em um ambiente controlado (protótipo em Python/PyTorch simulando agendamento) e validado parcialmente com checkpoints reais (Mamba-370M, Mamba-1.4B, Mamba-2.8B).

• Desempenho (Latência e Throughput):

  • O COREY reduziu consistentemente a latência e aumentou o throughput (tokens/segundo) em comparação com bases sem fusão e fusão estática.
  • A redução no tráfego de DRAM por token foi significativa, confirmando que a melhoria vem da redução de materialização de memória.
  • Em contextos ultra-longos (64k), o ganho de throughput foi de aproximadamente 1.34x em relação à fusão estática e 1.12x em relação à fusão estática pura.

• Estabilidade em Baixa Precisão:

  • Em configurações W8A8 e W4A8, a reparametrização de Hadamard fundida reduziu o "proxy de degradação de qualidade" (uma métrica diagnóstica interna baseada em sensibilidade à quantização).
  • Isso indica que o controle de limites baseado em entropia evita a fusão excessiva em regiões dominadas por outliers, mantendo a estabilidade numérica.

• Validação de Checkpoint:

  • Resultados em benchmarks como LongBench (NarrativeQA, Qasper, etc.) e perplexidade em WikiText-103/PG19 mostraram que a qualidade do modelo é preservada entre diferentes hardwares (RTX 3070 e RTX 3090) e que o agendamento não introduz regressões na precisão da tarefa.

5. Significado e Impacto

• Ponte entre Teoria e Sistemas: O trabalho demonstra que a otimização de sistemas (fusão de kernels) pode ser guiada por propriedades estatísticas das ativações (entropia), criando um ciclo de feedback onde a suavização da distribuição permite maior eficiência de hardware.
• Eficiência Energética: Ao reduzir o tráfego de memória, o método promete diminuir o custo energético da inferência de SSMs em contextos longos, um fator crítico para a sustentabilidade de modelos de IA.
• Viabilidade de Implantação: Embora o estudo atual seja um protótipo, ele fornece as bases teóricas e empíricas para a implementação de kernels fundidos em nível de produção (ex: via Triton), sugerindo que a fusão agressiva é viável se acompanhada de técnicas de suavização de ativação.
• Limitações Atuais: O artigo enfatiza que os resultados de baixa precisão são diagnósticos e que a validação final de checkpoint em nível de implantação (com kernels fundidos reais e comparações diretas de qualidade) ainda é trabalho futuro.

Em resumo, o COREY propõe uma mudança de paradigma: em vez de tratar a fusão de operadores apenas como um problema de alocação de memória, ele utiliza a entropia das ativações como um sinal de controle para equilibrar agressividade de fusão e estabilidade numérica, permitindo inferência mais rápida e eficiente em modelos de estado de estado seletivos.