Mixture of Depths (MoD)

Mixture of Depths (MoD) è una tecnica per Transformer in cui un router appreso decide, blocco per blocco, quali token elaborare e quali saltare il blocco tramite la connessione residua, così token diversi ricevono quantità diverse di calcolo. David Raposo, Adam Santoro e i loro colleghi di Google DeepMind l'hanno introdotta nell'aprile 2024 nell'articolo Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models. Una capacità fissa k stabilisce quanti token vengono elaborati da ogni blocco instradato; il budget di calcolo e le forme dei tensori sono quindi noti in anticipo, anche se la scelta dei token dipende dall'input. I modelli linguistici di grandi dimensioni (LLM) MoD hanno eguagliato la qualità di base a parità di FLOPs di addestramento e tempo effettivo, usando una frazione dei FLOPs per passaggio in avanti e completando i passaggi di generazione con una velocità superiore del 50%. A budget di addestramento fisso, il miglior modello MoD era più grande del miglior modello di base e raggiungeva una loss inferiore.

Come funziona l'instradamento di MoD#

Prima di un blocco instradato, un router lineare assegna a ogni token un punteggio scalare. I k token con il punteggio più alto attraversano l'attenzione automatica e la MLP del blocco; i relativi output vengono moltiplicati per i punteggi del router e aggiunti di nuovo al flusso residuo. La moltiplicazione per il punteggio inserisce il router nel percorso del gradiente, così la retropropagazione può insegnargli quali token traggono vantaggio dall'elaborazione. I token rimanenti passano invariati e vengono esclusi dall'attenzione di quel blocco, quindi i token successivi non possono prestarvi attenzione in quel punto e il costo dell'attenzione diminuisce con il quadrato della capacità: dimezzare i token costa un quarto dei prodotti query-key. Nell'articolo, l'impostazione migliore instradava un blocco su due con una capacità pari al 12,5% della sequenza. Una versione minima di un singolo blocco instradato, qui con la sola MLP, in PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Stampa Block ran on 16 of 128 tokens; 112 skipped it unchanged: la MLP del blocco ha eseguito un ottavo del lavoro abituale e il tensore di output mantiene la forma dell'intera sequenza per il blocco successivo.

Instradamento durante la generazione#

La selezione top-k non è causale: il fatto che un token rientri tra i k migliori della sequenza dipende dai punteggi dei token che lo seguono, i quali non esistono ancora quando un modello genera testo un token alla volta. Gli autori hanno reso causale l'instradamento in due modi: una loss ausiliaria di entropia incrociata binaria insegna all'output del router a prevedere l'appartenenza ai top-k, modificando la loss di modellazione linguistica solo dello 0,2-0,3%; in alternativa, un piccolo predittore MLP separato, con input stop-gradient, ha raggiunto rapidamente un'accuratezza del 97-99%. In entrambi i casi, il campionamento autoregressivo perdeva poca qualità. I router addestrati mostrano inoltre uno schema: alcuni token attivano ogni blocco instradato lungo la profondità del modello, mentre altri aggirano i blocchi ogni volta che possono.

MoD e tecniche correlate#

• Miscela di esperti (MoE): Instrada ogni token verso una o più MLP esperte all'interno di un livello, aumentando il numero di parametri ma mantenendo fisso il calcolo per token. L'articolo ha combinato anche le due tecniche in MoDE, in cui un token può essere instradato a un esperto oppure a un'operazione nulla; i vantaggi si sono sommati.
• Uscita anticipata: Un token si arresta dopo un certo livello e salta tutti i successivi. MoD può saltare un blocco e usare comunque quelli successivi.
• Unione dei token (ToMe): Unisce token simili in un Transformer visivo così che i livelli successivi ne elaborino meno, senza un router appreso.
• Potatura del modello: Rimuove definitivamente i pesi, così ogni input usa la stessa rete ridotta.
• Instradamento del modello: Invia un'intera richiesta a uno tra diversi modelli, applicando la stessa idea al livello dei modelli completi.

La maggior parte delle reti fissa la propria profondità al momento della creazione. Ultralytics YOLO26 la imposta in base alla scala del modello: il moltiplicatore di profondità nel suo file YAML del modello è 0.50 per le scale n, s e m e 1.00 per l e x; uno stadio C3k2 indicato con 2 ripetizioni esegue quindi 1 blocco in YOLO26n e 2 in YOLO26x, e ogni immagine attraversa tutti i blocchi.