TaH2 verbessert Test-Time-Skalierung von Loop-Transformern durch adaptives Looping

TaH2: Adaptives Looping für effizientere Test-Time-Skalierung

Loop-Transformer nutzen wiederverwendete Schichten für latente Berechnungen und zeigen dadurch eine vielversprechende Parametereffizienz. Bisher vergleichen Studien solche Modelle mit Nicht-Loop-Varianten meist bei gleichen Parametern oder FLOPs pro Token. Ob das Looping die Test-Time-Skalierung bei längeren Ausgaben tatsächlich verbessert, war bislang kaum untersucht. Der neue Ansatz TaH2 setzt hier an: Statt wie frühere Loop-Transformer mit fester Iterationstiefe zusätzliche Iterationen auf jedes Token anzuwenden, führt TaH2 einen adaptiven Iterationsentscheider ein. Dieser wird durch Lookahead-Tiefenüberwachung trainiert und lenkt zusätzliche Iterationen gezielt auf diejenigen Tokens, die vom Looping profitieren. Nach Angaben der Autoren zeigt die Analyse, dass bei festem Looping viele Tokens keinen Nutzen aus weiteren Iterationen ziehen. TaH2 verbessert dadurch die Genauigkeit pro Rechenaufwand gegenüber früheren Loop-Transformern mit fester Iterationstiefe.

Der adaptive Mechanismus funktioniert, indem der Iterationsentscheider für jedes Token vorhersagt, ob eine weitere Iteration die Vorhersagequalität verbessert. Als Lernsignal dienen Online-Labels, die während des Trainings erzeugt werden: Das Modell führt probeweise eine zusätzliche Iteration aus und vergleicht den Vorhersage-Score vor und nach dieser Iteration. Verbessert sich der Score, wird das Token als profitabel markiert; andernfalls als nicht profitabel. Der Entscheider lernt, diese Profitabilität anhand der Token-Repräsentationen zu erkennen, und der Backbone wird gleichzeitig so angepasst, dass die zusätzlichen Iterationen tatsächlich zu besseren Vorhersagen führen. Dadurch konzentriert TaH2 den Rechenaufwand auf Tokens, bei denen ein Genauigkeitszuwachs zu erwarten ist, und vermeidet verschwendete Iterationen auf Tokens, die bereits konvergiert sind. Die Accuracy-Compute-Steigung misst den Genauigkeitsgewinn pro Verdopplung der Test-Time-Decoding-FLOPs; TaH2 zielt darauf, diese Steigung zu erhöhen, ohne die absolute Genauigkeit bei gleichem Rechenaufwand zu opfern.

53% bessere Accuracy-Compute-Steigung auf AIME-Benchmarks

Auf den anspruchsvollen AIME-Benchmarks verbessert TaH2 die Accuracy-Compute-Steigung um 53 Prozent: von 1,79 auf 2,74 gegenüber der Nicht-Loop-Baseline. Bei gleichem Test-Time-Compute übertrifft TaH2 die Spitzengenauigkeit der Baseline um etwa 3,4 Punkte. Die Autoren hatten zuvor festgestellt, dass bestehende Loop-Transformer oft steilere Steigungen aufweisen als ihre Nicht-Loop-Baseline, bei gleichem Rechenaufwand jedoch schlechter abschneiden. TaH2 kehrt dieses Verhältnis um, indem es zusätzliche Iterationen nicht gleichmäßig verteilt, sondern auf profitable Tokens konzentriert.

Der Grund für diese Verbesserung liegt in der gezielten Iterationsverteilung: Fixed-Depth-Looping wendet zusätzliche Iterationen auf jedes Token an, obwohl viele Tokens – etwa bereits eindeutig bestimmte oder einfache Tokens – keinen Nutzen aus weiteren Durchläufen ziehen. TaH2 identifiziert durch den trainierten Entscheider genau die Tokens, deren Vorhersage sich durch eine weitere Iteration noch verbessern lässt, und investiert den zusätzlichen Rechenaufwand nur dort. Dadurch wird bei gleichem Test-Time-Compute eine höhere Genauigkeit erreicht, und mit steigendem Rechenbudget wächst die Genauigkeit schneller, weil die zusätzlichen Iterationen effizient auf die kritischen Tokens gelenkt werden. Die Accuracy-Compute-Steigung ist definiert als Genauigkeitsgewinn pro Verdopplung der Test-Time-Decoding-FLOPs. Damit zeigt der Ansatz, dass adaptives Looping sowohl die Effizienz als auch die erreichbare Genauigkeit der Test-Time-Skalierung steigert. Die Ergebnisse stammen aus Experimenten mit post-trainierten Loop-Transformern.

Training mit Lookahead-Tiefenüberwachung

Der Iterationsentscheider wird gemeinsam mit dem Backbone nachtrainiert. Dafür nutzt TaH2 eine Lookahead-Tiefenüberwachung: Online-Labels zeigen an, ob eine weitere Iteration die Vorhersage für ein Token verbessert. Diese Labels werden erzeugt, indem das Modell für ein Token eine zusätzliche Iteration durchläuft und den Vorhersage-Score (z. B. die Wahrscheinlichkeit der korrekten Klasse) vor und nach dieser Iteration vergleicht. Ist der Score nach der zusätzlichen Iteration höher, wird das Label positiv gesetzt; andernfalls negativ. Der Entscheider wird darauf trainiert, dieses Label aus der aktuellen Token-Repräsentation vorherzusagen, während der Backbone gleichzeitig so optimiert wird, dass die zusätzliche Iteration tatsächlich zu einer Verbesserung führt. Dieses gemeinsame Post-Training integriert das Lernsignal direkt in beide Komponenten und ermöglicht die adaptive Verteilung von Iterationen. Dadurch wird Rechenzeit nur dort investiert, wo sie die Vorhersage tatsächlich verbessert.