STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

Resumo Técnico: STARFISH

Definição do Problema
A poda (pruning) de redes neurais é uma técnica padrão para reduzir o tamanho do modelo, a pegada de memória e a latência de inferência através da remoção de pesos, neurônios ou canais. No entanto, uma poda agressiva frequentemente leva a uma degradação substancial na acurácia preditiva. Embora processos de recuperação (cura/healing) sejam tipicamente empregados para retomar o desempenho perdido, os métodos existentes enfrentam limitações significativas. Abordagens tradicionais muitas vezes exigem o retreinamento completo ou o ajuste fino (fine-tuning) no conjunto de dados de treinamento original, o que é computacionalmente caro e frequentemente impossível em cenários modernos onde os dados de treinamento são proprietários ou inacessíveis. Além disso, os métodos de recuperação de última geração costumam ter dificuldades para restaurar a acurácia sob condições de alta esparsidade (por exemplo, removendo >75% dos pesos), particularmente para arquiteturas modernas de Vision Transformer (ViT).

Metodologia: STARFISH
Os autores propõem o STARFISH (faST Accuracy Recovery in pruned networks From Internal State Healing), um método de recuperação pós-poda que opera sem acesso aos dados de treinamento originais ou aos rótulos. A hipótese central é que otimizar o alinhamento das representações intermediárias (estados internos) entre a rede podada e a rede densa original é um proxy mais eficaz para a restauração da acurácia do que otimizar apenas nas predições de saída finais.

O algoritmo procede em duas etapas usando um pequeno conjunto de calibração não rotulado ($S_{cal}$) amostrado da distribuição de dados, porém distinto do conjunto de treinamento:

1. Cache de Representações Alvo: O modelo denso original é executado em $S_{cal}$ para realizar o cache das saídas vetorizadas de cada bloco computacional (por exemplo, após normalizações, não-linearidades e adições residuais).
2. Otimização de Alinhamento: O modelo podado ($f_{\theta \odot m}$) é treinado para minimizar a perda de alinhamento entre suas representações intermediárias e os alvos em cache. A função de perda é definida como a distância de cosseno média entre os estados ocultos do modelo podado ($f^\ell_{\theta \odot m}(x)$) e os estados ocultos do modelo denso ($r^\ell_x$) em todos os blocos $\ell$:
   $$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \odot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$$
   Os pesos restantes do modelo podado são atualizados para minimizar a perda empírica sobre o conjunto de calibração.

A motivação teórica é fornecida ao demonstrar que a distância entre as representações ocultas finais (entradas para a cabeça de classificação) serve como um limite superior para a divergência de Kullback–Leibler (KL) das distribuições de saída. Isso sugere que alinhar os estados internos restringe inerentemente o erro de saída.

Principais Contribuições

1. Recuperação Baseada em Representação: O artigo introduz um método de recuperação independente de tipos de poda, máscaras ou arquiteturas específicas. Demonstra empiricamente que o alinhamento de representações intermediárias supera os objetivos baseados diretamente na saída (como Cross-Entropy ou correspondência de divergência KL) à medida que os níveis de esparsidade aumentam.
2. Eficiência em Alta Esparsidade: O método é o primeiro a restaurar eficientemente altas porcentagens da acurácia original em níveis extremos de esparsidade (ex: 85% de esparsidade). Ele alcança isso usando um minúsculo conjunto de calibração (tão pouco quanto 1.000 exemplos, ou 0,4% do tamanho do conjunto de treinamento) e não requer acesso aos dados de treinamento originais.
3. Modularidade e Velocidade: A abordagem é computacionalmente leve, exigindo apenas algumas passagens de forward e backward em um pequeno conjunto de dados. É aplicável a padrões de poda estruturados, semiestruturados e não estruturados.

Resultados Experimentais
Os autores avaliaram o STARFISH no ImageNet-1K usando várias arquiteturas de Vision Transformer (DeiT, ViT) e Redes Neurais Convolucionais (MobileNet).

• Alta Esparsidade (0.75–0.85): Em um modelo DeiT-B podado para 75% de esparsidade, o STARFISH recuperou 82% da acurácia original densa, enquanto métodos concorrentes (como o CORP) recuperaram apenas 40%. Em um modelo DeiT3-H com 85% de esparsidade, o STARFISH reteve 92,08% da acurácia densa, melhorando o desempenho do modelo podado em 72,36 pontos percentuais.
• Esparsidade Moderada (0.5): Em várias arquiteturas DeiT e escopos de poda (MLP, QK e conjunta MLP+QK), o STARFISH superou consistentemente os métodos de última geração. Por exemplo, em um modelo DeiT-T com 50% de poda estruturada em MLP+QK, o STARFISH alcançou 81,95% de retenção da acurácia densa, uma melhoria de 31,54 pontos percentuais sobre o método CORP.
• Comparação: O STARFISH demonstrou desempenho superior contra métodos como o SNOWS (que utiliza aproximação de Hessiana) e o CORP (que utiliza recuperação afim de forma fechada), particularmente em regimes onde a poda degrada severamente o modelo.

Significância e Alegações
O artigo afirma que o STARFISH representa um avanço significativo no campo da compressão de modelos ao deslocar o paradigma de recuperação do correspondência de saída para a cura de estados internos. Os autores enfatizam que esta abordagem é particularmente vital para modelos de grande escala modernos, onde os dados de treinamento não estão disponíveis e o retreinamento total é inviável. Ao alavancar o princípio de que as representações internas capturam o conhecimento do modelo de forma mais robusta do que as saídas finais sob alta esparsidade, o STARFISH permite a implantação de redes podadas altamente eficientes sem sacrificar o desempenho preditivo. O método é apresentado como uma solução geral, plug-and-play, que pode ser aplicada a qualquer rede podada, independentemente da estratégia de poda inicial.