Decodificare meglio i Transformer a cicli (quasi) gratis

Abstract degli autori

I Transformer a cicli usano i parametri in modo efficiente eseguendo ripetutamente un blocco condiviso attraverso cicli ricorrenti. Ogni ciclo produce una rappresentazione intermedia da cui è possibile decodificare lo stesso token successivo, ma la decodifica standard scarta gli stati precedenti. Poiché i cicli precedenti comportano meno calcolo, la ricorrenza fornisce di per sé coppie allineate di predizioni deboli e forti, senza modelli ausiliari né addestramento esterno. Presentiamo LoopCD, un metodo di decodifica contrastiva che non richiede addestramento e guida la scelta del token confrontando la predizione finale con quella di un passaggio ricorrente precedente. Può operare nello spazio dei logit, con un passaggio aggiuntivo di output (LoopCD-Logits), oppure nello spazio degli stati nascosti, senza alcun costo aggiuntivo in output (LoopCD-Hidden). In quattro famiglie di Transformer a cicli, LoopCD produce miglioramenti notevoli e costanti a piena profondità ricorrente: LoopCD-Logits porta il pass@1 di Ouro-2.6B-Thinking su AIME 2024 dal 61,88% al 73,33%, mentre LoopCD-Hidden porta il pass@1 di Huginn su HumanEval dal 22,56% al 31,71%. Soprattutto, questi miglioramenti consentono di dimezzare il numero di cicli ricorrenti, ottenendo comunque risultati pari o superiori a quelli delle baseline senza guida a piena profondità e riducendo i FLOPs del passaggio in avanti dal 22,5% al 48,2%. Trasformando gli stati ricorrenti intermedi in segnali di guida efficaci, LoopCD migliora la qualità della decodifica e riduce sensibilmente il calcolo necessario per l’inferenza.