Dissecting Jet-Tagger Through Mechanistic Interpretability

Dieser Beitrag wendet Techniken der mechanistischen Interpretierbarkeit auf einen Particle Transformer Jet-Tagger an und zeigt, dass ein spärlicher Sechs-Kopf-Schaltkreis, der auf energiekorrelationsbasierten Darstellungen und einer spezifischen Quelle-Relais-Auslese-Architektur beruht, die vollständige Klassifikationsleistung des Modells wiederherstellen kann, während gleichzeitig demonstriert wird, dass Gradientenabstieg physikalisch bedeutsame Jet-Substrukturmerkmale auf natürliche Weise entdeckt.

Das Wesentliche

Das große Ganze: Die Blackbox öffnen

Stellen Sie sich einen hochqualifizierten Detektiv (ein Computerprogramm namens Particle Transformer) vor, der trainiert wurde, ein chaotisches Tatort (einen „Jet" von Teilchen, der in einem Teilchenbeschleuniger erzeugt wird) zu untersuchen und zu entscheiden: „War dies ein Top-Quark (der Verdächtige) oder nur zufälliger Hintergrundrauschen (QCD)?"

Lange Zeit wussten wir, dass der Detektiv den Fall außerordentlich gut löste, aber wir wussten nicht, wie er es tat. Es war eine „Blackbox". Dieses Paper ist wie die Anstellung eines forensischen Teams, um das Gehirn des Detektivs zu öffnen, genau zu kartieren, welche Neuronen feuern, und die schrittweise Logik zu erklären, die sie verwenden, um zu einem Urteil zu gelangen.

Das Gehirn des Detektivs: Ein Team von Spezialisten

Die Forscher fanden heraus, dass der Detektiv nicht sein ganzes Gehirn nutzt, um den Fall zu lösen. Stattdessen verlässt er sich auf ein winziges, effizientes Team von nur sechs Spezialisten (von 16 verfügbaren), das 97 % der schweren Arbeit leistet. Sie nennen dieses Team die „Schaltung" (Circuit).

So funktioniert dieses Sechser-Team, unter Verwendung einer Staffel-Relay-Analogie:

1. Der Kundschafter (Primäre Quelle): Ein Spezialist in der allerersten Schicht des Gehirns fungiert als Kundschafter. Diese Person sucht nicht direkt nach den „Bösewichten". Stattdessen scannt sie die Menge nach dem „Hintergrundrauschen" (weiche, kollidierende Teilchen) ab. Indem sie das Rauschen versteht, bereitet sie den Boden für alle anderen vor. Sie ist die wichtigste Person; wenn man sie entfernt, verliert das Team fast seine gesamte Fähigkeit, den Fall zu lösen.
2. Der zweite Kundschafter (Sekundäre Quelle): Ein weiterer Spezialist in der ersten Schicht hilft dem Kundschafter. Er ist dem Kundschafter sehr ähnlich, konzentriert sich jedoch auf leicht unterschiedliche Details.
3. Die Staffel-Läufer (Mittlere Schichten): Drei Spezialisten in den mittleren Schichten fungieren als Läufer. Sie nehmen die Informationen der Kundschafter und suchen nach etwas Spezifischem: schwere, energiereiche Paare von Teilchen. In der Welt der Teilchenphysik zerfällt ein Top-Quark in ein „W-Boson", das sich dann in zwei schwere Teilchen aufspaltet. Diese Läufer sind Experten darin, diese schweren Paare zu erkennen.
   • Kritische Entdeckung: Obwohl der Detektiv supposed ist, ein „Top-Quark" zu finden (was eine 3-teilige Struktur ist), suchen diese Läufer tatsächlich nur nach dem „W-Boson" (eine 2-teilige Struktur). Das Paper legt nahe, dass der Detektiv einen Abkürzungsweg gefunden hat: „Wenn ich das schwere 2-teilige W-Boson finden kann, kann ich ziemlich sicher sein, dass es ein Top-Quark ist." Es ist wie ein Detektiv, der einen Mord löst, indem er die Mordwaffe findet, anstatt zu versuchen, die gesamte Tatort-Szene zu rekonstruieren.
4. Der Richter (Ausgabe): Ein Spezialist in der letzten Schicht fungiert als Richter. Er schaut nicht direkt auf die Teilchen. Stattdessen nimmt er die Berichte der Staffel-Läufer, fasst sie zusammen und trifft die endgültige Entscheidung: „Schuldig" (Top-Quark) oder „Nicht schuldig" (Hintergrund).

Der „Aha!"-Moment: Es ist keine neue Idee, nur eine neue Sprache

Eine der überraschendsten Erkenntnisse des Papers betrifft den Zeitpunkt, zu dem der Detektiv seine Entscheidung trifft.

Normalerweise denken wir, der Detektiv sammelt Hinweise Schicht für Schicht und ruft plötzlich am ganz Ende: „Ich weiß, wer es war!" Die Forscher fanden jedoch heraus, dass der Detektiv die Antwort fast sofort kennt (nach der ersten Schicht des Scannens).

Warum sieht der letzte Schritt dann so dramatisch aus?

• Die Analogie: Stellen Sie sich vor, der Detektiv hat die Antwort in einem geheimen Code (einer anderen Sprache) in seiner ersten Schicht geschrieben. Der letzte Schritt ist kein „Denken" oder „Finden neuer Hinweise"; es ist einfach das Übersetzen dieses geheimen Codes in klare Sprache, damit der endgültige Richter es lesen kann.
• Das Paper nennt dies eine „Basis-Rotation". Die Information war bereits da; sie musste nur in die richtige Orientierung gedreht werden, um vom endgültigen Output verstanden zu werden.

Was hat der Detektiv tatsächlich gelernt?

Die Forscher haben auch überprüft, welche Art von „Physik" der Detektiv gelernt hat. Sie verglichen die internen Notizen des Detektivs mit Standard-Physikformeln, die von menschlichen Experten verwendet werden.

• Das Ergebnis: Der Detektiv ignorierte die komplexen, 3-teiligen Formeln, die Menschen normalerweise verwenden. Stattdessen entdeckte und bevorzugte er natürlich einfachere, 2-teilige Formeln (genannt Energie-Korrelatoren).
• Die Erkenntnis: Der Computer brauchte keinen Menschen, der ihm sagte: „Suche nach dem W-Boson!" Er fand selbst heraus, dass das Finden des schweren 2-teiligen Zerfalls der einfachste und zuverlässigste Weg war, das Rätsel zu lösen. Er entdeckte eine bedeutungsvolle physikalische Wahrheit neu, nur indem er versuchte, das Spiel zu gewinnen.

Zusammenfassung

Dieses Paper beweist, dass wir eine komplexe, moderne KI, die in der Hochenergiephysik eingesetzt wird, reverse-engineern können, um eine einfache, logische Schaltung im Inneren zu finden.

1. Es ist effizient: Ein winziges Team von 6 „Neuronen" erledigt fast die gesamte Arbeit.
2. Es ist logisch: Das Team folgt einem klaren Pfad: Rauschen scannen $\rightarrow$ schwere Paare weiterleiten $\rightarrow$ Ergebnis beurteilen.
3. Es ist intelligent: Die KI fand heraus, dass das Lösen eines einfacheren Teilproblems (Finden des 2-teiligen W-Bosons) der beste Weg ist, das große Problem zu lösen (Finden des Top-Quarks).
4. Es ist eine Übersetzung: Der letzte Schritt der KI ist nur das Übersetzen ihres frühen, geheimen Wissens in eine endgültige Antwort, keine neue Entdeckung.

Die Autoren kommen zu dem Schluss, dass die Werkzeuge, die wir zum Verständnis von KI-Sprachmodellen (wie Chatbots) verwenden, auch perfekt funktionieren, um KI in der Teilchenphysik zu verstehen, und zeigen, dass diese Maschinen tiefe physikalische Wahrheiten auf eigene Faust lernen können.

Technische Zusammenfassung

Technische Zusammenfassung: Zerlegung von Jet-Tagger durch mechanistische Interpretierbarkeit

Problemstellung
Deep-Learning-Architekturen, insbesondere der Particle Transformer (ParT), haben in Jet-Tagging-Aufgaben, wie der Unterscheidung hadronischer Top-Quark-Zerfälle von QCD-Hintergrundjets, State-of-the-Art-Leistung erzielt. Die internen Rechenmechanismen, durch die diese Modelle Klassifikationsentscheidungen treffen, bleiben jedoch intransparent. Während frühere Arbeiten post-hoc-Zuordnungsmethoden (z. B. Shapley-Werte, Saliency-Karten) oder Attention-Visualisierung genutzt haben, identifizieren diese Ansätze zwar, welche Eingaben relevant sind, erklären aber nicht, wie das Netzwerk sie kombiniert, und isolieren nicht das minimale kausale Subnetzwerk (Schaltung), das für das Verhalten verantwortlich ist. Ziel dieses Papers ist es, diese Lücke zu schließen, indem der gesamte Werkzeugkasten der mechanistischen Interpretierbarkeit – ursprünglich für natürliche Sprachmodelle entwickelt – auf einen Jet-Physik-Klassifikator angewendet wird.

Methodik
Die Autoren trainierten einen kleinen Particle Transformer (4 Particle-Attention-Schichten, 4 Köpfe pro Schicht, ~1,3 Mio. Parameter) auf einem Teilbereich des Referenzdatensatzes für Top-Quark-Tagging (Signal: $t \to Wb \to q\bar{q}b$; Hintergrund: leichte Quarks/Gluonen). Die Analyse verwendet eine Suite von Interventions- und Sondierungstechniken:

1. Zero-Ablation: Systematisches Setzen des Outputs einzelner Attention-Köpfe auf Null, um ihre strukturelle Wichtigkeit über den Abfall der mittleren Logit-Differenz zu messen.
2. Path Patching: Eine kausale Intervention, bei der der Output eines spezifischen Kopfs auf einem „sauberen" Input in einen „korrupten" Input substituiert wird (unter Verwendung von Partikel-Ersetzung innerhalb des Batches oder ganzer Jet-Permutation). Dies isoliert direkte Effekte und Pfadeffekte (Informationsfluss) zwischen den Köpfen.
3. Logit-Lens & Proben pro Schicht: Die Standard-Logit-Lens projiziert intermediate Repräsentationen durch den finalen trainierten Klassifikationskopf. Um Basis-Fehlausrichtung zu lösen, trainierten die Autoren zudem pro Schicht logistische Regressions-Sonden auf der Repräsentation jeder Schicht, um die wahre lineare Zugänglichkeit von Klasseninformationen zu bestimmen.
4. Lineare Proben: Training von Ridge-Regressionsmodellen zur Vorhersage klassischer Jet-Substruktur-Observablen (z. B. $N$-Subjettiness, Energy Correlation Functions) aus dem Residualstream in verschiedenen Tiefen, um den physikalischen Inhalt interner Repräsentationen zu charakterisieren.

Hauptbeiträge und Ergebnisse

• Identifikation einer spärlichen Sechs-Kopf-Schaltung: Durch Zero-Ablation und Path Patching identifizierten die Autoren eine minimale Schaltung aus sechs Attention-Köpfen, die 97,3 % der AUC des vollständigen Modells wiederherstellt. Diese Schaltung ist signifikant leistungsfähiger als zufällig gesampelte Sechs-Kopf-Teilmengen (liegt im 96. Perzentil einer zufälligen Baseline-Verteilung).
• Kausale Struktur (Quelle-Relais-Lesekopf): Die Schaltung zeigt eine klare kausale Hierarchie:
  • Primäre Quelle ($L0H1$): Ein einzelner Kopf in der ersten Particle-Attention-Schicht, der als primäre kausale Quelle fungiert. Er allein stellt 88,6 % der AUC des vollständigen Modells wieder her und zeigt „Super-Wiederherstellung" im Path Patching. Er attendiert bevorzugt weiche und kollineare Partikelpaare und liefert eine kontextuelle Normalisierung.
  • Sekundäre Quelle ($L0H2$): Ein Kopf derselben Schicht mit einem ähnlichen Repräsentationsraum wie $L0H1$, aber einer unterschiedlichen kausalen Rolle, der ein komplementäres Signal beisteuert.
  • Relais-Köpfe ($L1H0, L1H1, L1H3$): Ein Cluster von Köpfen in der zweiten Schicht, die selektiv auf harte paarweise Substrukturen (hohe invariante Masse, hohes $k_T$) attendieren. Ihre Funktion ist vom upstream-Signal der Quellköpfe abhängig.
  • Lesekopf ($L3H3$): Ein einzelner Kopf in der vierten Schicht, der die weitergeleiteten Signale aggregiert.
• Basis-Rotation vs. Informationsgewinn: Die Standard-Logit-Lens-Analyse deutete darauf hin, dass Klasseninformationen erst im ersten Klassen-Attention-Block ($Cls0$) entstehen, was einen Sprung von AUC 0,111 auf 0,973 zeigt. Jedoch zeigten pro Schicht trainierte Sonden, dass das klassen-diskriminierende Signal bereits in der ersten Particle-Attention-Schicht ($L1$) linear zugänglich ist mit AUC $\approx$ 0,97. Der dramatische Sprung bei $Cls0$ wird daher nicht als Generierung neuer Informationen interpretiert, sondern als Basis-Rotation, die das latente Signal an die Basis des finalen Klassifikationskopfs anpasst.
• Physikalischer Inhalt: 2-Prong vor 3-Prong Kodierung: Lineare Proben des Residualstreams gegen klassische Observablen zeigten, dass das Modell 2-Prong-Energie-Korrelator-Observablen (z. B. $D^{(\beta=1)}_2$) gegenüber 3-Prong-Observablen (z. B. $C^{(\beta=1)}_3, N^{(\beta=1)}_3$) bevorzugt kodiert, obwohl die Aufgabe das 3-Prong-Top-Tagging ist. Das Modell faktorisiert das Problem implizit, indem es sich auf die Identifikation des hadronischen $W$-Boson-Zerfalls (eine 2-Prong-Substruktur) konzentriert, anstatt auf die vollständige 3-Körper-Topologie. Diese Präferenz bleibt auch nach Residualisierung gegen die Jet-Masse bestehen.
• Energie-Korrelator-Basis: Der Residualstream ist stärker an die Energie-Korrelator-Basis als an die $N$-Subjettiness-Basis ausgerichtet, was darauf hindeutet, dass das Modell durch Gradientenabstieg ohne explizite Überwachung physikalisch sinnvolle Strukturen wiederentdeckt hat.

Bedeutung und Behauptungen
Das Paper behauptet zu demonstrieren, dass mechanistische Interpretierbarkeitsmethoden, die für natürliche Sprachmodelle entwickelt wurden, erfolgreich auf Jet-Physik-Klassifikatoren übertragen werden können. Die Ergebnisse deuten darauf hin, dass:

1. Gradientenabstieg physikalisch bedeutsame Aspekte des Jet-Taggings (speziell die Faktorisation des Top-Zerfalls in ein $W$-Boson-Teilproblem) ohne explizite physikalische Zwänge in der Architektur wiederentdecken kann.
2. Die hier identifizierte Quelle-Relais-Lesekopf-Schaltungsstruktur ein charakteristisches Muster für Physik-domänenspezifische Transformer sein könnte, das sich von der spezifischen Struktur von Sprachmodellen unterscheidet.
3. Die Studie die Notwendigkeit von On-Manifold-Korruptionsstrategien für Path Patching in kinematisch engen physikalischen Datensätzen unterstreicht, da Off-Manifold-(Gaußsche)-Korruption zu strukturellen Inkompatibilitäten mit Standard-Wiederherstellungs-Score-Formulierungen führen kann.

Die Autoren bleiben bezüglich der Universalität ihrer Ergebnisse bescheiden und weisen darauf hin, dass die spezifische Sechs-Kopf-Schaltung an einem kleinen Modell identifiziert wurde und größere Modelle reichhaltigere Schaltungsstrukturen besitzen könnten. Sie erkennen zudem an, dass lineare Sonden eine untere Schranke für den Informationsgehalt darstellen, da nicht-lineare Kodierungen nicht erfasst werden.