Imparare a leggere i token contestuali nei transformer di diffusione

Abstract degli autori

I transformer di diffusione multimodali (MM-DiTs) elaborano congiuntamente rappresentazioni visive e testuali durante l’intero processo di generazione. Questi modelli aggiornano ripetutamente i token testuali attraverso l’attenzione multimodale, formando token contestuali dinamici la cui funzione non è ancora ben compresa. In questo lavoro presentiamo un metodo per leggere questo spazio contestuale ponendo domande in linguaggio naturale. Addestriamo una piccola rete con un collo di bottiglia che mappa i token contestuali intermedi nello spazio di input di un modello linguistico di grandi dimensioni (LLM) con parametri congelati. Questo permette all’LLM di rispondere a domande sull’immagine in formazione direttamente a partire da tali rappresentazioni nascoste. Il nostro sistema di lettura rivela che i token contestuali codificano una rappresentazione ricca e globale della scena in formazione: le informazioni semantiche specifiche della generazione, compresi gli attributi che il prompt lascia non del tutto specificati, sono accessibili sorprendentemente presto nel processo di rimozione del rumore, mentre col tempo diventano leggibili dettagli sempre più fini. È notevole che queste informazioni restino decodificabili anche quando l’MM-DiT riceve un prompt vuoto: ciò mostra che i token contestuali accumulano una quantità considerevole di informazioni specifiche dell’immagine dalla stessa rappresentazione visiva in evoluzione. Osserviamo inoltre che le immagini generate con rappresentazioni contestuali più leggibili tendono a ottenere punteggi più alti in termini di preferenze umane. Sulla base di queste osservazioni presentiamo Contextual Alignment, una tecnica di addestramento che rafforza esplicitamente le informazioni visivo-semantiche codificate nei token contestuali, migliorando la qualità della generazione e la copertura della distribuzione. Nel complesso, i nostri risultati mostrano che i token contestuali offrono sia una prospettiva interpretabile sulle dinamiche interne degli MM-DiTs sia un obiettivo efficace per migliorare i modelli generativi.