Resumo Técnico: AIRE-Prune

Declaração do Problema

Modelos de Espaço de Estados (SSMs) Profundos têm demonstrado eficácia na modelagem de dados sequenciais ao comprimir o histórico de entrada em estados internos. No entanto, a dimensão do estado $n$ permanece como um principal impulsionador dos custos computacionais e de memória. Embora as arquiteturas modernas frequentemente utilizem parametrizações diagonais ou diagonalizáveis (por exemplo, empilhamentos multi-SISO ou MIMO como o S5) para garantir estabilidade e eficiência, esses modelos frequentemente sofrem de sobreparametrização. Os mecanismos existentes de pós-treinamento para otimizar $n$ são limitados. Abordagens anteriores, como a poda de pesos não estruturada, não reduzem diretamente a dimensão do estado. Outros métodos estruturados, como o recente LAST (Gwak et al., 2025), baseiam-se no ganho de frequência de pior caso ($H_\infty$) para pontuar estados. Embora eficaz, medidas de pior caso podem ser excessivamente conservadoras para cargas de trabalho típicas, pois enfatizam amplificações de pico que as tarefas raramente excitam, potencialmente levando a razões de compressão subótimas.

Metodologia: AIRE-Prune

Os autores introduzem o AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)ing), um método de poda estruturado, pós-treinamento e adaptável por camada, projetado para SSMs diagonais (ou diagonalizáveis). Diferente de métodos anteriores que focam no ganho de pior caso, o AIRE-Prune alinha a poda com a energia assintótica da resposta ao impulso, um critério de caso típico que representa a energia total de saída que um estado contribui ao longo de um horizonte de tempo infinito.

Mecanismo Central

1. Classificação Baseada em Energia:
   Para uma camada SSM diagonal de tempo discreto estável definida por $(\Lambda, B, C)$, onde $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ e $|\lambda_i| < 1$, o método calcula a energia total de saída ao quadrado gerada por uma entrada de impulso unitário para cada estado (modo) $i$.
   O score de forma fechada para o $i$-ésimo modo é derivado como:
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   Este score integra três fatores:

   • Controlabilidade: Representada por $\|B_{i,:}\|_2^2$ (facilidade de excitação).
   • Observabilidade: Representada por $\|C_{:,i}\|_2^2$ (facilidade de medição).
   • Amortecimento/Memória: Representada pelo polo $\lambda_i$; estados com polos próximos ao círculo unitário (decaimento mais lento) contribuem mais para a energia de longo prazo.

2. Normalização entre Camadas:
   Como as escalas de energia podem variar significativamente entre diferentes camadas (por exemplo, devido aos ganhos de codificador/decodificador), os autores propõem uma estratégia de pontuação com prefixo normalizado para permitir a comparação global:

   • Dentro de cada camada $\ell$, os modos são ordenados por sua energia local $E^{(\ell)}_i$ em ordem decrescente.
   • Uma soma de prefixo $S^{(\ell)}_{(i)}$ é computada para as energias ordenadas.
   • O score global é definido como a razão de "taxa de risco" (hazard-rate):
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     Esta normalização garante que os scores sejam monotonicamente não crescentes dentro de uma camada, permitindo que um único limiar global $\tau$ seja aplicado a todo o modelo.

3. Execução da Poda:
   Dado um alvo de razão de poda ou limiar, o algoritmo retém o prefixo mais longo de estados em cada camada onde o score normalizado excede $\tau$. Isso permite a poda adaptável por camada, onde algumas camadas podem ser fortemente podadas ou até inteiramente removidas se seus estados caírem abaixo do limiar global, enquanto outras permanecem amplamente intactas.

Principais Contribuições

• Score de Energia Assintótica de Forma Fechada: O artigo fornece um score derivado matematicamente e de forma fechada para cada modo aprendido com base na energia de saída de horizonte infinito. Isso evita a necessidade de simulações numéricas dispendiosas ou retreinamento.
• Extensão da Truncagem Modal: O trabalho estende a truncagem modal clássica (tipicamente aplicada a sistemas lineares únicos) para empilhamentos profundos de SSMs não lineares, preservando a estrutura diagonal essencial para a eficiência.
• Caso Típico vs. Pior Caso: Ao mudar a métrica de importância do ganho $H_\infty$ de pior caso para a energia total da resposta ao impulso, o método visa o comportamento agregado do sistema em vez de amplificações de pico raras.
• Agnóstico à Arquitetura: O método é aplicável tanto a construções multi-SISO quanto a empilhamentos MIMO (como o estilo S5), bem como a outras famílias de SSM como S4D e Mamba.

Resultados Experimentais

Os autores avaliaram o AIRE-Prune no benchmark Long Range Arena (LRA) e no conjunto de dados Speech Commands usando backbones S5, S4D e Mamba.

• Eficiência de Compressão: Em diversos benchmarks, o AIRE-Prune alcançou uma poda de estado média de 60,8% com uma queda de precisão média de apenas 0,29% sem qualquer retreinamento.
• Comparação com Baselines:
  • No conjunto LRA, o AIRE-Prune superou significativamente os baselines incluindo Random, Uniform Magnitude, Global Magnitude, LAMP, Uniform $H_\infty$, Global $H_\infty$ e LAST.
  • Sob uma restrição estrita de $\le 1$ ponto percentual (pp) de perda de precisão, o AIRE-Prune permitiu razões de poda substancialmente maiores do que o baseline mais forte (LAST). Por exemplo, na tarefa Path-X, o AIRE-Prune podou com segurança 70% dos estados comparado a 30% para o LAST.
  • Na tarefa Speech Commands, 45% dos estados foram podados com perda de precisão desprezível (96,40% vs. 96,43% para o modelo completo).
• Eficiência de Inferência: A poda traduziu-se diretamente em ganhos no nível do sistema. Em uma GPU NVIDIA H100, o método rendeu acelerações de inferência variando de 1,23$\times$ a 2,86$\times$ e reduções de parâmetros de 19% a 65%, dependendo da tarefa e da razão de poda.
• Generalização: O método demonstrou desempenho robusto nas arquiteturas S4D e Mamba, alcançando até 90% de poda na tarefa Text para S4D e 80% na tarefa Image para Mamba com mínima degradação de precisão.

Significância e Alegações

O artigo afirma que o AIRE-Prune revela uma redundância substancial nos espaços de estado de SSMs treinados, que pode ser removida sem retreinamento. Ao alinhar a poda com a energia de resposta assintótica em vez do ganho de pior caso, o método alcança uma compressão agressiva com custo desprezível para o desempenho da tarefa.

Os autores destacam que o perfil de queda de precisão "em degraus" observado com o AIRE-Prune (onde a precisão permanece estável até que um alto limiar de poda seja atingido, então cai bruscamente) indica uma separação forte entre estados importantes e não importantes. Isso contrasta com os baselines que mostram uma degradação suave, sugerindo que seus métodos de pontuação misturam estados importantes e não importantes.

Praticamente, o método permite:

1. Orçamento Global: Definir um orçamento alvo único sem necessidade de ajuste por camada.
2. Redução em Nível de Bloco: A capacidade de podar camadas inteiras quando seus estados caem abaixo do limiar, convertendo esparsidade fina em reduções estruturadas em blocos, que são mais amigáveis ao hardware.
3. Otimização Sem Treinamento: Um mecanismo de pós-treinamento que não requer ajuste fino adicional, tornando-o imediatamente aplicável a modelos implantados.

O trabalho conclui que o AIRE-Prune oferece uma abordagem fundamentada, baseada em energia, para a redução de ordem de modelos em SSMs profundos, superando as estratégias de poda de estado da arte atuais tanto em razão de compressão quanto em preservação de precisão.