Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

Título: Ensembling Pruned Attention Heads for Uncertainty-Aware Efficient Transformers

Autores: Firas Gabetni, Giuseppe Curci, Andrea Pilzer, Subhankar Roy, Elisa Ricci, Gianni Franchi.

1. O Problema

A Quantificação de Incerteza (UQ) é crucial para a implantação segura de Redes Neurais Profundas (DNNs) em domínios críticos como saúde e direção autônoma. Atualmente, o método mais confiável para UQ são os Deep Ensembles (Conjuntos Profundos), que agregam previsões de múltiplos modelos treinados independentemente.

No entanto, os Deep Ensembles apresentam desafios significativos:

• Custo Computacional e de Memória: Exigem múltiplas rodadas de pré-treinamento e ajuste fino (fine-tuning), armazenamento de vários checkpoints e múltiplas passagens de inferência (uma por modelo), tornando-os inviáveis para modelos de base (foundation models) grandes como CLIP e BERT.
• Limitações de Métodos Alternativos: Abordagens mais leves (como MC Dropout, MIMO, BatchEnsemble) muitas vezes sacrificam a diversidade dos membros do conjunto ou ainda exigem re-treinamento completo.
• O Paradoxo da Poda (Pruning): Uma abordagem ingênua para reduzir custos seria podar pesos de modelos existentes. Contudo, o artigo demonstra que métodos de poda convencionais, embora preservem a precisão em dados limpos, podem degradar severamente a calibração da incerteza e a robustez em dados ruidosos ou fora de distribuição (OOD).

2. Metodologia: Hydra Ensembles

Os autores propõem o Hydra Ensembles, um framework eficiente que constrói um conjunto de modelos transformadores a partir de um único modelo pré-treinado, sem necessidade de re-treinamento do zero.

Conceito Central

A ideia baseia-se na mitologia da Hidra: ao cortar uma cabeça, novas crescem. No contexto de redes neurais, o método gera diversidade criando sub-redes distintas através da poda estruturada de cabeças de atenção (attention heads) de um único modelo base.

Passos da Metodologia:

1. Geração de Membros Diversos:
   • A partir de um único transformador pré-treinado, criam-se múltiplas cópias.
   • Em cada cópia, um conjunto diferente de cabeças de atenção é removido (podado).
   • Estratégias de Poda:
     • Taylor: Poda baseada na importância das cabeças calculada via expansão de Taylor de primeira ordem (sem dados de validação de incerteza).
     • Circuit (Headmap): Poda baseada na extração de "circuitos" neurais que são especificamente importantes para a tarefa de detecção de OOD ou precisão (requer um conjunto de validação focado em incerteza).
2. Fusão Eficiente (Merging):
   • Em vez de executar os modelos sequencialmente, os cabeças restantes de todos os membros são fundidos em um único modelo.
   • Atenção Múltipla Fundida (Fused MHA): Utiliza camadas Fully-Connected Agrupadas (Grouped Fully-Connected - GFC). As projeções de Query, Key e Value de todos os membros são concatenadas e processadas em uma única operação de atenção, mantendo a diversidade estrutural.
   • MLP Fundido: As camadas MLP (Feed-Forward) dos membros são fundidas através da média dos pesos e vieses, reduzindo o custo computacional.
3. Vantagens Operacionais:
   • Inferência: O modelo fundido executa uma única passagem forward, com custo de latência quase idêntico ao de um único modelo (especialmente em precisão bfloat16).
   • Treinamento: Não requer pré-treinamento do zero. Pode ser usado em modo zero-shot ou com um ajuste fino leve apenas nos membros podados.

3. Contribuições Principais

1. Análise Teórica e Empírica da Poda: Os autores provam teoricamente e demonstram empiricamente que a poda ingênua em dados ruidosos degrada a performance de UQ mais do que em dados limpos, pois amplifica a lacuna de perda entre dados limpos e corrompidos.
2. Hydra Ensembles: Introdução do primeiro framework de poda específico para UQ em modelos transformadores de grande escala. Ele preserva a diversidade dos membros (essencial para UQ) enquanto mantém o custo de inferência próximo ao de um único modelo.
3. Desempenho Superior sem Retreinamento: O método alcança estimativas de incerteza comparáveis ou superiores aos Deep Ensembles, mas com uma fração do custo computacional e de memória.

4. Resultados Experimentais

O método foi avaliado em três tarefas principais: classificação supervisionada de imagens, classificação de texto e classificação zero-shot.

• Classificação Supervisionada (ViT-B/16 em ImageNet-1K e CIFAR-100):

  • O Hydra Ensembles superou a maioria das baselines eficientes (MIMO, BatchEnsemble, LoRA) e foi competitivo com os Deep Ensembles em precisão.
  • UQ e OOD: Mostrou ganhos consistentes em métricas de detecção de OOD (AUROC, FPR95, AUPR) em comparação com os Deep Ensembles, especialmente no CIFAR-100 (melhoria de +3.4% em AUROC).
  • Custo: Em precisão bfloat16, o custo de inferência do Hydra Ensembles foi apenas 1.07x o de um modelo único, enquanto os Deep Ensembles foram quase 3x mais lentos.

• Classificação de Texto (BERT em SST-2):

  • O Hydra Ensembles igualou a precisão e calibração dos Deep Ensembles, mas superou-os significativamente na detecção de OOD (melhoria de +2.8% em AUROC e redução de -7.6% em FPR95).

• Classificação Zero-Shot (OpenCLIP-ViT em ImageNet-1K):

  • Destaque: O método superou o estado da arte (ViLU, que requer treinamento) em tarefas de detecção de OOD, mesmo sem treinamento adicional.
  • Ganhos: +1.3% em AUROC, -3.5% em FPR95 e +4.0% em AUPR em relação ao ViLU.
  • O modelo manteve alta precisão e calibração (ECE) sem necessidade de ajuste fino.

5. Significado e Conclusão

O trabalho Hydra Ensembles representa um avanço significativo na interseção entre eficiência computacional e confiabilidade de modelos de IA.

• Escalabilidade: Permite aplicar a robustez dos ensembles a modelos de base massivos (Foundation Models) que seriam proibitivamente caros para ensembles tradicionais.
• Eficiência de Recursos: Reduz drasticamente o uso de memória e tempo de inferência, tornando a quantificação de incerteza viável em cenários de produção com restrições de hardware.
• Mecanismo de Diversidade: Demonstra que a diversidade necessária para uma boa UQ pode ser extraída da estrutura interna do modelo (cabeças de atenção especializadas) em vez de depender apenas da aleatoriedade de inicialização e treinamento completo.

Em resumo, o artigo oferece uma solução prática para o dilema "precisão vs. custo vs. incerteza", permitindo que modelos grandes sejam não apenas precisos, mas também conscientes de suas próprias limitações, sem o ônus computacional tradicional dos ensembles.