Sakana AI presenta DiffusionBlocks, un framework di addestramento a blocchi che trasforma le reti residuali in moduli di denoising addestrabili in modo indipendente

Un team di ricerca formato da Sakana AI e l’Università di Tokyo ha proposto un nuovo framework per il training delle reti neurali chiamato DiffusionBlocks. L’idea centrale consiste nel suddividere la rete in blocchi di addestramento autonomi, riducendo drasticamente la memoria necessaria mentre mantiene le performance. Invece di addestrare l’intera rete con le backpropagation tradizionali, ogni blocco può essere addestrato separatamente, rendendo l’approccio più efficiente.

Il problema della memoria durante il training

Benché la backpropagation end-to-end funzioni bene, richiede di salvare le attivazioni intermedie su ogni layer, consumando una quantità di memoria che cresce linearmente con la profondità della rete. Aumentando il numero di strati, quindi, questo diventa un collo di bottiglia.

Una tecnica esistente per ridurre la memoria, detta activation checkpointing, evita di salvare le attivazioni temporanee ricalcolandole solo quando necessario. Ma pur ottenendo un miglioramento parziale, non riduce la memoria necessaria per i parametri, gradienti e stato dell’ottimizzatore.

Ad esempio, con l’ottimizzatore Adam ogni layer deve conservare la memoria per i parametri, gradienti e due stati dell’ottimizzatore (momentum e varianza), equivalente a 4 volte la dimensione dei parametri. L’activation checkpointing non riduce questo aspetto.

Adeguarsi ai blocchi

L’approssimazione che permette a DiffusionBlocks di funzionare si basa sulla suddivisione di una rete in B blocchi e sull’addestramento uno alla volta di questi blocchi riducendo la richiesta di memoria a circa 1/B. La riduzione è proporzionale al numero di blocchi utilizzati.

Tuttavia, il problema principale è definire un oggettivo locale per ciascun blocco che, quando applicato indipendentemente, produca comunque un modello coerente a livello globale. Molti approcci precedenti, come l'algoritmo Forward-Forward di Hinton e il training greedy layer-wise, si sono rivelati sottoperformanti.

DiffusionBlocks: un’idea strutturale

DiffusionBlocks supera questi limiti introducendo un nuovo approccio che si basa sull’idea di connessioni residue e sull’approssimazione di equazioni differenziali ordinarie (ODE). Una connessione residuale, che calcola z_ℓ = z_ℓ–1 + fθ_ℓ(z_ℓ–1), può essere interpretata come un passo di approssimazione di un’ODE.

In particolare, i ricercatori collegano gli aggiornamenti delle layer residuali all’ODE della probability flow presente nei modelli di diffusione score-based. L’equazione di diffusione reversa segue:

d𝐳/σ dσ = -σ ∇𝐳 log p_σ(𝐳_σ)

Applicando una discretizzazione di Euler a questa equazione, si ottiene uno schema strutturalmente simile all’aggiornamento in un network residuale. Stacking multiple blocks, ognuno opera su una finestra temporale specifica di rumorosità.

Convertire la rete: tre passaggi

I ricercatori hanno stabilito il seguente processo per l’adattamento:

• Partizionamento a blocchi: dividi la rete di L strati in B blocchi, ciascun blocco contiene uno strato continuo.
• Assegnazione degli intervalli di rumore: definisci una distribuzione di rumore p_noise e un intervallo di rumorosità [σ_min, σ_max]. Dividi l’intervallo in B intervalli e assegna un intervallo a ogni blocco.
• Condizionamento al rumore: espandi ciascun input di blocco aggiungendo una versione rumorosa del target e usa il AdaLN per aggiungere condizionamento per il livello del rumore.

Partizionamento equi-probabilistico

Invece di distribuire uniformemente gli intervalli di rumorosità, DiffusionBlocks utilizza un partizionamento equi-probabilistico. Ciascun blocco gestisce esattamente 1/B della massa totale della distribuzione p_noise lungo l’intervallo [σ_min, σ_max].

I blocchi che trattano livelli intermedi di rumorosità ricevono intervalli più stretti, mentre quelli estremi ricevono intervalli più larghi. Questo permette di bilanciare meglio i contributi di ogni blocco in base alla difficoltà di denoising.

Una sperimentazione su CIFAR-10 ha mostrato che il partizionamento equi-probabilistico raggiunge un FID di 38.03 contro 43.53 per partizionamento uniforme.

Risultati sperimentali

I ricercatori hanno valutato DiffusionBlocks su cinque architetture diverse, testando su tre tipi di task: visione, generazione di testi e modelli autoregressivi. Tutti i risultati confrontano DiffusionBlocks (addestrati blocco per blocco) con la stessa architettura addestrata con backpropagation full e end-to-end.

Ecco un riepilogo dei risultati:

• ViT, 12 layers, B=3 su CIFAR-100: Accuratezza del baseline 60.25%, DiffusionBlocks 59.30% con riduzione di memoria del 3x.
• DiT-S/2, 12 layers, B=3 su CIFAR-10: FID 39.83 vs 37.20 con riduzione del 3x.
• DiT-L/2, 24 layers, B=3 su ImageNet 256×256: 12.09 vs 10.63 con miglioramento della denoising qualità.
• MDM, 12 layers, B=3 su text8: BPC 1.56 vs 1.45.
• AR Transformer, 12 layers, B=4 su LM1B: MAUVE 0.50 vs 0.71.
• AR Transformer, 12 layers, B=4 su OpenWebText: MAUVE 0.85 vs 0.82.
• Huginn, architettura iterativa su LM1B: MAUVE 0.49 vs 0.70, con riduzione di 10x in calcolo totale.

Ogni risultato conferma che la performance globale si mantiene stabile pur riducendo fino al 4x la memoria o computo necessario durante l’addestramento.