Distilling Vision Transformers for Distortion-Robust Representation Learning

Technische Samenvatting: Distilling Vision Transformers for Distortion-Robust Representation Learning

1. Het Probleem (Problem Statement)

Het trainen van robuuste visuele representaties is een grote uitdaging wanneer beelden onderhevig zijn aan vervormingen (distorties) zoals ruis, onscherpte (blur), occlusie of ontbrekende pixels. In kritieke domeinen zoals medische beeldvorming, remote sensing en autonome systemen is hoogwaardige, "schone" data vaak schaars of onbeschikbaar.

Huidige methoden hebben twee grote tekortkomingen:

• Supervised learning: Modellen die getraind zijn op schone, gelabelde data presteren drastisch slechter wanneer ze worden geconfronteerd met gecorrumpeerde inputs.
• Image Restoration: Methoden die proberen de pixels te herstellen (denoising), richten zich op pixel-niveau in plaats van op het semantische niveau dat nodig is voor downstream taken zoals classificatie.

2. Methodologie (Methodology)

De auteurs stellen een asymmetrisch kennisdistillatie-framework (Knowledge Distillation) voor. In plaats van pixels te reconstrueren, proberen ze de "schone" semantische representaties te herstellen in de latente ruimte.

• Teacher-Student Setup: Er wordt gebruikgemaakt van twee Vision Transformers (ViT) die beide zijn geïnitialiseerd met een pretrained model (DINO-ViT).
  • De Teacher is bevroren (frozen) en verwerkt uitsluitend schone afbeeldingen ($x$).
  • De Student wordt getraind en verwerkt alleen de vervormde versies ($\tilde{x}$) van diezelfde afbeeldingen.
• Multi-Level Distillation Objective: Om de student effectief te laten leren, wordt er niet alleen gekeken naar de output, maar wordt de kennis op drie hiërarchische niveaus uitgelijnd:
  1. Global Semantic Alignment ($L_{cls}$): Minimaliseert het verschil tussen de class tokens van de teacher en student (MSE-loss) om globale semantiek te behouden.
  2. Local Spatial Alignment ($L_{patch}$): Minimaliseert het verschil tussen de patch tokens om de lokale ruimtelijke structuur te behouden.
  3. Attention Structure Alignment ($L_{attn}$): Gebruikt KL-divergentie om de attention maps van de teacher en student gelijk te trekken. Dit zorgt ervoor dat de student leert naar dezelfde relevante objectgebieden te kijken als de teacher.

De totale loss is een gewogen som van deze drie componenten: $\mathcal{L} = \lambda_{cls}\mathcal{L}_{cls} + \lambda_{patch}\mathcal{L}_{patch} + \lambda_{attn}\mathcal{L}_{attn}$.

3. Belangrijkste Bijdragen (Key Contributions)

• Nieuw Framework: Een zelfgesuperviseerd distillatie-raamwerk dat specifiek is ontworpen voor het leren van robuuste representaties zonder ooit directe toegang tot schone trainingsdata voor de student.
• Multi-target supervisie: De introductie van een hiërarchische aanpak (global, local, en attention) die rijker is dan standaard distillatiemethoden.
• Robuustheid en Transferabiliteit: Het bewijs dat een model dat getraind is op één type vervorming (bijv. ruis) effectief kan worden overgedragen naar andere domeinen (medisch, satellietbeelden) en andere vervormingen.

4. Resultaten (Results)

De resultaten op ImageNet-100 en diverse transfer-datasets (CIFAR, STL-10, RESISC45, CAMELYON17) laten het volgende zien:

• Outperformance: De voorgestelde methode presteert consistent beter dan zowel de supervised baseline als de bestaande Contrastive Inversion (CI) methode.
• Extreme Distorties: Het model blijft effectief, zelfs bij extreme gevallen zoals 90% pixel-masking of zeer sterke Gaussian blur.
• Label-efficiëntie: De methode is zeer efficiënt met labels. Zelfs met slechts 5% tot 10% van de beschikbare labels benadert de student de prestaties van een volledig gesuperviseerd model.
• Interpretabiliteit: De attention maps (Fig. 4 in het paper) laten zien dat de student, ondanks de zware vervormingen, nog steeds correct focust op de semantisch relevante objecten, in tegenstelling tot de supervised baseline die de focus verliest.

5. Betekenis (Significance)

Dit onderzoek is significant omdat het aantoont dat we de "kennis" van krachtige pretrained modellen kunnen gebruiken om de kloof tussen vervormde observaties en schone semantiek te overbruggen. Het biedt een oplossing voor scenario's waarin data van lage kwaliteit de standaard is, waardoor AI-systemen betrouwbaarder worden in de echte wereld (zoals in de medische sector of bij autonome voertuigen) zonder dat er enorme hoeveelheden handmatig gelabelde, schone data nodig zijn.