Resumo Técnico: Repensando a Poda de Profundidade para Vision Transformers: Uma Perspectiva Consciente da Heterogeneidade

1. Definição do Problema

Os Vision Transformers (ViTs) alcançaram o estado da arte em vários domínios, mas sofrem com altos custos computacionais e grandes contagens de parâmetros, levando a uma latência de inferência estendida. Embora a poda estruturada seja uma estratégia de compressão eficaz, os métodos existentes enfrentam limitações significativas:

• Poda de Largura (Width Pruning): Reduz canais ou cabeças de atenção dentro das camadas, mas frequentemente produz ganhos de velocidade limitados em hardwares convencionais.
• Poda de Profundidade (Depth Pruning): Remove camadas inteiras, oferecendo um potencial de aceleração significativamente maior (como mostrado na Figura 2a). No entanto, é notoriamente difícil recuperar a acurácia, especialmente com a remoção agressiva de camadas.
• O Gargalo Central: O artigo argumenta que a falha dos métodos de poda de profundidade existentes decorre de dois problemas negligenciados:
  1. Negligência da Heterogeneidade: Os ViTs contêm tipos de camadas distintas (por exemplo, camadas de Atenção vs. camadas de Função de Ativação dentro das Redes Feed-Forward). Os métodos existentes costem tratar essas camadas de forma uniforme ou dependem de métricas baseadas em gradientes que falham devido à disparidade significativa nas escalas de gradiente entre esses tipos.
  2. Desajuste de Dimensão (Dimension Mismatch): Podar simultaneamente as camadas de Atenção e as camadas Lineares da Rede Feed-Forward (FFN) cria incompatibilidades estruturais. Como as FFNs utilizam um mecanismo de expansão-contração (ex: $D \to 4D \to D$), remover uma camada Linear quebra a continuidade da forma do tensor necessária para as camadas subsequentes (ilustrado na Figura 3).

2. Metodologia: HetDPT

Os autores propõem o HetDPT (Heterogeneity-aware DePTh pruning), um framework de dois estágios projetado para abordar a heterogeneidade e evitar o desajuste de dimensão.

Insights Chave

O método baseia-se em duas observações críticas em relação à heterogeneidade das camadas:

1. Disparidade de Gradiente: As camadas de Atenção e as camadas de Ativação exibem magnitudes de gradiente vastamente diferentes durante o backpropagation. O uso de métricas de importância baseadas em gradientes padrão leva a decisões de poda enviesadas, favorecendo consistentemente a poda de um tipo sobre o outro.
2. Assimetria de Recuperação:
   • A poda de camadas de Atenção causa quedas iniciais moderadas de acurácia, mas requer um retreinamento extensivo para recuperação.
   • A poda de camadas de Ativação causa uma degradação inicial severa (frequentemente $\ge 90\%$ de perda de acurácia), mas permite uma recuperação rápida durante o fine-tuning.
   • Implicação: Métricas míopes que avaliam a acurácia imediata pós-poda são ineficazes.

O Pipeline de Dois Estágios

Estágio 1: Identificação de Camadas Redundantes
Este estágio determina quantas e quais camadas específicas devem ser podadas sem comparar diretamente tipos heterogêneos de camadas.

• Passo 1: Alocação de Orçamento de Poda (Model Accuracy Predictor - MAP): Em vez de comparar camadas diretamente, o método utiliza um processo de coleta de dados leve (poda iterativa $\to$ fine-tune rápido $\to$ avaliação) para construir um conjunto de dados. Uma função polinomial aproxima o Model Accuracy Predictor (MAP), que prevê a acurácia final de um modelo podado dada as razões de poda das camadas de Atenção ($\tilde{m}_a$) e de Ativação ($\tilde{m}_g$). Isso permite que o sistema encontre a alocação ideal do orçamento de poda entre os dois tipos de camadas.
• Passo 2: Remoção de Camadas Redundantes Específicas: Uma vez que o orçamento é alocado, o método identifica camadas específicas para remover dentro de cada grupo homogêneo (por exemplo, apenas entre as camadas de Atenção). Utiliza-se parâmetros de importância aprendíveis e propagação de gradiente para classificar as camadas, garantindo que não ocorra comparação entre tipos distintos.

Estágio 2: Otimização do Modelo Podado

• Fine-tuning: O modelo podado é submetido a fine-tuning, opcionalmente usando autodestilação guiada pelo ViT original não podado para restaurar a acurácia.
• Fusão (Merging): Para resolver o desajuste de dimensão e acelerar a inferência, camadas Lineares adjacentes que foram separadas por uma camada de Ativação podada são fundidas. Por exemplo, se uma camada de Ativação for removida, $L_{l,1}$ e $L_{l,2}$ são fundidas em uma única camada linear $L_l$. Isso reduz a profundidade do modelo enquanto mantém o alinhamento rigoroso de dimensões.

HetDPT+: Os autores integram ainda o HetDPT com métodos estabelecidos de poda de largura para criar o HetDPT+, um pipeline para compressão extrema de ViT que aborda as dimensões de profundidade e largura.

3. Principais Contribuições

1. Poda de Funções de Ativação: O artigo propõe a primeira estratégia para remover camadas de Função de Ativação (GELU) em ViTs. Ao remover essas não-linearidades, as camadas Lineares adjacentes podem ser naturalmente fundidas, resolvendo o problema de desajuste de dimensão inerente à poda de profundidade heterogênea.
2. Análise de Heterogeneidade e Solução: Os autores revelam e abordam a disparidade de gradiente e a assimetria de recuperação. Eles introduzem um framework de dois estágios (HetDPT) equipado com um Model Accuracy Predictor (MAP) para gerenciar essas heterogeneidades, apoiado por análise teórica que garante a existência de uma configuração ótima e a consistência do preditor sob fine-tuning rápido e ruidoso.
3. Recordes de Estado da Arte (SOTA):
   • Poda de Profundidade: O HetDPT alcança um ganho de velocidade de 1.58$\times$ para o DeiT-B e 1.39$\times$ para o DeiT-S com acurácia quase sem perdas, superando trabalhos existentes de poda de profundidade.
   • Compressão Extrema: O HetDPT+ estabelece um novo marco para a compressão extrema de ViT. Para a configuração Isomorphic-Pruning-2.6G, ele aumenta a razão de aceleração de 4.24$\times$ para 5.19$\times$ mantendo uma acurácia quase sem perdas.
4. Análise de Resultados Experimentais:
   • Troca Acurácia-Velocidade: No ImageNet-1K, o HetDPT supera consistentemente a poda isolada de camadas de atenção ou ativação. Para o DeiT-B, atinge 81.8% de acurácia Top-1 (correspondendo ao baseline) com um ganho de velocidade de 1.58$\times$.
   • Comparação com SOTA: O HetDPT+ supera métodos de compressão extrema existentes (ex: Isomorphic Pruning, NViT). Com ~81% de acurácia, o HetDPT+ alcança 5.19$\times$ de aceleração comparado aos 4.24$\times$ do Isomorphic Pruning. Em configurações ultra-extremas (1.1G MACs), alcança 9.19$\times$ de aceleração.
   • Transferibilidade: O método demonstra transferência de aprendizado robusta no CIFAR-100 e mantém o desempenho em tarefas de predição densa (detecção de objetos COCO e segmentação semântica ADE20K) com degradação mínima de acurácia.
   • Escalabilidade: A abordagem escala para modelos de bilhões de parâmetros, demonstrado no DINO-V2-Giant, onde poda 15 camadas com apenas 0.06% de queda na acurácia.
   • Ortogonalidade: O HetDPT é ortogonal a métodos de poda de tokens (ex: GTP-ViT, ToMe), proporcionando ganhos de velocidade adicionais quando combinados.

4. Significância e Alegações

O artigo alega que a principal contribuição é viabilizar o deploy mais eficiente de Vision Transformers ao abordar os desafios estruturais e dinâmicos específicos da poda de profundidade.

• Democratização: Ao reduzir substancialmente os requisitos computacionais (até 9.19$\times$ de aceleração), o método torna modelos de visão de estado da arte acessíveis a pesquisadores e profissionais com recursos limitados.
• Sustentabilidade: A redução do consumo de energia em modelos comprimidos contribui para um desenvolvimento de IA mais sustentável ambientalmente.
• Deploy em Edge: Os ganhos de eficiência facilitam o deploy em tempo real em dispositivos de borda (edge) para aplicações que exigem baixa latência, como diagnósticos médicos e sistemas autônomos.

Os autores observam uma limitação: o Model Accuracy Predictor (MAP) é atualmente reajustado para cada família de arquitetura, e estender o método para modelos de linguagem ou outras modalidades permanece como um trabalho futuro. No entanto, este trabalho estabelece um novo fundamento teórico e prático para a poda de profundidade heterogênea, indo além das abordagens de "tamanho único" da literatura anterior.