HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

O artigo apresenta o HiAP, um framework de auto-podagem estocástica multi-granular que otimiza Vision Transformers em uma única fase de treinamento, utilizando portas Gumbel-Sigmoid hierárquicas para descobrir sub-redes eficientes que equilibram precisão e recursos computacionais sem a necessidade de heurísticas manuais ou pipelines complexos.

O essencial

Imagine que você tem um restaurante de luxo (o modelo de Inteligência Artificial chamado "Vision Transformer") que é incrível para cozinhar pratos complexos, mas é tão grande, caro e lento que ninguém consegue levá-lo para um pequeno carrinho de comida de rua (os dispositivos móveis ou "edge devices").

O problema é que esse restaurante tem muitos chefs, muitas mesas e muitos ingredientes, e a maior parte do tempo e dinheiro é gasta apenas transportando tudo isso de um lugar para outro, e não necessariamente cozinhando.

Aqui está como o HiAP (o método proposto no artigo) resolve esse problema, explicado de forma simples:

1. O Problema: O Restaurante Superlotado

Os modelos atuais de IA são como restaurantes que têm:

• Muitos chefs de cabeça (Attention Heads): Às vezes, eles fazem a mesma coisa.
• Muitos ajudantes (Neurônios): Às vezes, eles não são necessários.
• O gargalo: O maior custo não é nem cozinhar (cálculo), é carregar os ingredientes pesados da despensa para a cozinha (memória). Se você apenas demitir alguns ajudantes (pruning tradicional), o restaurante ainda tem que carregar todas as mesas e cadeiras vazias, o que continua lento.

2. A Solução: O "Gerente Automático" (HiAP)

Antes, para cortar o tamanho do restaurante, os donos tinham que fazer uma lista manual de quem demitir, testar, ajustar e testar de novo (métodos complexos e manuais).

O HiAP é como um Gerente Inteligente e Automático que entra no restaurante e faz duas coisas ao mesmo tempo, durante o treinamento:

A. O Nível Macro (Demitição de Departamentos Inteiros)

O gerente olha para o restaurante e decide: "Esse departamento inteiro de chefs de salada não está sendo usado. Vamos fechar a porta e parar de gastar energia transportando os ingredientes desse setor."

• Na prática: Ele remove cabeças inteiras de atenção ou blocos inteiros da rede. Isso economiza o transporte pesado (memória).

B. O Nível Micro (Ajuste Fino das Cozinhas)

Dentro dos departamentos que ficaram, o gerente olha para cada chef individualmente e diz: "Você, Chef João, só precisa de 3 facas, não de 10. Vamos tirar as 7 facas extras."

• Na prática: Ele reduz o tamanho interno das ferramentas (dimensões dentro da cabeça de atenção) e o número de neurônios. Isso economiza o esforço de cozinhar (cálculo).

3. A Mágica: A "Porta Mágica" (Gumbel-Sigmoid)

Como o gerente sabe quem demitir sem demitir alguém e estragar o prato?
Ele usa uma porta mágica (chamada de Gumbel-Sigmoid).

• No começo, a porta é "borrada" (pode estar 50% aberta). O restaurante ainda funciona, mas o gerente está aprendendo.
• Conforme o tempo passa, a porta vai ficando mais rígida, até virar um 0 ou 1 (fechada ou aberta).
• O segredo é que o restaurante aprende a cozinhar bem mesmo com as portas fechando. Os chefs restantes se adaptam e ficam mais eficientes.

4. O Resultado: Um Restaurante Compacto e Pronto

Ao final do treinamento, o HiAP não entrega apenas uma lista de demissões. Ele entrega um novo restaurante físico, menor, mais leve e pronto para funcionar imediatamente.

• Sem retrabalho: Não precisa de uma segunda fase para "consertar" o sabor (fine-tuning).
• Equilíbrio Perfeito: Ele descobre sozinho onde economizar mais (se é fechar um departamento inteiro ou apenas reduzir as ferramentas) para ficar o mais rápido possível sem perder a qualidade do prato (precisão).

Resumo da Ópera

O HiAP é como ter um arquiteto que redesenha o prédio enquanto você mora nele, sem precisar de você dizer "demitam 30% dos funcionários". Ele decide sozinho o que é essencial e o que é luxo desnecessário, criando uma versão do modelo que é rápida, leve e cabe perfeitamente no seu celular, mantendo a inteligência original.

Em poucas palavras: O HiAP ensina a IA a se "podar" sozinha, de cima a baixo, transformando um gigante lento em um ágil e eficiente, sem precisar de ajuda humana para decidir o que cortar.

Resumo técnico

Título: HIAP: Um Framework de Auto-Poda Estocástica Multi-Granular para Vision Transformers

1. O Problema

Os Vision Transformers (ViTs) tornaram-se a arquitetura dominante na visão computacional, mas enfrentam desafios significativos para implantação em dispositivos de borda (edge devices) devido ao seu alto custo computacional e exigência de largura de banda de memória.

• Limitações das Métodos Atuais:
  • Granularidade Única: A maioria dos métodos de poda estruturada opera apenas em uma escala (ou apenas em "micro-estruturas" como dimensões internas, ou apenas em "macro-estruturas" como cabeças de atenção inteiras). A poda apenas micro-preserva a profundidade e o número de cabeças, mantendo o alto custo de acesso à memória (overhead de DRAM/SRAM). A poda apenas macro pode causar perda excessiva de capacidade de representação.
  • Dependência de Heurísticas Manuais: Métodos existentes frequentemente exigem pipelines complexos de múltiplos estágios, incluindo poda inicial, ranking de importância manual, definição de limites de esparsidade por camada e ajuste fino (fine-tuning) pós-hoc para recuperar a acurácia.

2. Metodologia: HiAP (Hierarchical Auto-Pruning)

O HiAP propõe um framework de relaxação contínua que descobre sub-redes ótimas em uma única fase de treinamento end-to-end, sem necessidade de heurísticas manuais ou alvos de esparsidade pré-definidos.

• Mecanismo de Portas Hierárquica (Gating):
  O método introduz portas estocásticas baseadas na distribuição Gumbel-Sigmoid em dois níveis de granularidade simultaneamente:

  1. Nível Macro: Controla a retenção de unidades inteiras, como cabeças de atenção completas ($g_{l,h}$) e blocos de FFN (Feed-Forward Networks) ($b_l$). Se uma porta macro for 0, o bloco inteiro é contornado ou removido, eliminando o overhead de memória de carregar matrizes grandes.
  2. Nível Micro: Opera dentro das estruturas macro ativas, podando seletivamente dimensões intra-cabeça ($d_{l,h,j}$) e neurônios individuais do FFN ($c_{l,k}$). Isso reduz as operações matemáticas (FLOPs).

• Modelagem de Custo Diferenciável:
  O HiAP formula um custo exato e diferenciável de operações Multiplicação-Acumulação (MACs) que é incorporado diretamente à função de perda.

  • O custo é decomposto em componentes estáticos (overhead de camadas não podadas) e dinâmicos (governados pelas portas).
  • A função de perda penaliza o custo computacional esperado, permitindo que o modelo "aprenda" a alocação de recursos para satisfazer restrições de hardware.

• Regularização de Viabilidade Estrutural:
  Para evitar o colapso estrutural (onde o modelo poda camadas inteiras prematuramente para minimizar o custo), o HiAP utiliza penalidades explícitas que garantem um número mínimo de cabeças, dimensões e neurônios ativos em cada camada, assegurando um fluxo de gradiente estável.

• Descoberta em Fase Única (Single-Phase):
  Diferente de métodos que exigem busca e ajuste fino separados, o HiAP treina a rede densa e as portas de poda simultaneamente. A temperatura $\tau$ da Gumbel-Sigmoid é reduzida (annealing) ao longo do treinamento, transformando as portas contínuas em decisões binárias determinísticas. Ao final, a sub-rede é extraída fisicamente, sem necessidade de ajuste fino secundário.

3. Contribuições Principais

1. Unificação Multi-Granular: O primeiro framework a unificar a poda de macro-estruturas (blocos e cabeças) e micro-estruturas (neurônios e dimensões) em um único sistema de portas diferenciável.
2. Auto-Descoberta de Arquitetura: Elimina a necessidade de heurísticas de importância manuais, métricas de proxy ou pipelines de múltiplos estágios. O modelo descobre autonomamente sua sub-arquitetura ótima.
3. Eficiência de Implantação: Gera sub-redes fisicamente comprimidas e densas prontas para inferência imediata em hardware padrão, sem depender de engines de convolução esparsa.
4. Validação Empírica: Demonstração de que o HiAP alcança fronteiras de Pareto competitivas entre acurácia e eficiência em CIFAR-10 e ImageNet.

4. Resultados Experimentais

Os experimentos foram realizados no ImageNet-1K (usando DeiT-Small) e CIFAR-10 (ViT-Tiny).

• Desempenho no ImageNet (DeiT-Small):

  • O HiAP reduziu o custo computacional de 4.6G MACs para 3.1G MACs (redução de ~33%) mantendo uma acurácia Top-1 de 79.10% (apenas -0.75% de perda em relação à base densa de 79.85%).
  • Competiu favoravelmente com métodos de ponta como GOHSP e ViT-Slim, mas com um pipeline significativamente mais simples (uma única fase vs. múltiplos estágios complexos).
  • Em configurações mais agressivas (2.5G MACs), manteve 77.95% de acurácia.

• Desempenho no CIFAR-10:

  • Superou heurísticas manuais (como ranking por norma L1 e proporção uniforme) em regimes de compressão moderada e agressiva.
  • No cenário de 33% de redução, o HiAP obteve +0.93% de ganho de acurácia sobre a linha de base de proporção uniforme.

• Eficiência de Hardware:

  • A poda física das dimensões resultou em um aumento real de velocidade de inferência. Em um GPU único, o tempo de latência caiu de 5.57 ms para 3.86 ms (~1.44x de aceleração), confirmando que as economias de MACs se traduzem diretamente em ganhos de throughput.

• Análise de Dinâmica de Busca:

  • O modelo demonstrou comportamento inteligente: nas fases iniciais, priorizou a poda macro (remover cabeças inteiras e blocos FFN redundantes) para reduzir o overhead de memória. Posteriormente, refinou a poda micro (dimensões internas) para ajustar o orçamento de FLOPs, criando uma distribuição de largura heterogênea e eficiente.

5. Significância e Conclusão

O HiAP representa um avanço significativo na otimização de Vision Transformers para dispositivos com recursos limitados. Ao tratar a poda como um problema de aprendizado de orçamento único e diferenciável, o método supera as limitações de granularidade única e dependência de intervenção humana dos métodos anteriores.

A principal contribuição prática é a capacidade de gerar sub-redes nativamente eficientes que reduzem tanto o tráfego de memória (via poda macro) quanto os cálculos (via poda micro), eliminando a complexidade de pipelines de treinamento multi-estágio. Isso torna a implantação de ViTs em hardware de borda mais viável e acessível, oferecendo um novo paradigma para a busca de arquiteturas (NAS) focada em eficiência real.