Die Infrastruktur hinter der KI-Suche in Figma

Auf der Config 2024 haben wir neue KI-gestützte Suchfunktionen vorgestellt, die es ermöglichen, alle Designs und veröffentlichten Komponenten in einem Team oder einer Organisation zu durchsuchen. Wir hatten unzählige Geschichten von Nutzer*innen gehört, die Schwierigkeiten hatten, ein bestimmtes Design oder eine Komponente zu finden, insbesondere in großen Organisationen mit komplexen Designsystemen. Dank der KI-gestützten Suche kannst du jetzt einfach mit einem Screenshot, einer Auswahl an Figma-Ebenen oder sogar einer textbasierten Beschreibung genau das finden, wonach du suchst.

Um diese Funktionen zu unterstützen, nutzen wir KI in zwei Suchabläufen: In der Suche nach Designs, einem völlig neuen Ablauf, und der Suche nach Komponenten, eine erheblich verbesserte Funktionalität für die Komponentensuche.

Suche nach Designs

Wir haben Frames in all deinen Dateien indiziert, damit du genau das findest, was du suchst… Selbst einen unbeschrifteten Frame in einer Datei mit Dutzenden von Seiten. Ab sofort kannst du ähnliche Designs entweder lexikalisch über eine textbasierte Beschreibung oder visuell über einen Screenshot oder eine Auswahl eines ähnlichen Designs finden.

Suche nach Komponenten

Wir haben die Funktionalität des Assets-Panels erweitert. Bisher basierte die Suche nach veröffentlichten Komponenten auf einer strengen Textübereinstimmung. Designer*innen mussten mögliche Schlüsselwörter manuell in Beschreibungen aufführen. Jetzt bietet die Assets-Suche ein KI-gestütztes semantisches Verständnis der Komponente. Beispielsweise kann eine 😀-Komponente über „Smiley“, „Happy“, „Face“ oder „Grinsen“ gefunden werden – egal, wie sie heißt. Keine manuelle SEO-Arbeit erforderlich! Ähnlich wie bei der Suche nach Designs kannst du auch visuell über einen Screenshot oder die Auswahl einer ähnlichen Komponente suchen.

Dieser Blick hinter die Kulissen zeigt, wie wir die Infrastruktur aufgebaut haben, um diese neuen KI-gestützten Suchfunktionen in Figma zu möglich zu machen.

Unser Einbettungsmodell

Das Herzstück der KI-gestützten Suche von Figma ist ein Einbettungsmodell, das Text oder ein Bild als aussagekräftige Zahlenmatrix darstellen kann. Zum Beispiel könnte das Bild einer Katze diese Einbettung erzeugen:

[0.066, 0.469, 0.103, 0.35, 0.382, 0.163, 0.587, 0.796, 0.311, 0.477]

Figma verwendet derzeit das Open-Source-Modell CLIP, ein sogenanntes multimodales Einbettungsmodell. Das Modell kann mehrere Formen von Eingaben (Bild und Text) verarbeiten und Einbettungen generieren, die sich im selben Bereich befinden. Dies bedeutet, dass eine Einbettung für die Zeichenfolge „Katze“ der obigen Einbettung numerisch ähnlich ist, obwohl die erste mit einem Bild als Eingabe generiert wurde.

Im Allgemeinen funktioniert unsere KI-gestützte Suche, indem sie einen Index mit Inhalten erstellt – zum Beispiel alle Komponenten in deinem Designsystem – und deren zugehörige Einbettungen. Anschließend wird eine Abfrage gestartet, um die Elemente zu finden, deren Embeddings der Abfrage am nächsten sind. Die Visualisierung einer Vektorsuche nach dem nächsten Nachbarn könnte folgendermaßen aussehen, allerdings mit viel mehr Dimensionen als den drei unten gezeigten:

Im Gegensatz zur traditionellen Suche, bei der die Abfrage direkt mit Einträgen im Suchindex verglichen wird, erfordert die Einbettungs-Suche zunächst die Generierung einer Einbettung für die Abfrage. Dieses wird anschließend mit einem Index von Einbettungen verglichen. Wenn du eine Suche mit einem Screenshot oder Text durchführst, generiert Figma im Handumdrehen eine Einbettung, die direkt in das Einbettungsmodell eingegeben werden kann. Suchst du über eine Auswahl von Figma-Layern, erstellt Figma einen neuen Screenshot deiner Auswahl und nutzt diesen als Eingabe für das Modell, um ein Query-Embedding zu generieren.

Aufbau des Vektorsuchindex

Sobald du dein Query-Embedding hast, weißt du, wonach du suchst, benötigst aber einen Index, um die Suche durchzuführen.

Identifizieren von Designs in Figma-Dateien

Um Designs zu finden, die tief in Figma-Dateien verborgen sein könnten, müssen wir alle Frames innerhalb einer Datei aufzählen, die durchsuchbar sein sollten. Für jeden dieser Frames müssen ein Thumbnail (ein gerenderter Screenshot der Figma-Layer) und eine zugehörige Einbettung erstellt und im Suchindex gespeichert werden. Das ist schwierig: Unveröffentlichte Frames in einer Figma-Datei können nicht ohne Weiteres gelistet werden. Um sie zu identifizieren, führen wir eine serverseitige Headless-Version des C++-Figma-Editors in einem asynchronen Job aus.

Weitere Informationen über unsere Sandboxing-Techniken zum Ausführen von serverseitigem C++ findest du hier.

Nachdem wir alle aktuellen indexierbaren Designs in einer Figma-Datei identifiziert haben, werden die Metadaten der Frames in DynamoDB gespeichert. Thumbnails werden gerendert und auf S3 hochgeladen. Der Prozess für die Identifizierung und das Thumbnailing startet den nächsten Schritt der Queue (Warteschlange) und wird erfolgreich abgeschlossen. Durch die Aufteilung der einzelnen Schritte der Pipeline in separate Jobs haben wir eine präzisere Kontrolle über Batch-Prozesse und Wiederholungslogiken.

Generieren von Einbettungen

Der nächste Schritt besteht darin, die Einbettungen zu generieren und zu speichern. Unser Einbettungsmodell wird in AWS SageMaker bereitgestellt. Einbettungsanforderungen werden in Batches an unseren SageMaker-Endpunkt gesendet, sodass wir Inferenzen parallel durchführen können. Die Eingabe ist eine Reihe von Thumbnail-URLs, die im vorherigen Schritt generiert wurden. Die Ausgabe ist eine Reihe von Einbettungen, eine für jedes Thumbnail.

Um Einbettungen effizient zu generieren, war es zum Zeitpunkt der Inferenz wichtig, sowohl das Herunterladen der Bilder als auch die Größenanpassung und Normalisierung der Bilder zu parallelisieren. Die Bestimmung der optimalen Stapelgröße erforderte einige Experimente. Ab einer bestimmten Schwelle stieg die Latenz linear mit der Batchgröße an, anstatt die Vorteile einer sublinearen Batch-Verarbeitung zu nutzen. Nachdem die Einbettungen erzeugt und gespeichert sind, wird der letzte Schritt der Pipeline gestartet.

Suchindizierung

Nach der Generierung der Einbettungen erfolgt die Speicherung im OpenSearch-Index, in dem die eigentlichen Vektorsuchanfragen ausgeführt werden. OpenSearch wird bei Figma bereits umfassend für herkömmliche Suchfunktionen eingesetzt. Daher war es sinnvoll, OpenSearch auch auch für die Einbettungs-Suche bei Figma zu nutzen.

Zusammen mit den Einbettungen werden zusätzliche Metadaten in den Suchindex geschrieben, wie z. B.: Frame-Name, Datei-ID und -Name, die den Frame enthalten, und zugehöriges Projekt, Team und Organisation. Diese zusätzlichen Metadaten ermöglichen facettierte Suchanfragen (Filter) zusätzlich zur Vektorsuche über die „Nearest Neighbor“-Logik.

Suche nach Komponenten

Jedes Mal, wenn eine Bibliothek veröffentlicht wird, wird ein asynchroner Job gestartet, um Einbettungen für jedes Thumbnail zu berechnen. Wir verwenden ein ähnliches Modell wie bei der Designs-Suche, allerdings ist es speziell auf öffentlich zugängliche Community-UI-Kits abgestimmt. Wie zuvor wurde dieses Modell nicht anhand privater Figma-Dateien oder Kundendaten trainiert.

Vor der KI-gestützten Suche gab es bereits die lexikalische Suche (mittels Fuzzy-String-Matching) über Komponentennamen und -beschreibungen. Um die KI-gestützte Suche sicher einzuführen und gleichzeitig die lexikalischen Ergebnisse beizubehalten, werden Suchen gleichzeitig sowohl im lexikalischen Index als auch im neuen Einbettungsindex durchgeführt. Da die Rohwerte aus den unabhängigen OpenSearch-Indizes nicht direkt vergleichbar sind, werden den Ergebnismengen auf der Grundlage einer Min-Max-Normalisierung neue Ergebnisse zugewiesen, wobei genaue lexikalische Übereinstimmungen einen Boost erhalten. Die Ergebnisse werden dann entsprechend ihrer aktualisierten Scores ineinander verschachtelt.

Jetzt liefert eine Abfrage nicht nur lexikalische Ergebnisse, sondern auch passende Ergebnisse basierend auf einem semantischen Verständnis. Beispielsweise gibt „Maus“ nicht nur ein Icon mit der Bezeichnung „Maus“ zurück, sondern auch Computermaus-/Cursor-ähnliche Icons.

Skalierungsbedingte Herausforderungen

Um eine KI-gestützte Suche in der Größenordnung von Figma zu ermöglichen, müssen Einbettungen für Milliarden von Einträgen erstellt und indiziert werden. Das ist nicht nur ein zeitaufwendiger, sondern auch ein kostenintensiver Prozess. Ein großer Teil des Projekts konzentrierte sich daher auf Optimierungen, um die Kosten zu senken.

Eine knifflige Herausforderung bei der Einführung dieser Funktionalität war, dass selbst für einen einzelnen Nutzer, der die Suchfunktionen wie vorgesehen nutzen möchte, alle Daten seines Teams indexiert sein müssen. Um die KI-gestützte Suche auch nur für eine kleine Anzahl von Nutzer*innen in der frühen Testphase verfügbar zu machen, mussten wir die Daten ganzer Teams indexieren. Paradoxerweise führte selbst eine kleine Anzahl aktiver Nutzer*innen schnell dazu, dass fast alle Teams bei Figma indexiert werden mussten – unsere Teams sind zwar klein, aber es gibt sehr viele von ihnen! Daher war es umso wichtiger, die Kosten für die Indexierung und das Backfilling niedrig zu halten.

Optimierung der Indizierung

Ein Blick auf die zuvor beschriebene Pipeline zeigte, dass die Hauptkosten nicht durch die Generierung von Einbettungen verursacht wurden, sondern durch die Identifikation und das Thumbnailing relevanter Designs in Figma-Dateien. Mit diesem Fokus haben wir mehrere Optimierungen vorgenommen:

• Ruby → C++ Unser anfänglicher Ansatz beruhte darauf, die gesamte Figma-Datei als JSON zu serialisieren und in Ruby zu parsen. Das war extrem langsam und speicherintensiv. Durch das Umschreiben dieser Logik von Ruby nach C++ und das Eliminieren jeglicher Zwischenserialisierung wurden enorme Laufzeitverbesserungen und Speicherreduktionen erzielt.
• Software-Rendering Wir haben die Erstellung unserer Thumbnails vom GPU-basierten Rendern auf einem älteren AWS-Instanztyp auf CPU-basiertes Rendern mit llvmpipe auf einem neueren Instanztyp umgestellt. Dies führte zu enormen Kosteneinsparungen. Die Maschinen vom Typ „CPU-Instanz“ sind viel günstiger. Und da sie neuer sind, bewältigen sie unsere Arbeitslast schneller.
• Verringerung der Indizierungsaktualität Ursprünglich wurde unsere Identifikations- und Indizierungspipeline bei jeder Änderung in einer Datei ausgelöst. Aber Figma-Nutzer*innen nehmen während einer längeren Bearbeitungssitzung oft viele Änderungen vorn. Und eine schnelle Aktualisierung des KI-gesteuerten Index ist nicht bei jeder einzelnen Änderung erforderlich. Trotzdem wollen wir unseren Index einigermaßen aktuell halten. Eine Analyse der Daten zeigte, dass wir durch eine Drosselung der Indexaktualisierungen auf maximal alle vier Stunden nur 12 % der Daten verarbeiten mussten!
• Automatische Clusterskalierung Die Nutzung von Figma ist stark tageszeitabhängig. Durch das Herunterskalieren unserer Cluster während Zeiten mit geringem Datenverkehr konnten wir unnötige Rechenkosten vermeiden.

Skalierung der Suchinfrastruktur

Reduzierung der Indexgröße

Ein weiterer großer Kostenfaktor war OpenSearch, das einen großen Cluster benötigte, um unsere umfangreichen Indizes im Speicher zu halten. Um dem entgegenzuwirken, haben wir genauer untersucht, welche Inhalte wir überhaupt indexieren sollten. Was zählt als relevantes Design? Am Ende entfernten wir Entwurfsdateien, doppelte Designs innerhalb von Dateien und Dateien, die als Kopien ohne neue Änderungen gespeichert waren. Mit dieser Maßnahme konnten wir den Index halbieren. Vieles davon waren letztendlich auch Produktverbesserungen: Das Vermeiden von Duplikaten innerhalb von Dateien führt zu einer deutlich besseren Nutzererfahrung.

Als nächstes nutzten wir die Vektorquantisierung. Standardmäßig stellt das kNN-plugin von OpenSearch jedes Element der Einbettung als Vier-Byte-Float dar. Mithilfe der Vektorquantisierung kann die Größe von Einbettungen jedoch komprimiert werden, um den Speicherbedarf zum Speichern und Durchsuchen zu reduzieren. Dies geschieht auf Kosten einer kleinen Verringerung der Suchgenauigkeit des nächsten Nachbarn.

Tücken bei OpenSearch

Während unserer Arbeit mit OpenSearch (sogar mit der neuesten in AWS unterstützten Version) stießen wir auf mehrere interessante Eigenheiten und Bugs im Zusammenhang mit der kNN-Suche.

• Bei End-to-End-Tests unserer Suchfunktionen stellten wir periodisch nicht-deterministische Ergebnisse fest. Nach einer ausführlichen Debugging-Session stellten wir fest, dass Anfragen, die zu Replikaten in OpenSearch geroutet wurden, andere Ergebnisse lieferten als Anfragen, die zu den Primärinstanzen gingen. Interessanterweise trat bei diesen Replikationsabfragen tief in OpenSearch ein Fehler auf (Reader cannot be cast to class SegmentReader).
  
  In enger Zusammenarbeit mit dem OpenSearch-Team von AWS gingen wir der Sache auf den Grund. Nach mehreren Anrufen fand das OpenSearch-Team die Ursache: Ein Typkonvertierungsfehler im Löschpfad beeinträchtigte Replikate auf Clustern, die Segmentreplikation nutzen. OpenSearch stellte hier eine Korrektur bereit.
• Um Speicherplatz und Abfragelatenz zu verbessern, entfernen wir den Einbettungsvektor selbst aus dem Feld _source im Suchindex. Das Feld _source ist der ursprüngliche Dokumenttext, der an den Suchindexer übergeben wird. Es ist nicht durchsuchbar, wird aber bei Antworten an den Client zurückgegeben – und Einbettungsvektoren sind groß!
  
  Es stellt sich jedoch heraus, dass OpenSearch bei Dokumentaktualisierungen auf _source angewiesen ist, um das aktualisierte Dokument zu diffundieren und zu schreiben, anstatt jeden Index an Ort und Stelle zu aktualisieren. Es greift auf die vorhandenen Felder aus _source zu, wendet die Änderungen an und schreibt das Dokument neu. Nachdem wir also die Einbettungen aus _source entfernt hatten, wurde jedes Mal, wenn wir versuchten, ein Dokument zu aktualisieren (z. B. wenn sich Dateinamen ändern), die Einbettung versehentlich aus dem Dokument gelöscht, da sie in _source nicht vorhanden war. Um dies zu beheben und gleichzeitig unsere _source-Optimierung beizubehalten, rufen wir bei Aktualisierungen die Einbettungen erneut aus DynamoDB ab.

Nächste Schritte

Diese KI-Suchfunktionen befinden sich derzeit in einer frühen Beta-Phase und werden in den kommenden Monaten schrittweise für weitere Nutzer*innen verfügbar gemacht. Wir freuen und auf euer Feedback und sind gespannt zu sehen, wie KI-gestützte Suche eure Workflows transformiert.