Policy-Mask: adattamento delle feature dei modelli di diffusione ai compiti discriminativi in funzione del budget

Abstract degli autori

Selezionare le feature intermedie di un modello di diffusione i cui parametri restano fissi per svolgere compiti discriminativi non è semplice: il segnale utile è distribuito lungo la traiettoria di denoising, tra i diversi livelli della rete e tra i canali. In pratica, occorre scegliere un timestep, un blocco e un sottoinsieme di canali. I metodi esistenti fissano queste scelte manualmente oppure ottimizzano solo alcuni iperparametri di estrazione, il più delle volte il timestep, per poi applicare la stessa regola a input e budget diversi. Presentiamo Policy-Mask, una strategia condizionata dal budget e addestrata sull’obiettivo del compito finale, che seleziona congiuntamente tutti e tre gli elementi nel rispetto di un budget di canali prestabilito. Senza normalizzazione, il contributo dei punteggi dei canali cresce con il numero di canali selezionati e prevale sugli aggiornamenti relativi al timestep e al blocco. Lo normalizziamo in base alla cardinalità del sottoinsieme, mantenendo le tre teste della strategia su scale comparabili e stabilizzando l’addestramento congiunto. A parità di budget, Policy-Mask supera i metodi di riferimento basati sull’estrazione fissa e sul filtraggio delle feature in compiti di segmentazione, classificazione e corrispondenza semantica: raggiunge 38,9 mIoU su ADE20K, il 75,4% di accuratezza top-1 su CIFAR-100 con 128 dimensioni e 76,3 PCK@0,1 su SPair-71k. Le analisi dei singoli componenti mostrano che la scelta del blocco incide sensibilmente sull’accuratezza nei compiti finali e che il punteggio più alto registrato per un budget di canali si ottiene mantenendo 256 canali, individuando così una configurazione pratica per rappresentazioni compatte dei modelli di diffusione. Adattiamo inoltre Policy-Mask ai diffusion transformer, tra cui PixArt-$Σ$ e SD3.5 MMDiT, sostituendo i canali delle U-Net con le dimensioni degli stati nascosti dei token immagine.