Sintesi Tecnica: Distillazione della Rappresentazione Discriminativa e Consistente (DCD)

1. Definizione del Problema

La Distillazione della Conoscenza (KD) mira a trasferire la conoscenza da modelli teacher ad alta capacità a modelli student compatti. Sebbene gli obiettivi contrastivi si siano dimostrati efficaci per l'apprendimento di rappresentazioni strutturate nell'apprendimento auto-supervisionato, la loro applicazione alla distillazione affronta tre principali criticità pratiche:

1. Overhead di Memoria: I metodi esistenti di distillazione contrastiva (ad es., CRD) si affidano a banchi di memoria esterni per memorizzare i vettori delle feature per il campionamento dei negativi, introducendo costi di memoria significativi e complessità di implementazione.
2. Iperparametri Fissi: L'uso di iperparametri di temperatura fissi limita l'adattabilità dell'obiettivo contrastivo attraverso i diversi stadi di addestramento e i gap di capacità tra teacher e student.
3. Supervisione Strutturale Incompleta: Gli obiettivi di discriminazione dell'istanza tipicamente supervisionano solo la diagonale della matrice di similarità tra i modelli (allineando un'istanza dello student con il suo corrispettivo nel teacher). Ciò lascia la struttura fuori diagonale non vincolata, permettendo allo student di corrispondere al proprio teacher pur relazionandosi in modo incoerente con altri campioni nel batch.

2. Metodologia

Gli autori propongono la Distillazione della Rappresentazione Discriminativa e Consistente (DCD), un framework che combina la discriminazione dell'istanza contrastiva con un termine di regolarizzazione della consistenza sulla matrice di similarità tra i modelli.

2.1 Obiettivo Core

La funzione di perdita totale per l'addestramento della rete student $\theta_S$ è definita come:
$$\theta_S = \arg\min_{\theta_S} \sum_{i=1}^N \left( \mathcal{L}_{\text{sup}} + \lambda \cdot \mathcal{L}_{\text{distill}} + \beta \cdot \mathcal{L}_{\text{dcd}} \right)$$
dove $\mathcal{L}_{\text{dcd}}$ è la perdita nello spazio delle feature proposta, composta da due termini:
$$\mathcal{L}_{\text{dcd}} = \mathcal{L}_{\text{ctr}} + \alpha \cdot \mathcal{L}_{\text{cst}}$$

2.2 Termine Contrastivo ($\mathcal{L}_{\text{ctr}}$)

Questo termine esegue la discriminazione a livello di istanza. Per un mini-batch di $N$ campioni, la rappresentazione dello student $z_i^S$ deve identificare la sua corrispondente rappresentazione del teacher $z_i^T$ tra tutte le $N$ rappresentazioni del teacher nel batch. Questo è formulato come un problema di classificazione a $N$-vie utilizzando la Stima del Contrasto del Rumore (NCE):
$$\mathcal{L}_{\text{ctr}} = -\frac{1}{N} \sum_{i=1}^N \log \frac{\exp(\ell_{ii})}{\sum_{j=1}^N \exp(\ell_{ij})}$$
Qui, $\ell_{ij}$ rappresenta la similarità scalata tra lo student $i$ e il teacher $j$. Fondamentalmente, questo termine utilizza il campionamento in-batch, eliminando la necessità di banchi di memoria esterni.

2.3 Termine di Consistenza ($\mathcal{L}_{\text{cst}}$)

Il termine contrastivo vincola solo la diagonale della matrice di similarità. Per vincolare la struttura fuori diagonale, gli autori introducono un termine di regolarizzazione della consistenza che impone l'accordo tra due viste complementari della stessa matrice di similarità tra modelli $\ell$:

1. Vista dello Student ($p_i^S$): Probabilità normalizzate per riga, che rappresentano l'affinità di uno student $i$ verso tutti i teacher.
2. Vista del Teacher ($p_i^T$): Probabilità normalizzate per colonna (trasposte), che rappresentano l'affinità di tutti gli student verso il teacher $i$.

La perdita di consistenza minimizza la divergenza di Kullback-Leibler (KL) tra queste due viste:
$$\mathcal{L}_{\text{cst}} = \frac{1}{N} \sum_{i=1}^N D_{\text{KL}}(\mathbf{p}_i^T \,\|\, \mathbf{p}_i^S)$$
Gli autori dimostrano che questo termine svanisce se e solo se la matrice di similarità tra i modelli è simmetrica ($\langle z_i^S, z_j^T \rangle = \langle z_j^S, z_i^T \rangle$). Ciò assicura che l'affinità dello student $i$ verso il teacher $j$ concordi con l'affinità dello student $j$ verso il teacher $i$, catturando una consistenza strutturale bidirezionale.

2.3 Scala e Bias Apprendibili

A differenza dei metodi che utilizzano temperature fisse, DCD parametrizza la scalatura della similarità come $s = \exp(\tau)$ e aggiunge un bias apprendibile $b$. Questi parametri ($\tau$ e $b$) sono ottimizzati durante l'addestramento per controllare in modo adattivo la nitidezza e l'offset del segnale di distillazione, rimuovendo la necessità di una regolazione manuale per ogni coppia teacher-student.

2.4 Dettagli di Implementazione

• Campionamento In-Batch: I campioni negativi sono estratti naturalmente dal corrente mini-batch, riducendo l'overhead di memoria da ~655MB (per ImageNet in CRD) a ~0.13MB per iterazione.
• Teste di Proiezione: Le feature dello student e del teacher sono mappate in uno spazio di embedding condiviso di 128 dimensioni tramite strati di proiezione lineare. Questi strati sono addestrati congiuntamente ma scartati all'inferenza, aggiungendo parametri trascurabili (66K) e nessun costo di inferenza.
• Velocità di Addestramento: Il metodo eguaglia la velocità di addestramento della standard logit-based KD (8ms per batch).

3. Contributi Chiave

1. Framework Unificato: Un framework di distillazione che combina la discriminazione dell'istanza contrastiva con una regolarizzazione della consistenza basata sulla divergenza KL che vincola la struttura fuori diagonale della matrice di similarità tra i modelli.
2. Efficienza: Una strategia di campionamento in-batch efficiente che elimina i banchi di memoria esterni, riducendo l'overhead di memoria e il tempo di addestramento mantenendo un throughput comparabile alla standard KD.
3. Parametri Adattivi: L'introduzione di parametri di scala e bias apprendibili che regolano automaticamente il segnale di distillazione, evitando la sintonizzazione manuale degli iperparametri.
4. Validazione Estesa: Esperimenti completi su CIFAR-100, ImageNet e MS-COCO, che dimostrano prestazioni competitive in classificazione, rilevamento di oggetti e trasferimento tra dataset.

4. Risultati Sperimentali

Il metodo è stato valutato su cinque benchmark (CIFAR-100, ImageNet, STL-10, Tiny ImageNet, MS-COCO) su 13 coppie teacher-student.

• Classificazione (CIFAR-100): DCD combinato con KD (DCD + KD) raggiunge prestazioni competitive, superando la rete teacher di +0.45% in configurazioni con stessa architettura e di +0.90% in configurazioni con architetture diverse. Sebbene metodi come ReviewKD, SimKD e CRCD ottengano un'accuratezza superiore in configurazioni specifiche (ad es., ReviewKD raggiunge il 76.12% rispetto al 76.06% di DCD+KD nella coppia WRN-40-2 $\to$ WRN-16-2), DCD+KD supera costantemente la standard KD e molti altri baseline in diverse combinazioni architettoniche.
• Classificazione (ImageNet): Il metodo ha migliorato costantemente le prestazioni rispetto ai baseline (KD, AT, CRD, ecc.) attraverso diverse architetture, confermando che i guadagni si estendono a configurazioni su larga scala.
• Rilevamento Oggetti (MS-COCO): Utilizzando Faster R-CNN con FPN, DCD ha migliorato la Mean Average Precision (AP) in tutte le coppie teacher-student testate (ad es., ResNet-101 $\to$ ResNet-18), dimostrando l'applicabilità oltre la classificazione.
• Trasferibilità: Le rappresentazioni distillate hanno mostrato una superiore trasferibilità cross-dataset su STL-10 e Tiny ImageNet rispetto ai baseline, indicando una migliore generalizzazione.
• Efficienza: DCD ha aggiunto solo 66K parametri (tutti scartati all'inferenza) e ha eguagliato la latenza di addestramento di 8ms/batch della standard KD, rendendolo significativamente più veloce e leggero rispetto a CRD (19ms/batch) e altri metodi basati sulle feature.

5. Significato e Rivendicazioni

Il paper sostiene che DCD affronti le specifiche limitazioni degli attuali metodi di distillazione contrastiva:

• Rimozione dei Banchi di Memoria: Operando interamente all'interno del mini-batch per ridurre il consumo di memoria e la complessità.
• Imposizione della Consistenza Strutturale: Andando oltre il semplice allineamento dell'istanza per imporre la simmetria nella matrice di similarità tra i modelli, catturando così una conoscenza strutturale più ricca.
• Adattabilità: Utilizzando parametri apprendibili per adattarsi a diversi stadi di addestramento e capacità dei modelli senza la sintonizzazione manuale degli iperparametri.

Gli autori posizionano DCD come un metodo che raggiunge l'efficienza della semplice distillazione basata sui logit pur mantenendo i benefici di accuratezza della distillazione delle feature contrastive, senza richiedere cambiamenti architettonici o infrastrutture di memoria esterne.