📊 statistics

Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models

Dieses Paper schlägt PACE vor, einen Optimizer-Wrapper für AdamW, der das Training als ein optimales Steuerungsproblem behandelt, um den Fehler von iterativ gemittelten Sprachmodellen zu minimieren, wobei sowohl theoretische Konvergenzgarantien als auch empirische Verbesserungen bei Supervised Fine-Tuning- und Pretraining-Aufgaben nachgewiesen werden.

Ursprüngliche Autoren: Kwok Chun Au, Adam Block

Veröffentlicht 2026-06-25

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Kwok Chun Au, Adam Block

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 trainieren ein riesiges, digitales Gehirn (ein Sprachmodell), um eine neue Fähigkeit zu erlernen, wie zum Beispiel das Schreiben von Poesie oder das Programmieren. Dies tun Sie, indem Sie ihm tausende Beispiele zeigen und es Schritt für Schritt seine internen „Gewichte“ (sein Wissen) anpassen lassen.

Normalerweise, wenn Sie das Training beenden, nehmen Sie die allerletzte Version des Gehirns und sagen: „Dies ist das fertige Produkt.“ Aber viele moderne Forscher haben eine Geheimnis entdeckt: Der Durchschnitt aller Versionen, die das Gehirn während des Trainings durchlaufen hat, ist oft klüger und stabiler als die letzte Version allein. Es ist so, als würde man sagen: „Der Durchschnitt all meiner Übungstests ist ein besserer Prädiktor für meine Abschlussprüfung als der eine, den ich am allerletzten Tag abgelegt habe.“

Das Problem ist, dass der Trainingsprozess nicht weiß, dass er nach diesem „Durchschnitt“ beurteilt wird. Er versucht nur, das Ziel so schnell wie möglich zu erreichen, was den Pfad wackelig machen und den endgültigen Durchschnitt weniger perfekt machen kann.

Dieses Paper stellt eine neue Methode namens PACE (Pullback Averaging Control for Efficient Optimization) vor, um dies zu beheben. So funktioniert es, unter Verwendung einfacher Analogien:

1. Das Problem: Der wackelige Seiltanz

Stellen Sie sich das Training eines Sprachmodells wie einen Seiltanz vor, bei dem man auf ein bestimmtes Ziel zugeht (die perfekte Antwort).

Standardmäßiges Training (AdamW): Sie machen Schritte nach vorne. Manchmal treten Sie zu weit, manchmal wackeln Sie. Wenn Sie am Ende ankommen, sind Sie vielleicht leicht aus der Mitte geraten.
Der „Durchschnitt“-Trick: Anstatt darauf zu schauen, wo Sie gelandet sind, entscheiden wir uns, den Durchschnitt aller Orte anzusehen, an denen Sie während des Spaziergangs gestanden haben.
Das Problem: Wenn Sie wissen, dass Sie nach Ihrem Durchschnittsstandort beurteilt werden, sollten Sie anders laufen. Sie sollten nicht einfach nur zum Ende eilen; Sie sollten versuchen, Ihren Pfad die ganze Zeit über stabil und zentriert zu halten. Aber Standard-Trainingsalgorithmen wissen das nicht; sie eilen einfach nur vorwärts.

2. Die Lösung: Der „Geisterführer“ (Optimale Steuerung)

Die Autoren stellten eine grundlegende Frage: Wenn wir wissen, dass wir nach dem Durchschnitt beurteilt werden, wie sollten wir die Art und Weise ändern, wie wir laufen, um diesen Durchschnitt so gut wie möglich zu gestalten?

Sie nutzten fortgeschrittene Mathematik (speziell die „optimale Steuerungstheorie“, die zur Führung von Raketen und Robotern verwendet wird), um dies zu lösen. Sie stellten sich eine vereinfachte Welt vor, in der das Ziel ein perfektes Tal ist und das „Rauschen“ des Trainings wie Wind ist, der den Wanderer vom Kurs abbringt.

Sie berechneten die perfekte Strategie: Gehen Sie nicht nur vorwärts; schauen Sie gelegentlich auf Ihren „Geisterführer“ (den Durchschnitt dessen, wo Sie gewesen sind) und ziehen Sie sich sanft wieder zurück zu ihm.

3. Das Werkzeug für die reale Welt: PACE

Die perfekte mathematische Lösung war zu komplex, um sie auf echten, massiven Computern anzuwenden. Also erschufen die Autoren eine praktische, leichtgewichtige Version namens PACE.

Wie es funktioniert: PACE ist ein „Wrapper“, der über dem Standard-Trainingswerkzeug (AdamW) sitzt.
Der Mechanismus: Alle paar Schritte überprüft PACE den Unterschied zwischen dem Ort, an dem sich das Modell derzeit befindet, und dem „Durchschnitt“ seiner vergangenen Schritte.
Der Rückzug (Pullback): Wenn das Modell zu weit von seinem Durchschnitt abge drifted ist, stößt PACE es sanft wieder zurück. Es ist wie ein Trainer, der die Leine an einem Hund hält, der ständig vom Weg abläuft; der Trainer stoppt den Hund nicht, aber er zieht ihn sanft zurück zur Mitte der Gruppe.
Kluge Stöße: Die Stärke dieses Zuges ist nicht zufällig. Sie wird „geclippt“ (begrenzt), damit sie das Modell nicht zu hart herumreißt, und sie passt sich an, wie sicher sich das Modell in seinem aktuellen Schritt ist.

4. Die Ergebnisse: Eine sanftere Fahrt

Das Paper testete PACE auf mehreren verschiedenen Sprachmodellen (von 1 Milliarde bis 2 Milliarden Parametern).

Die Analogie: Stellen Sie sich zwei Läufer vor. Einer läuft wild und zickzackförmig zum Ziel (Standardmäßiges Training). Der andere läuft geschmeidig und korrigiert ständig seinen Pfad, um zentriert zu bleiben (PACE).
Das Ergebnis: Der „geschmeidige“ Läufer (PACE) beendete mit einer durchweg besseren „Durchschnittsposition“ als der wilde Läufer.
Wichtige Erkenntnis: PACE funktionierte besser als die Standardmethode, selbst wenn die Standardmethode erlaubt war, ihre Lernrate am Ende zu verlangsamen (ein gängiger Trick, um das Training zu stabilisieren). PACE erreichte diese Stabilität während des gesamten Trainingsprozesses, nicht erst am Ende.

Zusammenfassung

Kurz gesagt: Das Paper argumentiert, dass man anders trainieren sollte, wenn man plant, den Durchschnitt der Trainingsschritte als sein finales Modell zu verwenden. Man sollte nicht nur auf die Ziellinie abzielen; man sollte darauf abzielen, seine gesamte Reise zentriert zu halten. PACE ist ein einfaches, neues Werkzeug, das das Modell sanft zu seiner eigenen Geschichte zurückzieht, was zu einem klügeren, stabileren Endprodukt führt.

Was das Paper NICHT behauptet:

Es behauptet nicht, dass dies für medizinische Diagnosen oder klinische Anwendungen funktioniert.
Es behauptet nicht, dass dies für Modelle mit mehr als 2 Milliarden Parametern funktioniert (sie haben nur bis zu dieser Größe getestet).
Es behauptet nicht, dass es alle anderen Methoden ersetzt, sondern vielmehr die Standardmethode „AdamW“ in Kombination mit dem Mittelwertverfahren verbessert.

Technisches Resümee: Training für das zurückgegebene Modell

Problemstellung

Moderne Sprachmodell-Pipelines (LM) geben häufig ein gemitteltes Modell zurück – typischerweise ein exponentiell gleitender Durchschnitt (Exponential Moving Average, EMA) der Trainingsiterationen – anstatt der letzten Iteration selbst. Während die Iterationsmittelung bekannt dafür ist, die Optimierung zu stabilisieren und die Generalisierung zu verbessern, bleibt eine grundlegende Frage offen: Wenn das endgültige Ergebnis ein Iterationsdurchschnitt ist, wie sollte der Trainingsprozess selbst modifiziert werden, um die Leistung dieses spezifischen Schätzers zu optimieren?

Vorherige Arbeiten haben dies implizit adressiert, indem sie beobachteten, dass die Mittelung aggressivere Lernraten ermöglicht. Diese Arbeit verfolgt jedoch einen direkten Ansatz: Sie sucht nach einem Optimierungsalgorithmus, der die Trainings-Trajektorie explizit steuert, um den Fehler dieses zurückgegebenen Iterations-Durchschnitt-Schätzers zu minimieren.

Methodik

Theoretische Formulierung

Die Autoren rahmen das Problem als ein Optimalsteuerungsproblem innerhalb eines stetigen, zeitkontinuierlichen stochastischen quadratischen Modells.

Dynamik: Die Parameter-Trajektorie $\theta_t$ folgt einer stochastischen Differentialgleichung (SDE), die einen ungeordneten Gradientenabstieg in der Nähe eines quadratischen Minimums (Ornstein–Uhlen-Prozess) darstellt, ergänzt durch einen Steuerungs-Input $u_t$ .
Zielsetzung: Das Ziel ist es, eine Steuerungsstrategie $u_t$ zu finden, die den mittleren quadratischen Fehler zwischen dem zurückgegebenen Zeit-Durchschnitts-Schätzer $\hat{\theta}_{EMA, T}$ und dem wahren Optimum $\mu^*$ minimiert, unter Berücksichtigung eines quadratischen Strafterms auf die Magnitude der Steuerungsintervention (Regularisierungsparameter $\lambda$ ).
Ableitung: Unter Verwendung der Linear-Quadratic-Gaussian (LQG) Regelungstheorie und des Prinzips der Gewissheitsäquivalenz (Certainty Equivalence) leiten die Autoren einen optimalen Controller ab. Die exakte Lösung hängt vom Posterior-Mittelwert des Optimums gegeben die Trajektorie ab, was für großskalige LMs unpraktikabel ist.

Praktische Approximation: PACE

Um die Lösung implementierbar zu machen, führen die Autoren eine Reihe heuristischer Approximationen ein:

Schätzer-Substitution: Ersetzen des unbekannten Posterior-Mittelwerts des Optimums durch den aktuellen Iterations-Durchschnitt ( $\hat{\theta}_{EMA, t}$ ).
Zeit-Approximation: Annahme, dass die aktuelle Zeit $t$ viel kleiner als die gesamte Trainingszeit $T$ ist ( $t \ll T$ ), was die zeitabhängigen Koeffizienten vereinfacht.
Diskretisierung und Preconditioning: Diskretisierung der kontinuierlichen Steuerung und Ersetzung der unbekannten Hesse-Matrix durch den diagonalen Preconditioner, der in AdamW verwendet wird (den zweiten Moment-Schätzer $v_t$ ).

Dies ergibt PACE (Pullback Averaging Control for Efficient Optimization), einen leichtgewichtigen Wrapper um AdamW. Die Update-Regel fügt einen „Pullback“-Term hinzu, der die Live-Gewichte in Richtung der EMA-Gewichte bewegt:
$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$
wobei $C$ eine diagonale Matrix ist, die aus dem Adam-Preconditioner abgeleitet wird, und die Stärke des Pullbacks begrenzt wird, um sicherzustellen, dass das Update eine konvexe Kombination des Standard-Schritts und des Pullbacks bleibt.

Zentrale Beiträge

Optimalsteuerung-Formulierung: Die Arbeit liefert die erste fundierte Ableitung eines Optimierungsalgorithmus, der speziell darauf ausgelegt ist, den Fehler eines Iterations-Durchschnitt-Schätzers zu minimieren, anstatt der letzten Iteration.
Theoretische Garantien:
- Konvexer Kontext: Die Autoren beweisen, dass eine stilisierte Version von PACE gegen das Minimum jeder konvexen Verlustfunktion mit unverzerrten stochastischen Gradienten mit der Standardrate von $O(1/\sqrt{T})$ konvergiert, bis auf einen konstanten Faktor, der von der EMA-Stärke abhängt. Dies stellt sicher, dass die Methode niemals signifikant schlechter als das Standard-SGD ist.
- Quadratischer Kontext: Im quadratischen Setting, das für die Ableitung verwendet wurde, verbessern die gesteuerten Dynamiken den begrenzenden mittleren quadratischen Fehler des Iterations-Durchschnitt-Schätzers im Vergleich zu ungesteuerten Dynamiken (Standard-SGD mit EMA) strikt. Die Verbesserung kann auf bestimmten quadratischen Instanzen beliebig groß sein.
Empirische Validierung: Die Methode wird evaluiert auf:
- Fine-tuning: SmolLM2-1.7B, Qwen3-1.7B und Gemma3-1B Modelle auf dem smol-smoltalk Datensatz.
- Pretraining: GPT-2 (124M) auf FineWeb bei einem Chinchilla-optimalen Token-Budget.
- Baselines: PACE wird gegen Vanilla AdamW, AdamW mit EMA-Evaluierung und den Schedule-Free Algorithmus verglichen.

Ergebnisse

Leistung: PACE verbessert die Leistung gegenüber sowohl Vanilla AdamW als auch AdamW mit EMA-Evaluierung über alle getesteten Modelle (1B–2B Parameter) und Datensätze hinweg strikt.
Robustheit: Die Verbesserung ist robust gegenüber einer breiten Palette von Lernraten, Decay-Schedules (Constant, Cosine, Warmup-Stable-Decay) und Hyperparametern (Pullback-Stärke $c$ , EMA-Power $\kappa$ ).
Vergleich zu Schedule-Free: PACE ist wettbewerbsfähig mit dem Schedule-Free Algorithmus (Defazio et al., 2024), der ebenfalls Iterationsmittelung nutzt, um die Notwendigkeit von Lernraten-Decay zu eliminieren. PACE erreicht dies jedoch mit einer anderen Update-Regel.
Lernraten-Decay: PACE funktioniert auch mit konstanten Lernraten gut, was darauf hindeutet, dass es das Training in Regimen stabilisieren kann, in denen ein Lernraten-Decay normalerweise erforderlich wäre. In Token-Budget-Vergleichen folgt die PACE-Constant-LR-Trajektorie dem „Flussboden“ der Loss-Landschaft so effektiv wie WSD-Zweige und übertrifft diese oft sogar.

Bedeutung und Limitationen

Bedeutung:
Die Arbeit zeigt, dass die Optimierung für das zurückgegebene Modell (den Durchschnitt) anstatt für die letzte Iteration spürbare Leistungssteigerungen führt. Indem sie die Trainings-Trajektorie als ein Kontrollproblem behandelt, bietet PACE eine theoretisch fundierte Methode, um das Training zu stabilisieren und die Generalisierung zu verbessern, ohne komplexe Schedule-Abstimmungen zu erfordern. Die Autoren merken an, dass PACE eine verbesserte theoretische Konvergenz im quadratischen Setting im Vergleich zu Standardmethoden bietet, was jedoch mit einer etwas langsameren theoretischen Garantie im allgemeinen konvexen Setting aufgrund des EMA-Faktors einhergeht.

Limitationen:

Theoretische Annahmen: Die Ableitung beruht auf einem vereinfachten quadratischen Modell mit additivem Rauschen. Während die Konvergenz für allgemeine konvexe Funktionen bewiesen ist, bleibt die rigorose Einbeziehung von Momentum und adaptiver Preconditioning (zentral für AdamW) in die theoretische Analyse eine offene Forschungsrichtung.
Skalierung: Die Experimente beschränkten sich auf Modelle bis zu 2B Parametern beim Fine-tuning und 124M beim Pretraining. Die Wirksamkeit bei größeren Skalen (z. B. 70B+) ist nicht verifiziert.
Speicher-Overhead: PACE erfordert das Vorhalten einer zusätzlichen Kopie der Modellgewichte (den EMA-Zustand) während des Trainings, was einen nicht unerheblichen Speicher-Overhead verursacht, der in der Praxis jedoch durch große Batch-Größen abgemildert wird.
Seed-Sensitivität: Aufgrund von Rechenbeschränkungen verwendeten die Hauptläufe einen einzelnen Seed, obwohl Studien zur Robustheit über mehrere Seeds die Stabilität der Ordnung der Methode bestätigten.

Die Autoren kommen zu dem Schluss, dass PACE einen vielversprechenden Schritt darstellt, um Optimierungsalgorithmen mit der tatsächlichen Deployment-Strategie moderner LMs in Einklang zu bringen, bei denen das gemittelte Modell das Endprodukt ist.