Love Me, Love My Label: Rethinking the Role of Labels in Prompt Retrieval for Visual In-Context Learning

1. Problemstellung

Visuelles In-Context Learning (VICL) ermöglicht es visuellen Grundmodellen (Foundation Models), verschiedene Aufgaben durch die Verwendung von demonstrierenden Prompts (Beispielbildern) zu lösen. Die Leistung von VICL hängt maßgeblich von der Auswahl dieser Prompts ab.

• Aktueller Stand: Bisherige Forschungsarbeiten konzentrieren sich stark auf das Retrieval und die Neuordnung (Reranking) von Prompts basierend auf Bildähnlichkeit.
• Das Problem: Diese label-agnostischen Ansätze ignorieren oft die semantischen Labels der Prompts. Dies führt dazu, dass visuell ähnliche, aber label-inkonsistente Prompts ausgewählt werden (z. B. ein Bild einer Katze wird mit einem Prompt ausgewählt, der eine Katze zeigt, aber fälschlicherweise als „Blume" gelabelt ist).
• Konsequenz: Solche Inkonsistenzen verschlechtern die VICL-Leistung erheblich. Die Autoren zeigen empirisch, dass eine hohe Übereinstimmung zwischen dem Query-Label und dem Prompt-Label positiv mit der VICL-Performance korreliert.
• Herausforderung: Während die Prompts im Training sowohl Bild als auch Label haben, ist das Label des Query-Bildes zum Zeitpunkt der Inferenz (Testzeit) oft unbekannt.

2. Methodik: LaPR (Label-aware Prompt Retrieval)

Die Autoren schlagen LaPR vor, ein Framework, das Labels explizit in den Retrieval-Prozess integriert, um label-konsistente Prompts zu finden.

A. Label-Bewusste Repräsentation

Das Framework nutzt eine Mixture-of-Experts (MoE)-Architektur für beide Seiten (Query und Prompt), um verschiedene „Modi" (z. B. spezifische Merkmale wie „Horn", „Schnabel") zu erfassen.

1. Prompt-Seite (Label-Injektion):

   • Prompts bestehen aus einem Bild $I^p$ und einem Label $L^p$ (z. B. Segmentierungsmaske).
   • Beide werden durch einen Encoder verarbeitet und zu einer joint embedding (gemeinsamen Repräsentation) fusioniert: $z_i = f(I^p) + f(L^p)$.
   • Eine Menge von Experten $\bar{E}_k$ erzeugt modusspezifische Embeddings für diese joint Repräsentation.

2. Query-Seite (Schätzung des unbekannten Labels):

   • Da das Query-Label $y_q$ unbekannt ist, wird ein Router $R$ verwendet.
   • Der Router analysiert das Query-Bild $x_q$ und berechnet eine Wahrscheinlichkeitsverteilung $\pi_q$ über die $K$ Experten.
   • Diese Gewichte dienen dazu, die Experten-Ausgaben zu gewichten und ein label-aware Query-Embedding $\tilde{u}_q$ zu erzeugen, das implizit das wahrscheinlichste Query-Label widerspiegelt.

3. Retrieval-Prozess:

   • Die Query-relevanten Prompt-Embeddings werden durch Gewichtung der Prompt-Experten mit den Query-Router-Gewichten ($\pi_q$) extrahiert.
   • Die Ähnlichkeit wird zwischen dem label-aware Query-Embedding und den gewichteten Prompt-Embeddings berechnet (Cosine-Similarity).

B. Optimierungsstrategie (Alternating Optimization)

Da Experten und Router unterschiedliche Rollen spielen, wird ein entkoppelter Zwei-Schritt-Optimierungsansatz pro Mini-Batch verwendet:

1. Schritt 1: Experten-Optimierung (Performance-Guided):

   • Der Router wird eingefroren.
   • Die Experten werden mittels eines kontrastiven Verlusts optimiert, der auf der tatsächlichen VICL-Leistung (Task-Score) basiert.
   • Ziel: Experten lernen, modusspezifische Merkmale zu kodieren, die die Vorhersageleistung verbessern.

2. Schritt 2: Router-Optimierung (Label-Guided):

   • Die Experten werden eingefroren.
   • Der Router wird mittels eines label-geführten kontrastiven Verlusts optimiert. Hier werden Prompts mit hoher Label-Kompatibilität als Positiv-Beispiele genutzt.
   • Zusätzlich wird ein Load-Balancing-Loss hinzugefügt, um sicherzustellen, dass alle Experten genutzt werden und keine unterausgelastet sind.
   • Ziel: Der Router lernt, die Mischungsgewichte so anzupassen, dass sie mit dem Ground-Truth-Label des Queries übereinstimmen.

3. Hauptbeiträge

1. Erster label-bewusster Ansatz: LaPR ist der erste Ansatz, der Labels explizit für das Prompt-Retrieval in VICL nutzt und das Problem der Label-Inkonsistenz adressiert.
2. Architektur-Design: Einführung von joint Image-Label-Embeddings für Prompts und einer MoE-Architektur mit einem query-adaptiven Router für Queries, um implizite Labels zu schätzen.
3. Entkoppelte Optimierung: Ein innovatives Trainingsschema, das Experten und Router getrennt optimiert, um spezialisierte Repräsentationen und präzise Label-Schätzungen zu gewährleisten.
4. State-of-the-Art Ergebnisse: LaPR erzielt konsistent neue Bestwerte auf drei verschiedenen VICL-Aufgaben.

4. Ergebnisse

Die Methode wurde auf drei Aufgaben getestet: Foreground Segmentation (Pascal-5i), Single-Object Detection (Pascal VOC) und Image Colorization (ImageNet).

• Quantitative Leistung:
  • LaPR erreicht ohne Voting-Strategie die besten Ergebnisse aller retrieval-basierten Methoden.
  • Im Vergleich zum starken Reranking-Baseline (RH-Partial2Global) erzielt LaPR Verbesserungen von 6,00 % bei der Segmentierung und 3,46 % bei der Detektion.
  • Mit Voting-Strategie erreicht LaPR Segmentierungsergebnisse, die mit den besten Reranking-Methoden vergleichbar sind, und übertrifft diese bei der Detektion um 4,09 %.
• Cross-Fold Transferierbarkeit:
  • Beim Training auf einem Fold und Testen auf anderen (Cross-Fold) übertrifft LaPR bestehende Methoden (SupPR, Partial2Global) deutlich (z. B. +14,8 % gegenüber SupPR). Dies zeigt, dass die Experten feinere Modi als reine Kategorien lernen.
• Robustheit:
  • Die Leistung bleibt stabil, wenn der Feature-Extractor von CLIP auf DINOv2 gewechselt wird.
• Qualitative Analyse:
  • Visualisierungen zeigen, dass LaPR konsistent Prompts mit kompatiblen Labels auswählt, während label-agnostische Methoden oft falsche Kategorien (z. B. „TV" statt „Katze") auswählen, was zu schlechten Vorhersagen führt.

5. Bedeutung und Fazit

Das Paper unterstreicht, dass Labels ein entscheidender Hilfsfaktor für das Prompt-Retrieval in visuellen In-Context-Learning-Szenarien sind. Die reine Bildähnlichkeit reicht nicht aus; die semantische Konsistenz zwischen Query und Prompt ist essenziell für hohe Genauigkeit.

LaPR etabliert einen neuen Standard, indem es zeigt, wie man unbekannte Query-Labels durch adaptive Router-Mechanismen schätzen und diese Information nutzen kann, um relevantere Prompts zu finden. Die Arbeit liefert damit eine neue Grundlage für zukünftige Forschung im Bereich des label-bewussten Retrievals für Foundation Models. Der Code ist öffentlich verfügbar.