Distilling Vision Transformers for Distortion-Robust Representation Learning

Resumo Técnico: Destilação de Vision Transformers para Aprendizado de Representações Robustas a Distorções

1. O Problema

O aprendizado de representações visuais robustas é um desafio central na visão computacional, especialmente quando os dados estão sujeitos a ruídos, oclusões, desfoque (blur) ou perda de pixels. Em aplicações críticas como imagens médicas, sensoriamento remoto e sistemas autônomos, dados "limpos" (sem distorções) são frequentemente escassos ou inexistentes.

Modelos pré-treinados em dados limpos (como os Vision Transformers - ViT) perdem sua capacidade semântica quando confrontados com entradas corrompidas. O desafio proposto é: é possível recuperar as representações de um codificador limpo observando apenas versões distorcidas das imagens, sem nunca ter acesso direto aos dados originais sem ruído?

2. Metodologia

Os autores propõem um framework de destilação de conhecimento assimétrica baseado em Vision Transformers (ViT). A abordagem trata a robustez como um problema de "recuperação de representação no espaço latente".

• Configuração Professor-Aluno: Ambos os modelos são inicializados com o mesmo checkpoint pré-treinado (DINO-ViT). O Professor é mantido congelado e processa imagens limpas ($x$), enquanto o Aluno é treinado para processar as versões distorcidas ($\tilde{x}$) dessas mesmas imagens.
• Objetivo de Destilação Multi-nível: Para garantir que o aluno capture a semântica completa, o método não se limita apenas à saída final, mas alinha três níveis hierárquicos de representação:
  1. Alinhamento Semântico Global ($L_{cls}$): Utiliza o erro quadrático médio (MSE) para alinhar o class token do professor e do aluno, garantindo que a semântica de alto nível (categoria do objeto) seja preservada.
  2. Alinhamento Espacial Local ($L_{patch}$): Alinha os patch tokens (tokens de patches) via MSE para manter a consistência da estrutura espacial e detalhes locais, impedindo que a distorção colapse a estrutura da imagem.
  3. Alinhamento de Estrutura de Atenção ($L_{attn}$): Utiliza a Divergência KL para alinhar os mapas de atenção de cada cabeça do Transformer. Isso força o aluno a "olhar" para as mesmas regiões de interesse que o professor olharia em uma imagem limpa.

A função de perda total é uma combinação ponderada desses três termos: $\mathcal{L} = \lambda_{cls}\mathcal{L}_{cls} + \lambda_{patch}\mathcal{L}_{patch} + \lambda_{attn}\mathcal{L}_{attn}$.

3. Principais Contribuições

• Framework de Destilação Assimétrica: Um método eficaz para treinar codificadores robustos sem a necessidade de dados limpos durante o treinamento.
• Supervisão Multi-nível: A introdução do alinhamento de mapas de atenção e patches, proporcionando uma supervisão muito mais rica do que métodos que focam apenas em embeddings globais.
• Abordagem Auto-supervisionada: O método não utiliza rótulos de classe durante a fase de destilação, dependendo inteiramente das representações internas do professor.

4. Resultados

Os experimentos foram realizados no ImageNet-100 e em diversos datasets de transferência (CIFAR, STL-10, RESISC45, CAMELYON17) sob três tipos de distorção: mascaramento de pixels, ruído Gaussiano e desfoque Gaussiano.

• Desempenho In-Domain: O método superou consistentemente o ajuste fino supervisionado (supervised baseline) e o método de Inversão Contrastiva (CI) em quase todos os cenários de distorção severa.
• Generalização (Transfer Learning): O modelo demonstrou uma capacidade excepcional de transferir o conhecimento de robustez para domínios completamente diferentes, como imagens aéreas (RESISC45) e histopatologia médica (CAMELYON17), superando significativamente os baselines.
• Eficiência de Rótulos: O modelo mostrou-se altamente eficiente em cenários de poucos dados (low-label regime). Mesmo com apenas 5% a 10% dos rótulos disponíveis, o desempenho aproximou-se do modelo treinado com 100% dos rótulos.
• Análise de Atenção: Mapas de atenção visuais confirmaram que, enquanto modelos supervisionados comuns perdem o foco em imagens distorcidas, o modelo destilado continua a focar nas regiões semânticas relevantes do objeto.

5. Significância

Este trabalho é significativo porque demonstra que a robustez não precisa ser aprendida apenas através de aumentos de dados (data augmentation) ou reconstrução de pixels (denoising), mas pode ser alcançada através do alinhamento de estruturas semânticas latentes. A capacidade de aprender representações robustas e transferíveis sem acesso a dados limpos abre portas para o uso de IA em setores onde a qualidade da imagem é inerentemente baixa ou imprevisível, como medicina e exploração espacial.