Decoding Looped Transformers Better for (Almost) Free

L'article présente LoopCD, un cadre de décodage contrastif sans entraînement qui exploite les états intermédiaires des boucles récurrentes dans les Transformers bouclés. En contrastant la prédiction finale avec une passe récurrente antérieure (dans l'espace des logits ou des états cachés), LoopCD guide efficacement la sélection des jetons. Cette approche offre des gains substantiels de performance de décodage tout en permettant de réduire de moitié le nombre de boucles récurrentes.

LoopCD-Logits fait passer le pass@1 de AIME 2024 pour Ouro-2.6B-Thinking de 61,88 % à 73,33 %.

LoopCD-Hidden élève le pass@1 de HumanEval pour Huginn de 22,56 % à 31,71 %.

L'approche permet de diviser par deux le nombre de boucles récurrentes, réduisant les FLOPs directs de 22,5 % à 48,2 %.