🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

Das Papier stellt Token-Selective Attention (TSA) vor, einen leichten, end-to-end differenzierbaren Gating-Mechanismus, der Transformern ermöglicht, redundante Schichtberechnungen für einzelne Token basierend auf kontextueller Schwierigkeit dynamisch zu überspringen, wodurch erhebliche Effizienzgewinne bei minimalem Qualitätsverlust und ohne Notwendigkeit expliziter Tiefenregularisierung erzielt werden.

Ursprüngliche Autoren: Ahmed Abdelmuniem Abdalla Mohammed

Veröffentlicht 2026-05-08

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Ahmed Abdelmuniem Abdalla Mohammed

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie betreiben eine Fabrik, in der jedes einzelne Produkt, das vom Fließband kommt, exakt die gleiche Behandlung erfährt. Ob es sich um ein einfaches, perfektes Bauteil oder ein komplexes, defektes handelt, jedes Item verbringt an jeder einzelnen Station genau die gleiche Zeit. So funktionieren derzeit Standard-KI-Modelle (Transformer): Sie verarbeiten jedes Wort in einem Satz durch exakt die gleiche Anzahl von „Denk-Schichten", unabhängig davon, wie einfach oder schwer dieses Wort zu verstehen ist.

Diese Arbeit stellt ein neues System namens Token-selektive Aufmerksamkeit (Token-Selective Attention, TSA) vor. Stellen Sie sich TSA als die Installation eines intelligenten, automatischen Wächters an jeder einzelnen Station der Fabrik vor.

Das Problem: Die „Einheitsgröße"-Fabrik

In einem Standard-KI-Modell behandelt das Modell, wenn Sie es bitten, „Sein oder Nichtsein" zu schreiben, das Wort „Sein" und das Wort „oder" mit der gleichen Menge an Rechenkraft.

Einfache Wörter (wie „der" oder „ist") sind wie einfache Bauteile. Sie benötigen wenig Verarbeitung.
Schwierige Wörter (wie komplexe Konzepte oder seltene Namen) sind wie defekte Bauteile. Sie benötigen zusätzliche Zeit und Aufmerksamkeit, um sie zu reparieren.

Derzeit verschwendet die Fabrik Energie, indem sie die einfachen Bauteile mit der gleichen Intensität verarbeitet wie die schwierigen.

Die Lösung: Der intelligente Wächter (TSA)

Die Autoren fügten zwischen jede Schicht der KI einen winzigen, leichten „Wächter" (ein kleines neuronales Netzwerk) ein. Dieser Wächter betrachtet jedes Wort (Token) einzeln und fragt: „Benötigt dieses Wort die nächste Verarbeitungsstation, oder kann es diese überspringen?"

Die Entscheidung: Der Wächter sagt nicht einfach nur „Ja" oder „Nein". Er gibt einen Wahrscheinlichkeitswert aus (wie ein Dimmschalter). Ist ein Wort einfach, könnte der Wächter sagen: „Du kannst den nächsten Schritt überspringen" oder „Du musst nur die Hälfte der Arbeit leisten". Ist ein Wort schwierig, sagt er: „Gehe den ganzen Weg durch."
Die Magie: Das Beste ist, dass der Wächter sich selbst unterrichtet. Er braucht keinen Menschen, der ihm sagt, welche Wörter schwierig sind. Er lernt rein durch den Versuch, Fehler zu minimieren. Wenn das Überspringen eines Schritts für ein bestimmtes Wort einen Fehler verursacht, lernt der Wächter, dieses Wort beim nächsten Mal aktiv zu halten. Wenn das Überspringen gut funktioniert, lernt er, Energie zu sparen.

Funktionsweise in der Praxis

Die Arbeit testete dies an zwei Hauptaspekten:

Einfache Logikrätsel: Wenn die KI eine Liste von Zahlen kopieren musste, stellte der Wächter fest: „Das ist einfach", und übersprang etwa 65 % der Arbeit. Wenn sie Zahlen sortieren musste (was schwieriger ist), übersprang sie nur etwa 27 % der Arbeit. Sie erkannte auf natürliche Weise, dass schwierigere Aufgaben mehr Schritte benötigen.
Geschichten schreiben: Wenn die KI gebeten wurde, wie Shakespeare zu schreiben oder Wikipedia-Artikel zusammenzufassen, sparte sie 14 % bis 23 % ihrer Rechenleistung.
- Sie lernte, dass Satzzeichen, Leerzeichen und gebräuchliche Wörter einfach sind und schnell verarbeitet werden können.
- Sie lernte, dass einzigartige Inhaltsstoffe die volle „Fabrik"-Behandlung benötigen.

Die Ergebnisse: Schneller und intelligenter

Kein Qualitätsverlust: Die KI schrieb genauso gut wie das Standardmodell, verwendete jedoch deutlich weniger Energie (Rechenleistung).
Besser als „Early Exit": Andere Methoden versuchen, den gesamten Satz frühzeitig zu stoppen, wenn die KI sich „sicher" fühlt. TSA ist intelligenter; es verhindert, dass einzelne Wörter unnötige Schritte durchlaufen, während schwierige Wörter weiterlaufen dürfen. Die Arbeit ergab, dass TSA bessere Ergebnisse lieferte als diese älteren Methoden, wenn diese so eingestellt waren, dass sie die gleiche Menge an Energie sparen.
Echte Geschwindigkeit: Als die Forscher den Code tatsächlich auf einem Computer ausführten, sahen sie eine echte Geschwindigkeitssteigerung (etwa 2,3 % schneller), weil der Computer für die übersprungenen Wörter buchstäblich keine Mathematik mehr durchführen musste.

Das Fazit

Diese Arbeit präsentiert einen Weg, KI-Modelle auf die klügste mögliche Weise „faul" zu machen. Anstatt jedes Wort zur gleichen Menge an Arbeit zu zwingen, lässt TSA die KI entscheiden, welche Wörter einfach sind und welche zusätzliche Hilfe benötigen, und zwar in Echtzeit. Es ist wie eine Fabrik, in der die einfachen Produkte durch die Linie rasen, während die schwierigen die volle VIP-Behandlung erhalten, was Zeit und Energie spart, ohne die Qualität zu beeinträchtigen.

Technische Zusammenfassung: Adaptive Berechnungstiefe durch gelerntes Token-Routing in Transformern

Problemstellung

Standard-Transformer-Architekturen wenden eine feste Anzahl von Schichten auf jedes Token in einer Sequenz an, unabhängig von der kontextuellen Schwierigkeit oder Informationsdichte dieses Tokens. Dieses Design priorisiert architektonische Einfachheit gegenüber der Anpassungsfähigkeit pro Token, was zu erheblicher Rechenineffizienz führt. In der Praxis benötigen häufige Tokens in vorhersehbaren Kontexten weitaus weniger Verarbeitung als seltene Tokens in neuartigen Konstruktionen, erhalten jedoch bei jedem Vorwärtspass identische Rechenressourcen. Diese Ineffizienz ist insbesondere im Inferenzmaßstab kostspielig, wo die dominierenden Kosten den vollständigen Vorwärtspass durch alle Schichten für alle Tokens betreffen. Bestehende Lösungen wie Adaptive Computation Time (ACT) für RNNs, Universal Transformers, Mixture-of-Depths (MoD) und Early-Exit-Mechanismen versuchen dies zu adressieren, verlassen sich jedoch oft auf harte Selektion, akkumulierte Budgets oder nachgelagerte Konfidenzschwellen, die möglicherweise nicht vollständig differenzierbar oder optimal sind.

Methodik: Token-Selektive Aufmerksamkeit (TSA)

Der Artikel schlägt Token-Selektive Aufmerksamkeit (TSA) vor, einen Mechanismus, der ein gelerntes, pro-Token-Gatter für Residual-Updates zwischen aufeinanderfolgenden Transformer-Blöcken einführt.

Architektur und Routing

TSA fügt nach jedem Block $f_l$ (für $l = 0, \dots, L-2$ ) in einem Decoder-only-Transformer mit Pre-Norm einen leichten Router $r_l$ ein. Der Stem-Block $f_0$ wird immer bedingungslos ausgeführt.

Router-Design: Jeder Router ist ein zweischichtiger Multi-Layer Perceptron (MLP) mit sigmoider Ausgabe, der eine kontinuierliche Halte-Wahrscheinlichkeit $p_l \in (0, 1)$ für jedes Token erzeugt. Die verborgene Dimension beträgt $d/4$ (abgerundet auf 16) und fügt dem Basismodell einen Parameter-Overhead von etwa 1,7 % hinzu.
Gesteuertes Update: Die Residual-Updates für den nachfolgenden Block $f_{l+1}$ $f_{l + 1}$ werden mit $(1 - p_l)$ $(1 - p_{l})$ skaliert.
- Wenn $p_l = 0$ , ist das Update standardmäßig.
- Wenn $p_l = 1$ , bleibt der Zustand unverändert, wodurch der Block effektiv übersprungen wird.
- Die Interpolation ist glatt und stellt sicher, dass der Mechanismus vollständig differenzierbar ist, ohne Straight-Through-Schätzer, Gumbel-Sampling oder Reinforcement Learning zu benötigen.

Training und Regularisierung

Während des Trainings werden alle Schichten vollständig ausgeführt; das Gatter skaliert die Residual-Updates, überspringt jedoch keine Berechnungen. Der Gesamtverlust ist definiert als $L = L_{task} + L_{depth}$ .

Tiefen-Regularisierung: Ein Term $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ wird hinzugefügt, um ein frühes Halten sanft zu fördern.
Emergentes Verhalten: Bemerkenswerterweise zeigt der Artikel, dass der Routing-Mechanismus lernt, Operationen auch dann zu überspringen, wenn $\lambda = 0$ (kein expliziter Tiefendruck), angetrieben ausschließlich durch den Task-Verlust-Gradienten.

Inferenz

Bei der Inferenz wechselt der Mechanismus in einen „dünnbesetzten" (sparse) Modus. Tokens mit $p_l > 0.5$ werden über Gather/Scatter-Operationen aus dem Feed-Forward-Netzwerk (FFN) ausgeschlossen, während die Aufmerksamkeit dicht bleibt, um die exakte Key-Value-Semantik zu bewahren. Dies führt zu echten FLOPs-Einsparungen und Beschleunigungen der Wandzeit.

Hauptbeiträge

Differenzierbares Token-Routing: Ein einfaches, kontinuierliches weiches Gatter für Residual-Updates, das ein Routing pro Token und pro Schicht ermöglicht, ohne architektonische Änderungen am Basis-Transformer.
Gradientengetriebene Emergenz: Nachweis, dass ein schwierigkeit-proportionales Routing natürlich allein aus dem Task-Verlust-Gradienten entsteht. Selbst ohne Tiefen-Regularisierung ( $\lambda=0$ ) lernt der Router, etwa 20 % der Token-Schicht-Operationen zu überspringen.
Effizienz-Qualitäts-Abwägung: Validierung am Beispiel der Sprachmodellierung auf Zeichenebene, die Einsparungen von 14–23 % bei Token-Schicht-Operationen (TLOps) bei einem Qualitätsverlust von weniger als 0,5 % zeigt.
Überlegenheit gegenüber Early Exit: Bei angepasster Effizienz erzielt TSA einen niedrigeren Validierungsverlust als Early-Exit-Baselines, was darauf hindeutet, dass end-to-end erlernte Routing-Entscheidungen nachgelagerten Konfidenzschwellen überlegen sind.
Echte Beschleunigung: Demonstration einer echten Beschleunigung der Wandzeit (bis zu 2,3 %) auf handelsüblicher Hardware (Apple M1 Pro) durch sparse Inferenz.

Experimentelle Ergebnisse

Synthetische Aufgaben

Bei Copy- und Sortier-Aufgaben (Sequenzen der Länge 10):

Copy-Aufgabe: Der Router lernte, dass die Aufgabe eine Identitätsabbildung war, und übersprang 54,9 % der TLOps ( $\alpha = 0,341$ ) ohne Genauigkeitsverlust.
Sortier-Aufgabe: Der Router wies mehr Rechenleistung zu und übersprang 22,5 % der TLOps ( $\alpha = 0,730$ ), was die höhere Schwierigkeit widerspiegelt.

Sprachmodellierung auf Zeichenebene

Tiny-Shakespeare: TSA sparte 22,8 % der TLOps ( $\alpha = 0,726$ ) bei einer Erhöhung des Validierungsverlusts um nur 0,006 Nats (+0,4 %).
enwik8: Auf einem diverseren Korpus (Wikipedia) sparte TSA 13,9 % der TLOps ( $\alpha = 0,833$ ) ohne Qualitätsverschlechterung (marginally besserer Validierungsverlust).
Robustheit: Der Mechanismus ist robust gegenüber dem Regularisierungs-Hyperparameter $\lambda$ über zwei Größenordnungen hinweg ( $[0, 0,1]$ ). Bei $\lambda=0,05$ wurden 50,4 % TLOps-Einsparungen bei vernachlässigbarem Qualitätskosten erreicht.

Vergleich mit Early Exit

Bei einem angepassten aktiven Anteil ( $\alpha \approx 0,726$ ) erzielte TSA einen 0,71 % niedrigeren Validierungsverlust als eine Early-Exit-Baseline. Dies zeigt, dass das end-to-end Lernen von Routing-Entscheidungen bessere Entscheidungen trifft als die Anwendung einer Konfidenzschwelle zur Inferenzzeit.

Durchsatz der Wandzeit

Soft-TSA (Training): Fugte unabhängig von $\alpha$ einen Overhead von ca. 1 % hinzu.
Sparse-TSA (Inferenz): Erzielte eine Beschleunigung von 2,3 % bei Tiny-Shakespeare ( $\alpha=0,726$ ) und eine Break-even-Leistung bei enwik8 ( $\alpha=0,833$ ). Beschleunigungen erforderten eine Batch-Größe von $\ge 64$ , um den Overhead der CPU-GPU-Synchronisation zu überwinden.

Bedeutung und Behauptungen

Der Artikel behauptet, dass TSA eine minimale, architektonisch einfache Lösung für die Ineffizienz von Tiefen-festen Transformern bietet. Seine primäre Bedeutung liegt darin, zu demonstrieren, dass adaptive Berechnungstiefe direkt aus dem Task-Ziel gelernt werden kann, ohne explizite Überwachung bezüglich der Schwierigkeit oder komplexer Reinforcement-Learning-Signale.

Die Autoren betonen, dass der Router als „gelerntes Noise-Gate" fungiert, das Updates abschwächt, wenn Residual-Beiträge redundant sind. Dies ermöglicht dem Modell, Rechenleistung dynamisch basierend auf der Token-Schwierigkeit zuzuteilen und erhebliche Effizienzgewinne (bis zu 55 % bei synthetischen Aufgaben, 14–23 % bei Sprachmodellierung) bei minimalen Qualitätskompromissen zu erzielen. Die Arbeit legt nahe, dass end-to-end differenzierbares Routing eine effektivere Strategie für adaptive Tiefe ist als diskrete Selektion oder nachgelagerte Schwellenwertbildung und den Weg für effizientere Inferenz auf handelsüblicher Hardware ebnet.

Der Artikel bleibt hinsichtlich des Maßstabs bescheiden und stellt fest, dass alle Experimente an Modellen mit etwa 5–6 Millionen Parametern durchgeführt wurden und das Skalierungsverhalten auf größere Modelle (10M–100M+) eine offene Frage bleibt. Zudem sind die aktuellen Beschleunigungen batch-größenabhängig, und die Aufmerksamkeitssparsamkeit wurde noch nicht implementiert.