HEIMDALL: Disentangling tokenizer design for robust transfer in single-cell foundation models

Die Arbeit stellt HEIMDALL vor, ein Framework zur systematischen Analyse und Neugestaltung von Tokenisierungsstrategien für Single-Cell-Foundation-Modelle, das zeigt, dass die gezielte Anpassung spezifischer Designachsen wie Genidentität und Expressionskodierung entscheidend für eine robuste Generalisierung über Verteilungsverschiebungen hinweg ist.

Das Wesentliche

Das Problem: Der "Übersetzer" ist verwirrt

Stell dir vor, du hast einen genialen, super-intelligenten Roboter (ein KI-Modell), der gelernt hat, die Sprache der Zellen zu verstehen. Diese Zellen sind wie winzige Bibliotheken, in denen Tausende von Büchern (Genen) stehen. Jedes Buch hat einen Titel (den Gen-Namen) und einen Inhalt (wie aktiv das Gen ist).

Das Problem ist: Die KI kann keine Bücher direkt lesen. Sie braucht sie in einer bestimmten Form, wie eine Liste von Wörtern, die sie versteht. Dieser Prozess, bei dem man die Zelle in eine Liste für die KI umwandelt, nennt man Tokenisierung.

Bisher haben verschiedene Forschergruppen diese Umwandlung ganz unterschiedlich gemacht. Manche sortierten die Bücher nach dem alphabetischen Titel, andere nach der Farbe des Buchrücken, wieder andere nach dem Inhalt. Es gab keine einheitliche Regel. Das Ergebnis? Manchmal funktionierte die KI super, manchmal gar nicht. Niemand wusste genau, ob die KI dumm war oder ob der "Übersetzer" (der Tokenizer) sie einfach falsch bedient hatte.

Die Lösung: HEIMDALL – Der Werkzeugkasten

Die Forscher haben ein neues Werkzeug namens HEIMDALL entwickelt. Stell dir HEIMDALL wie einen riesigen, modularen Lego-Baukasten vor.

Statt ganze, fertige Roboter zu vergleichen, haben sie den "Übersetzer" in drei einfache Bausteine zerlegt:

1. Der Namensschilder-Macher (Gene Identity): Wie nennt die KI das Gen? (Nur der Name? Oder eine Beschreibung, was das Gen tut?)
2. Der Lautstärke-Messer (Expression Encoding): Wie laut ist das Gen? (Ist es leise oder schreit es?)
3. Der Regal-Einrichter (Cell Construction): In welcher Reihenfolge stellt die KI die Bücher ins Regal? (Nach Wichtigkeit? Nach Zufall? Nach Farbe?)

Mit HEIMDALL können die Forscher jetzt wie Wissenschaftler im Labor experimentieren: Sie nehmen den Baustein "Name" von Modell A, den "Lautstärke"-Baustein von Modell B und den "Reihenfolge"-Baustein von Modell C. Dann bauen sie einen neuen, perfekten Roboter zusammen und testen ihn.

Was haben sie herausgefunden?

Die Ergebnisse waren überraschend und wichtig:

1. Im "Heimatklima" ist es egal, aber im "Fremdland" zählt alles.
Wenn die KI Zellen aus dem Darm trainiert und dann wieder Zellen aus dem Darm erkennt, ist es fast egal, wie man die Bücher sortiert. Die KI schafft es trotzdem.
Aber: Wenn die KI Zellen aus dem Darm lernt und dann plötzlich Zellen aus dem Gehirn oder sogar von einer Maus sehen soll (das nennt man "Distribution Shift"), dann wird die Art und Weise, wie die Bücher sortiert sind, zum entscheidenden Faktor. Ein falscher Sortieralgorithmus lässt die KI komplett scheitern.

2. Es gibt nicht "den einen" perfekten Übersetzer.
Früher dachte man, es gäbe eine beste Methode, die immer gewinnt. HEIMDALL zeigt: Nein! Es kommt darauf an, wohin die KI reisen soll.

• Will sie von Mensch zu Maus wechseln? Dann braucht sie einen Übersetzer, der die Gene nach ihrer DNA-Sequenz (dem "Text" im Buch) erkennt, nicht nur nach dem Namen.
• Will sie neue Gene erkennen, die sie vorher nie sah? Dann hilft es, wenn die KI weiß, welche Gene oft zusammenarbeiten (wie Freunde, die immer in der gleichen Gruppe sind).

3. Die Mischung macht's.
Die Forscher haben gezeigt, dass man die besten Teile verschiedener Modelle kombinieren kann. Ein "Hybrid-Übersetzer", der die Stärken aller anderen vereint, ist oft besser als jedes einzelne Originalmodell.

Warum ist das wichtig?

Stell dir vor, du willst einen Arzt-Roboter bauen, der Krankheiten erkennt. Wenn du ihm die Patienten-Daten falsch formatierst, wird er den Krebs übersehen, auch wenn sein Gehirn (die KI-Architektur) das genialste der Welt ist.

Diese Arbeit sagt uns: Der Weg, wie wir Daten in die KI stecken, ist genauso wichtig wie die KI selbst. HEIMDALL gibt den Entwicklern jetzt eine Anleitung (eine Bauanleitung), wie sie ihre "Übersetzer" so bauen, dass die KI robust ist und auch in fremden Umgebungen (neue Gewebe, neue Spezies, neue Messmethoden) funktioniert.

Kurz gesagt: Die KI ist der Motor, aber HEIMDALL hat uns gezeigt, wie man den Kraftstoff (die Daten) richtig aufbereitet, damit der Motor auch auf unbekannten Straßen nicht stehen bleibt.

Technische Zusammenfassung

1. Problemstellung

Single-Cell Foundation Models (scFMs) versprechen, als allgemeine Werkzeuge für die biomedizinische Forschung zu dienen, indem sie auf Millionen von Einzelzell-Transkriptomprofilen vortrainiert werden. In der Praxis ist ihre Leistung jedoch inkonsistent, insbesondere bei Transfer-Learning-Aufgaben unter Distribution-Shift (z. B. Übertragung auf neue Gewebe, Spezies oder Gen-Panels).

Ein zentrales, aber bisher wenig untersuchtes Problem ist die Tokenisierung: Wie wird das hochdimensionale Genexpressionsprofil einer Zelle in eine Sequenz von „Tokens" umgewandelt, die als Eingabe für Transformer-Modelle dient?

• Im Gegensatz zu Text oder Bildern gibt es für Einzelzell-Daten keinen kanonischen Tokenisierungsstandard.
• Bestehende scFMs verwenden unterschiedliche, oft heuristische Schemata, die verschiedene biologische Annahmen kodieren.
• Da Tokenisierung, Architektur, Trainingsdaten und Modellgröße in bestehenden Benchmarks oft gleichzeitig variieren, ist es unmöglich, Leistungsunterschiede spezifisch auf die Tokenisierung zurückzuführen.
• Es fehlt an Prinzipien, um zu verstehen, welche Tokenisierungskomponenten unter realistischen Verschiebungen (z. B. Cross-Species) robust sind.

2. Methodik: Das HEIMDALL-Framework

Die Autoren stellen HEIMDALL vor, ein modulares Framework zur systematischen Zerlegung, Evaluierung und Neugestaltung von Tokenizern in scFMs.

Modulare Dekomposition:
HEIMDALL zerlegt jeden Tokenizer in drei funktionale Module, die unabhängig voneinander ausgetauscht werden können:

1. $F_G$ (Gene Identity Encoding): Kodiert die Identität eines Gens.
   • Varianten: Zufällige Initialisierung, vortrainierte Embeddings (ESM2 basierend auf Proteinsequenz, Gene2vec basierend auf Ko-Expression, GenePT basierend auf NLP-Textbeschreibungen, HyenaDNA basierend auf DNA-Sequenz).
2. $F_E$ (Expression Encoding): Kodiert den Expressionswert eines Gens innerhalb einer Zelle.
   • Varianten: No-Op (Ignorieren), kontinuierliche Abbildung (MLP), Quantile-Binning, Integer-Binning, Autobinning.
3. $F_C$ (Cell Construction / Aggregation): Integriert $F_G$ und $F_E$ zu einer Zellrepräsentation.
   • Unterteilt in ORDER (Reihenfolge der Gene: z. B. nach Expression, Chromosomenposition oder zufällig), SEQUENCE (Auswahl der Gene: z. B. Truncation oder gewichtetes Sampling) und REDUCE (Kombination von Identität und Expression: z. B. Summe oder Identity).

Experimentelles Design:

• Reimplementierung: Die Tokenizer von fünf führenden scFMs (scGPT, Geneformer, scFoundation, scBERT, UCE) wurden im HEIMDALL-Rahmenwerk nachgebaut.
• Kontrollierte Bedingungen: Um Tokenisierungseffekte isoliert zu messen, wurden alle Modelle von Grund auf neu trainiert (ohne Vortraining), mit fixierter Architektur (Transformer), Kontextlänge und Hyperparametern.
• Ablationsstudien: Durch das systematische Austauschen von Modulen (z. B. Ersetzen des $F_G$-Moduls eines Modells durch das eines anderen) konnte der Beitrag einzelner Design-Entscheidungen zur Generalisierungsfähigkeit quantifiziert werden.
• Benchmarks: Vier Transfer-Learning-Szenarien wurden getestet:
  1. Cross-Tissue (Gewebe-Shift).
  2. Cross-Species (Spezies-Shift: Mensch zu Maus).
  3. Cross-Gen-Panels (räumliche Transkriptomik mit unterschiedlichen Gen-Panels).
  4. Reverse Perturbation Prediction (Vorhersage von Störungen aus einem Zielzustand).

3. Wichtige Ergebnisse

A. Tokenisierung ist entscheidend bei Distribution-Shift, weniger bei In-Distribution

• In Szenarien mit übereinstimmenden Trainings- und Testdaten (In-Distribution) hatte die Wahl des Tokenizers nur minimale Auswirkungen auf die Leistung.
• Unter Distribution-Shift (neue Gewebe, Spezies, Gen-Panels) wurde die Tokenisierung zum entscheidenden Faktor. Modelle mit suboptimaler Tokenisierung scheiterten oft, während andere robust generalisierten.

B. Spezifische Erkenntnisse pro Benchmark

• Cross-Tissue: Die Leistung wurde primär durch die Reihenfolge (ORDER) bestimmt. Eine Sortierung nach Expressionsstärke (wie bei Geneformer) war der beste Ansatz. Änderungen an der Gen-Identität ($F_G$) hatten wenig Einfluss.
• Cross-Species: Hier war die Gen-Identität ($F_G$) kritisch.
  • Tokenizer, die Gen-Identitäten über Proteinsequenzen kodierten (ESM2, wie bei UCE), waren ohne weitere Anpassung robust gegenüber Spezies-Shift.
  • Tokenizer mit zufälligen oder artspezifischen Embeddings benötigten eine Orthologie-Mapping-Strategie, um vergleichbare Gene zwischen Mensch und Maus abzubilden.
  • Eine Kombination aus Orthologie-Mapping und artspezifischen $F_G$-Embeddings (ESM2) erzielte die besten Ergebnisse.
• Cross-Gen-Panels (Spatial Transcriptomics): Bei unvollständigen Gen-Panels war die Gen-Identität ($F_G$) der stärkste Treiber.
  • scBERT-tok (mit Gene2vec Embeddings, die auf Ko-Expressionsmustern basieren) schnitt am besten ab. Dies deutet darauf hin, dass Ko-Expressions-Priors die Repräsentation für Gene stabilisieren, die nur im Testset vorkommen.
  • Auch hier waren kontinuierliche Expressionskodierungen ($F_E$) und eine gute Sortierung ($ORDER$) vorteilhaft.
• Reverse Perturbation: Für die Vorhersage von Störungen waren Expressionskodierung ($F_E$) und Reihenfolge ($ORDER$) am wichtigsten.
  • Modelle, die keine expliziten Expressionsinformationen kodierten (wie UCE), verbesserten sich drastisch, sobald eine Expressionskodierung hinzugefügt wurde.
  • Die Kombination aus der $F_E$-Strategie von scBERT und der $ORDER$-Strategie von Geneformer erzielte die höchsten Leistungen.

C. Keine universell beste Tokenisierung
Es gibt keinen einzelnen „global optimalen" Tokenizer. Robuste Transferfähigkeit hängt von einer kleinen Anzahl von Design-Achsen ab, die unterschiedliche biologische Priors exponieren:

1. Gen-Identität (Sequenz vs. Ko-Expression vs. Text).
2. Expressionskodierung (Diskret vs. Kontinuierlich).
3. Reihenfolge (Sortiert vs. Zufällig).

4. Hauptbeiträge

1. HEIMDALL Framework: Ein einheitliches, modulares Interface zur Zerlegung von scFM-Tokenizern in interpretierbare Komponenten ($F_G, F_E, F_C$).
2. Systematische Abtrennung: Der Nachweis, dass Tokenisierung ein bisher unterschätzter, aber kritischer Designfaktor ist, der die Generalisierungsfähigkeit von scFMs unter realistischen Bedingungen maßgeblich bestimmt.
3. Design-Prinzipien: Identifikation spezifischer Tokenisierungskomponenten, die für verschiedene Shift-Szenarien (Gewebe, Spezies, Gen-Panels) optimal sind.
4. Hybride Tokenizer: Die Demonstration, dass hybride Tokenizer, die die besten Elemente verschiedener existierender Designs kombinieren (z. B. Gene2vec für Identität + Expressions-Sortierung), einzelne State-of-the-Art-Modelle übertreffen können.
5. Open Source: Bereitstellung des Codes als Python-Package (sc-heimdall) zur Reproduzierbarkeit und Weiterentwicklung.

5. Bedeutung und Ausblick

Die Arbeit etabliert die Tokenisierung als eine kritische Designachse für Single-Cell Foundation Models. Sie zeigt, dass die Robustheit von scFMs nicht nur von der Modellarchitektur oder der Datenmenge abhängt, sondern davon, wie biologische Informationen dem Modell präsentiert werden.

• Für Entwickler: HEIMDALL bietet eine Prinzipien-basierte Grundlage für das Design neuer Tokenizer, die auf spezifische Anwendungsfälle (z. B. räumliche Transkriptomik oder nicht-modell-Organismen) zugeschnitten sind, anstatt sich auf ad-hoc-Entscheidungen zu verlassen.
• Für Anwender: Es liefert Kriterien, um Modelle basierend auf dem erwarteten Einsatzszenario (z. B. Cross-Species vs. Intra-Species) auszuwählen.
• Zukunft: Das Framework legt den Grundstein für die Integration multimodaler Daten (Genomik, Epigenomik, Proteomik) in „virtuelle Zellen", indem es zeigt, wie verschiedene biologische Signale konsistent in Tokenisierungsschemata übersetzt werden können.

Zusammenfassend beweist HEIMDALL, dass die Entwicklung robuster scFMs eine bewusste und systematische Gestaltung der Schnittstelle zwischen biologischen Daten und dem Modell erfordert, wobei die Tokenisierung der Schlüssel zur universellen Übertragbarkeit ist.