Dynamic Weight Grafting: Localizing Finetuned Factual Knowledge in Transformers

Die Autoren stellen die Methode des dynamischen Gewichts-Transfers vor, um zu zeigen, dass feinabgestimmte Faktenwissen in Transformer-Modellen sowohl durch eine kontextuelle Anreicherung der Restverbindung während der Token-Verarbeitung als auch durch einen spezifischen Abrufmechanismus in den letzten Schichten vor der Vorhersage wiedergewonnen wird.

Das Wesentliche

Das große Rätsel: Wo landet das neue Wissen?

Stell dir vor, du hast einen sehr klugen, aber etwas verstaubten Bibliothekar (das ist dein KI-Modell, das vor dem Training schon viel weiß). Eines Tages bringst du ihm eine neue, wichtige Information bei: „Der neue Papst heißt Leo XIV."

Die Forscher stellen sich die Frage: Wo genau in diesem Bibliothekar wird diese Information gespeichert?

• Wird sie sofort in den Kopf des Bibliothekars geschrieben, sobald er den Namen „Leo" hört? (Wie ein Notizzettel, der sofort an den Regalrand geheftet wird).
• Oder denkt er erst am Ende des Satzes: „Aha, ich muss mir das jetzt gerade merken, bevor ich antworte"? (Wie ein kurzfristiges Gedächtnis, das nur kurz aktiviert wird).
• Oder passiert beides?

Bisher war das schwer zu untersuchen, weil die alten Methoden wie ein schwerer Hammer waren: Wenn man versuchte, einen Teil des Gehirns zu untersuchen, hat man dabei oft versehentlich alles andere zerstört oder überschrieben.

Die neue Methode: Der „Chirurgische Gewichtstransfer"

Die Forscher haben eine neue Technik erfunden, die sie „Dynamisches Gewichtstransferieren" (Dynamic Weight Grafting) nennen.

Stell dir das KI-Modell nicht als Gehirn, sondern als riesiges, modulares Lego-Schloss vor.

• Das alte Modell ist das Originalschloss.
• Das neue, trainierte Modell ist eine Version, bei der man ein paar spezielle Lego-Steine ausgetauscht hat, damit es die neue Information über den Papst kennt.

Die Forscher gehen jetzt nicht so vor, dass sie das ganze Schloss neu bauen. Stattdessen nehmen sie nur ganz bestimmte Lego-Steine aus dem neuen Schloss und tauschen sie während des Baus gegen die alten Steine im Originalschloss aus.

• Szenario A: Sie tauschen nur die Steine aus, die benutzt werden, wenn das Wort „Leo" gelesen wird.
• Szenario B: Sie tauschen nur die Steine aus, die benutzt werden, wenn das Modell die Antwort geben soll.
• Szenario C: Sie tauschen alles andere aus, aber lassen die wichtigen Teile des Originals.

So können sie genau sehen: Welche Steine sind wirklich nötig, damit das Modell die Antwort „Leo XIV" gibt?

Was haben sie herausgefunden?

Die Forscher haben zwei Hauptwege entdeckt, wie das Modell die Information nutzt. Man kann sich das wie zwei verschiedene Strategien vorstellen, um eine Frage zu beantworten:

1. Der „Anreicherungs-Weg" (Enrichment)

Stell dir vor, das Modell liest den Namen „Leo". An diesem Moment wird dieser Name im Inneren des Modells „aufgeladen" mit der Information „ist der Papst". Es ist, als würde man dem Namen einen glühenden Leuchtturm aufsetzen.

• Wenn das Modell später den Satz beendet, muss es nur noch diesen leuchtenden Namen finden und kann die Antwort ableiten.
• Ergebnis: In manchen Fällen reicht es, nur diesen ersten Moment zu „reparieren", und das Modell weiß die Antwort.

2. Der „Erinnerungs-Weg" (Recall)

Manchmal passiert nichts Besonderes, wenn der Name „Leo" gelesen wird. Der Name bleibt „dunkel". Aber ganz am Ende des Satzes, kurz bevor das Modell das Wort „Papst" aussprechen soll, schaltet es einen Super-Schalter ein.

• In diesem letzten Moment „ruft" das Modell die Information ab: „Moment mal, wer war das nochmal? Ach ja, Leo!"
• Ergebnis: In anderen Fällen reicht es, nur diesen letzten Schalter zu reparieren, und das Modell findet die Antwort, auch wenn der Name vorher „dunkel" war.

Die Überraschung: Oft braucht das Modell beide Wege, um sicher zu sein. Wenn man beide Wege „repariert" (also die richtigen Lego-Steine an den richtigen Stellen hat), funktioniert das Modell perfekt. Wenn man beide weglässt, vergisst es die Information komplett.

Wo sitzt dieser Schalter genau?

Die Forscher haben noch tiefer geschaut und herausgefunden, welche Lego-Steine genau diesen „Erinnerungs-Weg" steuern:

• Es sind bestimmte Aufmerksamkeits-Mechanismen (wie ein Suchscheinwerfer), die sich auf den Namen konzentrieren.
• Und es sind bestimmte Verarbeitungs-Zellen (die Feedforward-Netzwerke) ganz am Ende des Modells, die die Information tatsächlich in die Antwort umwandeln.

Warum ist das wichtig?

Früher dachte man, KI-Modelle speichern Fakten wie in einer Datenbank. Diese Studie zeigt aber, dass es viel dynamischer ist. Das Modell kann entscheiden, ob es eine Information sofort „einfärbt" oder sie sich erst am Ende „ins Gedächtnis ruft".

Der große Vorteil dieser neuen Methode:
Frühere Methoden waren wie ein Bulldozer: Man hat Teile des Modells zerstört, um zu sehen, was passiert. Diese neue Methode ist wie ein Chirurg mit einem Skalpell: Man tauscht nur winzige Teile aus, ohne den Rest des Gehirns zu stören. So verstehen wir endlich, wie KI wirklich „denkt" und lernt, ohne ihr dabei das Gehirn zu zerstören.

Zusammenfassung in einem Satz

Die Forscher haben eine neue Methode entwickelt, um genau zu sehen, dass KI-Modelle neue Fakten entweder sofort beim Lesen „einfärben" oder sie sich erst am Ende des Satzes „ins Gedächtnis rufen" – und dass sie dafür ganz spezifische, austauschbare Bauteile in ihrem Inneren nutzen.

Technische Zusammenfassung

1. Problemstellung

Die Arbeit adressiert eine fundamentale Frage im Bereich der interpretierbaren KI (Interpretability): Wenn ein Large Language Model (LLM) während des Fine-Tunings eine neue Tatsache lernt (z. B. eine neue Filmrolle oder einen neu gewählten Papst), wo genau wird diese Information im Modell gespeichert und wie wird sie während der Generierung abgerufen?

Bisherige Ansätze zur Lokalisierung von Wissen, wie z. B. Activation Patching (das Ersetzen von Aktivierungen im Residual-Stream) oder Ablationen, weisen eine kritische Einschränkung auf:

• Sie überschreiben Informationen, die bereits im Residual-Stream vorhanden sind.
• Dadurch wird der Zugriff auf vorherige Berechnungen blockiert. Es ist unmöglich zu unterscheiden, ob eine Komponente aktiv neue Informationen extrahiert oder lediglich Informationen weiterleitet, die bereits früher im Prozess „angereichert" wurden.
• Dies verhindert eine kausale Analyse darüber, welche Teile des Modells für das Abrufen von neuem Wissen notwendig und ausreichend sind.

2. Methodik: Dynamic Weight Grafting

Um diese Lücke zu schließen, schlagen die Autoren Dynamic Weight Grafting (Dynamisches Gewichts-Veredeln) vor. Dies ist eine Analyse-Technik, die es erlaubt, selektiv Teilmengen der Gewichte eines feinabgestimmten Modells (Fine-Tuned Model, $\theta_{SFT}$) in ein vortrainiertes Modell (Pretrained Model, $\theta_{PRE}$) einzufügen, während die Generierung stattfindet.

Kernprinzipien:

• Gewichts- statt Aktivierungs-Eingriff: Im Gegensatz zum Patching von Aktivierungen werden hier Parametermatrizen (Gewichte) ausgetauscht. Dies erhält die vorherigen Berechnungen im Residual-Stream intakt.
• Dynamische Selektivität: Das Grafting erfolgt tokenweise und komponentenweise. Man kann entscheiden, welche Gewichte (z. B. Attention-Matrizen oder Feedforward-Netzwerke) an welchen Positionen (z. B. beim ersten Entity-Token oder beim letzten Token vor der Vorhersage) aus dem feinabgestimmten Modell stammen.
• Formale Definition: Für ein gegebenes Token $t$ und eine Komponente $c$ wird das Gewicht $\tilde{\theta}_c(t)$ definiert als:
  $$\tilde{\theta}_c(t) = \begin{cases} \theta_{PRE, c} & \text{wenn } \gamma_c(t) = 0 \\ \theta_{SFT, c} & \text{wenn } \gamma_c(t) = 1 \end{cases}$$
  wobei $\gamma$ eine Maske ist, die steuert, welche Gewichte getauscht werden.

Experimentelles Setup:

• Modelle: Llama3, Pythia, GPT-2 XL, Gemma 1.1.
• Daten: Synthetische Datensätze mit fiktiven Filmen und echten/fiktiven Schauspielern, um spezifische Relationen zu lernen (z. B. „Wer hat in welchem Film mitgespielt?").
• Aufgabe: Relation Completion (Vervollständigung von Relationstupeln wie $(s, r, o)$).

3. Wichtige Beiträge und Ergebnisse

Die Anwendung von Dynamic Weight Grafting offenbarte zwei separate Pfade, über die feinabgestimmtes Faktenwissen abgerufen wird:

A. Zwei Abrufpfade (Pathways)

1. Der „Enrichment"-Pfad (Anreicherung):
   • Tritt auf, wenn das erste Entity-Token (z. B. „Zendaya") verarbeitet wird.
   • Das feinabgestimmte Modell „bereichert" die Repräsentation dieses Tokens mit relationalem Wissen in den unteren und mittleren Schichten.
   • Diese angereicherte Information wird durch den Residual-Stream weitergeleitet und kann später von generischen Mechanismen genutzt werden.
2. Der „Recall"-Pfad (Abruf):
   • Tritt am letzten Token vor der Generierung der Antwort auf.
   • Selbst wenn das Entity-Token nicht mit dem neuen Wissen angereichert wurde (d. h. es wurden nur die Gewichte der unteren Schichten des vortrainierten Modells verwendet), kann das Modell die Information am Ende des Satzes „erinnern".
   • Dies deutet darauf hin, dass die späteren Schichten des feinabgestimmten Modells einen Mechanismus besitzen, um Relationen aus dem Kontext zu extrahieren, auch wenn die Eingabe nicht explizit mit dem neuen Wissen modifiziert wurde.

Ergebnis: Das Grafting nur des ersten Entity-Tokens oder nur des letzten Tokens reicht in vielen Fällen aus, um die Leistung des vollständig feinabgestimmten Modells nahezu vollständig wiederherzustellen. Das Grafting aller anderen Teile (außer diesen beiden Pfaden) führt zu einer Genauigkeit nahe Null (ähnlich dem vortrainierten Modell).

B. Lokalisierung auf Komponenten-Ebene (Component Grafting)

Die Autoren lokalisierten den „Recall"-Pfad auf spezifische Komponenten:

• Aufgabe-spezifische Attention: Die Attention-Mechanismen auf dem ersten Entity-Token und dem letzten Token sind entscheidend, um die Aufgabe zu verstehen und den Kontext zu verknüpfen.
• Relation-spezifische Extraktion: Im Feedforward-Netzwerk (FFN) und insbesondere in der Output-Projektionsmatrix ($O$-Matrix) der letzten Schichten (vor der Vorhersage) findet die eigentliche Extraktion der Relation statt.
• Hybrid-Experimente: Durch das Mischen von Modellen, die nur auf der Aufgabenstruktur (Task-Model) vs. auf den spezifischen Relationen (Relation-Model) trainiert wurden, zeigte sich:
  • Die Attention-Mechanismen müssen die Aufgabenstruktur verstehen (Task-Model).
  • Die FFNs und die $O$-Matrix müssen die spezifischen Relationen kennen (Relation-Model).
  • Das Grafting von $O$-Matrix und FFN des Relation-Modells auf das Task-Model rekonstruierte fast die gesamte Leistung.

C. Generalisierung

Die Ergebnisse gelten sowohl für synthetische Daten als auch für reale Wikipedia-Artikel über Filme, die nach dem Release-Datum des Modells erschienen sind. Allerdings war der „Recall"-Pfad bei neueren Architekturen (Gemma, Llama3) stärker ausgeprägt als bei älteren Modellen (GPT-2 XL, Pythia), was auf fortschrittlichere Attention-Mechanismen hindeutet.

4. Signifikanz und Implikationen

• Überwindung der Limitationen von Activation Patching: Dynamic Weight Grafting ermöglicht eine kausale Analyse, ohne die Informationsflüsse zu zerstören. Es unterscheidet klar zwischen dem Speichern/Anreichern von Information und dem Abrufen derselben.
• Mechanistische Einsicht: Die Arbeit zeigt, dass LLMs nicht nur passiv Wissen speichern, sondern aktive Mechanismen implementieren, die je nach Kontext (Anreicherung vs. Abruf) funktionieren.
• Effizienz: Da oft nur ein kleiner Teil des Modells (bestimmte Token-Positionen und Komponenten) für das Abrufen neuer Fakten notwendig ist, könnte dies Ansätze für effizienteres Fine-Tuning oder gezieltes Knowledge Editing inspirieren.
• Architektur-Unterschiede: Die Studie liefert Hinweise darauf, wie sich verschiedene Transformer-Architekturen (z. B. RoPE vs. Learned Positional Embeddings, Norm-Typen) auf die Fähigkeit zur Wissensspeicherung und -abrufung auswirken.

Fazit:
Das Paper etabliert Dynamic Weight Grafting als ein leistungsfähiges Werkzeug für die mechanistische Interpretierbarkeit. Es beweist, dass das Abrufen von neuem Faktenwissen in Transformer-Modellen durch zwei komplementäre, aber teilweise redundante Pfade erfolgt: eine frühe Anreicherung der Entity-Repräsentation und einen späten, komponentenspezifischen Abrufmechanismus in den Feedforward- und Output-Schichten.