OptiMer: Optimal Distribution Vector Merging Is Better than Data Mixing for Continual Pre-Training

Die Arbeit stellt OptiMer vor, ein Verfahren, das die optimale Mischung von Trainingsdaten für das kontinuierliche Vor-Training von Sprachmodellen durch eine effiziente, nachträgliche Optimierung von Verteilungsvektoren bestimmt und dabei teures Hyperparameter-Tuning während des Trainings vermeidet.

Das Wesentliche

Das Problem: Der "Rezept-Fehler" beim Kochen

Stell dir vor, du hast einen riesigen, sehr klugen Koch (das ist dein KI-Modell, z. B. Gemma 3). Dieser Koch kann schon gut kochen, aber er kennt die japanische Küche oder Mathe-Aufgaben noch nicht perfekt.

Um ihn darin zu schulen, musst du ihm neue Zutaten geben (das sind die Trainingsdaten). Das Problem ist: Du musst entscheiden, wie viel von welchem Gericht du ihm gibst.

• Gibst du ihm zu viel Mathematik? Dann vergisst er vielleicht, wie man höflich auf Japanisch spricht.
• Gibst du ihm zu viel Code? Dann wird er vielleicht schlecht im Schreiben von Geschichten.

Bisher mussten die Forscher vor dem eigentlichen Lernen raten: "Ich mische 30 % Mathe und 70 % Japanisch."

• Das Risiko: Wenn sie sich irren, haben sie wochenlang Rechenzeit (und viel Geld) verschwendet, bis sie merken: "Oh nein, das schmeckt nicht!" Dann müssen sie von vorne anfangen. Das ist wie ein Koch, der einen ganzen Tag lang einen Kuchen backt, nur um festzustellen, dass er vergessen hat, Zucker reinzutun.

Die Lösung: OPTIMER – Das "Zutaten-Post-Processing"

Die Forscher aus Japan haben eine clevere Idee namens OPTIMER entwickelt. Statt den Koch sofort mit der perfekten Mischung zu füttern, machen sie etwas anderes:

1. Einzelne Meisterkurse: Sie schicken den Koch erst mal nur mit einer Zutat in die Schule.

   • Einmal nur mit Mathe.
   • Einmal nur mit Japanisch.
   • Einmal nur mit Programmieren.
   • Einmal nur mit Code.
     Jeder Kurs ist kurz und teuer, aber sie machen sie alle parallel.

2. Der "Geschmacks-Fingerabdruck": Nach jedem Kurs nehmen sie einen "Abdruck" davon, wie sich der Koch durch diese spezielle Zutat verändert hat. In der Wissenschaft nennen sie das einen Verteilungsvektor. Stell dir das vor wie einen kleinen Stempel, der genau zeigt: "Hier habe ich Mathe gelernt, hier habe ich Japanisch gelernt."

3. Der Zaubertrank (Die Mischung): Jetzt kommt der Clou. Statt den Koch noch einmal wochenlang zu trainieren, nehmen sie diese "Stempel" (die Vektoren) und mischen sie nachträglich zusammen.

   • Sie nutzen einen intelligenten Algorithmus (Bayesian Optimization), der wie ein super-schneller Tester funktioniert.
   • Der Algorithmus probiert in wenigen Minuten aus: "Was passiert, wenn ich 60 % Mathe-Stempel und 40 % Japanisch-Stempel mische?"
   • Er findet die perfekte Mischung, ohne dass der Koch noch einmal kochen muss.

Warum ist das so genial?

• Geschwindigkeit: Das alte Verfahren (alles auf einmal mischen und hoffen) dauert Wochen. OPTIMER findet die perfekte Mischung in Minuten. Es ist 15- bis 35-mal schneller.
• Flexibilität: Stell dir vor, du hast einmal die "Stempel" für Mathe, Japanisch und Code gespeichert.
  • Morgen willst du einen Koch, der nur Mathe kann? Kein Problem! Du mischst die Stempel neu (mehr Mathe, weniger anderes).
  • Übermorgen willst du einen Koch, der alles kann? Du mischst sie anders.
  • Du musst den Koch nicht neu ausbilden. Du hast eine "Werkbank" mit fertigen Zutaten, aus denen du jeden Tag ein neues Gericht zaubern kannst.

Ein einfaches Bild: Das Lego-Set

Stell dir das KI-Modell als ein riesiges Lego-Haus vor.

• Der alte Weg: Du versuchst, das Haus zu bauen, indem du alle Lego-Steine (Daten) in einen Eimer wirfst und hoffst, dass sie sich von selbst zu einem schönen Haus zusammenfügen. Wenn das Haus schief wird, musst du den ganzen Eimer leeren und von vorne anfangen.
• Der OPTIMER-Weg: Du baust erst ein kleines, perfektes Mathe-Haus. Dann ein perfektes Japanisch-Haus. Dann nimmst du dir die Pläne dieser Häuser (die Vektoren).
  • Jetzt kannst du die Pläne mischen, um ein neues, riesiges Haus zu entwerfen, das beides perfekt kann.
  • Wenn du merkst, dass das neue Haus zu viel Japanisch hat, änderst du einfach den Plan (die Gewichtung), ohne die Steine neu zu schmelzen.

Das Ergebnis

Die Forscher haben gezeigt, dass dieser neue Weg (OPTIMER) nicht nur schneller ist, sondern auch bessere Ergebnisse liefert. Die KI versteht sowohl Japanisch als auch Mathe besser als mit den alten Methoden. Und das Beste: Man kann die KI jederzeit auf neue Ziele anpassen, ohne sie erneut stundenlang trainieren zu müssen.

Kurz gesagt: Statt zu raten, wie man Zutaten mischt, bevor man kocht, kocht man erst kleine Proben, schmeckt sie ab und mischt dann die perfekten Rezepte zusammen – alles in Rekordzeit.

Technische Zusammenfassung

1. Problemstellung

Die Anpassung von Large Language Models (LLMs) an spezifische Sprachen und Domänen erfolgt häufig durch kontinuierliches Vor-Training (Continual Pre-Training, CPT). Ein zentrales Problem dabei ist die Bestimmung des optimalen Mischungsverhältnisses (Mixing Ratio) der verschiedenen Trainingsdatensätze (z. B. Mathematik, Code, verschiedene Sprachen).

• Herausforderung: Das Mischungsverhältnis ist ein kritischer Hyperparameter, der vor Trainingsbeginn festgelegt werden muss.
• Kosten: Eine suboptimale Wahl führt zu schlechteren Modellleistungen und verschwendet enorme Rechenressourcen (Wochen an GPU-Zeit), da das Training neu gestartet werden muss, um das Verhältnis zu korrigieren.
• Bestehende Ansätze: Methoden wie DoReMi oder RegMix versuchen, Verhältnisse über Proxy-Modelle zu schätzen, erfordern jedoch immer noch eine feste Vorgabe vor dem eigentlichen Training und können nicht nachträglich korrigiert werden.

2. Methodik: OPTIMER

Die Autoren schlagen OPTIMER vor, einen Ansatz, der die Auswahl der Datenverhältnisse vom eigentlichen Modelltraining entkoppelt. Statt ein einziges Modell auf einer gemischten Datenmenge zu trainieren, wird ein modulares Vorgehen gewählt:

• Unabhängiges Training: Für jeden einzelnen Datensatz $D_i$ wird ein separates CPT-Modell $\theta_{CPT_i}$ von einem Basis-Modell $\theta_{pt}$ aus trainiert.
• Extraktion von Verteilungsvektoren (Distribution Vectors): Aus jedem trainierten Modell wird ein Vektor $\tau_i$ extrahiert, der die durch den Datensatz induzierte Parameteränderung darstellt:
  $$\tau_i = \theta_{CPT_i} - \theta_{pt}$$
  Zusätzlich wird ein Vektor $\tau_{it}$ aus einem instruktionsfeingetunten (IT) Modell extrahiert, um die Instruktionstauglichkeit zu erhalten.
• Post-Hoc-Optimierung: Anstatt die Daten zu mischen, werden die Vektoren linear kombiniert, um ein neues Modell zu erzeugen:
  $$\theta_{merge} = \theta_{pt} + \alpha_{it} \cdot \tau_{it} + \sum_{i=1}^{n} \alpha_i \cdot \tau_i$$
  Dabei sind $\alpha_i$ die zu optimierenden Gewichte.
• Bayessche Optimierung (TPE): Da die Bewertungsfunktion (Score auf einem Entwicklungssatz) keine Gradienten liefert, wird die Suche nach den optimalen Gewichten $\alpha^*$ als Black-Box-Optimierungsproblem behandelt. Statt einer ineffizienten Gittersuche (Grid Search) wird der Tree-structured Parzen Estimator (TPE) verwendet, um die Gewichte effizient zu optimieren.

3. Wichtige Beiträge

1. Konzept der Verteilungsvektoren: Einführung von Verteilungsvektoren für CPT, die es ermöglichen, Datenverhältnisse nachträglich (post-hoc) zu optimieren, ohne das Modell erneut zu trainieren.
2. OPTIMER-Framework: Ein automatisiertes Framework, das die Suche nach optimalen Mischgewichten durch Bayessche Optimierung löst und so die Abhängigkeit von vorab festgelegten Hyperparametern beseitigt.
3. Flexibilität: Ein einmaliger Vektorpool kann für unterschiedliche Ziele (z. B. Fokus auf Mathematik vs. Fokus auf Japanisch) neu optimiert werden, um maßgeschneiderte Modelle „on demand" zu erzeugen.
4. Interpretierbarkeit: Die gefundenen optimalen Gewichte können als effektive Datenmischungsverhältnisse interpretiert werden.

4. Ergebnisse und Evaluation

Die Methode wurde auf dem Gemma 3 27B-Modell über verschiedene Sprachen (Japanisch, Chinesisch) und Domänen (Mathematik, Code) getestet.

• Leistung: OPTIMER übertrifft konsistent Baseline-Methoden (einfache Datenmischung, Task Arithmetic, TIES, DARE) in allen getesteten Kombinationen. Die Verbesserungen liegen zwischen 2,1 und 6,7 Punkten im Durchschnitt über mehrere Benchmarks.
• Effizienz: Die Suchkosten für OPTIMER sind 15- bis 35-mal niedriger als bei der Suche nach optimalen Datenmischungsverhältnissen durch wiederholtes Training (DataMix). Während ein DataMix-Versuch Tage bis Wochen dauert, sind OPTIMER-Experimente in Minuten abgeschlossen.
• Qualität der Vektoren: Die Analyse zeigt, dass Verteilungsvektoren verschiedener Datensätze annähernd orthogonal sind (Kosinus-Ähnlichkeit 0,03–0,31). Dies ermöglicht eine lineare Kombination ohne starke Interferenzen.
• Negativgewichte: Ein interessanter Befund ist, dass OPTIMER manchmal negative Gewichte zuweist, um den Einfluss bestimmter Verteilungen (z. B. Englisch bei einem japanisch-optimierten Modell) aktiv zu reduzieren und so Störungen zu eliminieren.
• Generalisierung: Die Methode wurde erfolgreich auf das SEA-LION v4 27B-Modell übertragen, was die Generalisierbarkeit über verschiedene Architekturen hinweg demonstriert.

5. Bedeutung und Fazit

Die Arbeit etabliert ein neues Paradigma für das kontinuierliche Vor-Training:

• Paradigmenwechsel: Die Entscheidung über Datenmischungsverhältnisse muss nicht mehr vor dem Training getroffen werden, sondern kann als post-hoc-Optimierungsproblem über Verteilungsvektoren reformuliert werden.
• Ressourceneffizienz: Durch die Vermeidung von teuren Neutraining-Versuchen bei suboptimalen Mischungen wird Rechenzeit massiv gespart.
• Zukunftsperspektive: Der Ansatz ermöglicht die Erstellung mehrerer, zielgerichteter Modelle aus einem einzigen Pool von Vektoren ohne weiteren Trainingsaufwand.

Zusammenfassend bietet OPTIMER eine robustere, effizientere und flexiblere Methode zur Anpassung von LLMs an neue Domänen und Sprachen im Vergleich zu traditionellen Datenmischungsansätzen.