FigmaのAI検索を支えるインフラ

Config 2024では、AIを活用した新しい検索機能を紹介し、チームや組織内のすべてのデザインと公開されているコンポーネントを横断的に検索できるようにしました。複雑なデザインシステムを持つ大規模な組織で作業している場合は特に、特定のデザインやコンポーネントを見つけるのに苦労しているというユーザーの話を数え切れないほど聞いてきたからです。AIを活用した検索機能により、スクリーンショットやFigmaレイヤーの選択、テキストによる説明だけで探しているものを見つけることができます。

これらの機能を強化するため、Figmaでは、まったく新しいフローである「デザインの検索」とコンポーネントの検索機能を大幅に向上させる「コンポーネントの検索」という2つの検索フローでAIを活用しています。

デザインの検索

すべてのファイルのフレームにインデックスが付けられているので、数十ページのファイルに紛れ込んだラベルのないフレームであっても、確実に見つけることができます。テキストによる説明で類似のデザインを語彙的に見つけたり、スクリーンショットや類似のデザインを選択して視覚的に見つけたりできるようになりました。

コンポーネントの検索

アセットパネルの機能を強化しました。これまでは、公開されているコンポーネントの横断的な検索に厳密なテキストマッチングを使用していたため、デザイナーは可能性のあるキーワードを説明文に手作業で列挙する必要がありました。現在のアセット検索では、AIを活用してコンポーネントを意味的に理解できるようになりました。たとえば、😀 コンポーネントは、その名前に関係なく、「スマイリー」、「ハッピー」、「顔」、「ニヤリ」で見つけることができ、手動のSEO作業は必要ありません。また、デザインの検索と同様に、スクリーンショットや類似コンポーネントの選択によって視覚的に検索することも可能です。

ここでは、Figmaの新しいAI検索機能を実現するため、私たちがどのようにインフラを構築したのか、その舞台裏をご紹介します。

当社の埋め込みモデル

FigmaのAIを活用した検索の中心となるのは、テキストや画像を意味のある数字の配列として表現できる埋め込みモデルです。例えば、猫の画像では以下のような埋め込みが生成される可能性があります。

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

Figmaは現在、マルチモーダル埋め込みモデルと呼ばれるオープンソースのCLIPモデルを使用しています。このモデルでは、複数の形式の入力(画像とテキスト)を受け取り、同じ空間にある埋め込みを出力することができます。つまり、文字列「cat」の埋め込みは、画像を入力として最初の埋め込みが生成された場合でも、数値的には上記の埋め込みとほぼ同じになります。

大まかに言うと、AIを活用した検索は、コンテンツ(設計システム内のすべてのコンポーネントなど)とそれに関連する埋め込みのインデックスを作成し、クエリを実行して埋め込みがクエリに最も近い項目を見つけるというものです。最近傍ベクトル検索を視覚化すると以下のようになりますが、次元は以下の3つよりもはるかに多くなります。

ユーザーのクエリを検索インデックスのエントリと直接比較する従来の検索とは異なり、埋め込み検索では、まずユーザーのクエリの埋め込みを生成し、それを埋め込みのインデックスと比較する必要があります。スクリーンショットまたはテキストを使用して検索を実行すると、Figmaは埋め込みモデルに直接入力できる埋め込みをその場で生成します。Figmaレイヤーを選択して検索する場合は、選択したレイヤーのクリーンショットを新たに生成し、それをモデルへの入力として使用してクエリの埋め込みを生成します。

ベクトル検索インデックスの取り込み

クエリの埋め込みが生成され、探しているものがわかるようになりましたが、検索を実行するにはインデックスを取り込む必要があります。

Figmaファイル内のデザインの識別

Figmaファイルの奥深くに眠るデザインを見つけるには、検索可能なFigmaファイル内の全フレームを列挙できる必要があります。また、それぞれについてサムネイル(Figmaレイヤーのレンダリングされたスクリーンショット)と埋め込みを生成し、検索インデックスに書き込む必要があります。ここが難しいところです。Figmaファイル内の未公開フレームは簡単に列挙できません。これらを識別するには、非同期ジョブでC++ Figmaエディターのヘッドレスサーバー側バージョンを実行します。

サーバーサイドのC++を実行するためのサンドボックス手法について詳しくは、こちらをご覧ください。

現在インデックス可能なFigmaファイル内のデザインをすべて特定したら、フレームのメタデータがDynamoDBに保持されます。また、サムネイルがレンダリングされ、S3にアップロードされた後、この識別およびサムネイル作成ジョブは、次のジョブをキューに入れて正常に終了します。パイプラインの個々のステップを個別のジョブに分割することで、バッチ処理と再試行の動作をより正確に制御できます。

埋め込みの生成

次に、実際に埋め込みを生成し、それを保持する必要があります。当社の埋め込みモデルはAWS SageMakerにデプロイされています。埋め込みリクエストはSageMakerエンドポイントにバッチで送信されるため、推論を並行して実行できます。入力は、前のステップで生成された一連のサムネイルURLであり、出力は入力画像ごとに1つずつ用意された、一連の埋め込みです。

埋め込みを効率的に生成するには、推論時に、画像のダウンロードに加え、画像のサイズ変更と正規化の両方を並行処理することが重要でした。最適なバッチサイズを決定するには、ある程度の実験が必要でした。そして、この実験により、しきい値を超えると、バッチエフェクトが遅れて現れるのではなく、バッチサイズとともにレイテンシが直線的に増加することがわかったのです。埋め込みが生成され、保持されると、パイプラインの最後のステップがキューに入れられます。

検索インデックスの作成

埋め込みが生成されたら、OpenSearchインデックスに埋め込みを書き込んで、実際のベクトル検索クエリを実行します。OpenSearchは従来の検索機能としてFigma全体に広く導入されており、Figmaでの埋め込み検索においてOpenSearchを活用するのは理にかなっています。

埋め込みは、フレームの名前、含まれているファイルのIDと名前、フレームを含むプロジェクト、チーム、組織などの追加メタデータとともに検索インデックスに書き込まれます。この追加メタデータにより、最近傍ベクトル検索に加えてファセット検索(フィルター)にも対応できます。

コンポーネントの検索

ライブラリが公開されるたびに、各サムネイルの埋め込みを計算するための非同期ジョブが開始されます。当社では、デザインの検索に使用されるものと非常によく似たモデルを使用しており、これは、公開されているコミュニティUIキットに特化して微調整されています。これまでと同様、このモデルはプライベートなFigmaファイルや顧客データに基づいてトレーニングされていません。

コンポーネント名と説明による語彙的な検索(あいまい検索)は、AIによる検索よりも以前から行われていました。AIを活用した検索をコンポーネントに安全に展開し、貴重な語彙的結果を保持するため、語彙インデックスと新しい埋め込みインデックスの両方に対して同時に検索が実行されます。独立したOpenSearchインデックスから得た未加工のスコアは直接比較できないため、結果セットでは正規化に基づいて新しいスコアが割り当てられ、正確な語彙の一致が強調されます。その後、結果は更新されたスコアに従ってインターリーブされます。

これで、クエリは語彙的な結果だけでなく、意味的な理解に基づいた適切な結果も返すようになりました。たとえば「マウス」の場合、「マウス」というタイトルのアイコンだけでなく、カーソルに隣接するアイコンも返されます。

大規模展開における課題

Figmaの規模でAIを活用した検索を可能にするには、何十億もの埋め込みを生成し、インデックスを作成する必要があります。これには時間だけでなく、費用もかかるため、プロジェクトの大部分は、コストを抑えるための最適化に重点が置かれていました。

この展開において注意が必要なのは、意図したとおりの検索機能をユーザー1人が体験するだけでも、チームのすべてのデータをインデックス化する必要があるということです。つまり、AIを活用した検索を少数のユーザーに提供し、早期にテストしてもらうためには、全チームのデータをインデックス化する必要があるのです。当社のチームのほとんどは小規模ですが、その数は膨大。ごく一部のユーザーが参加するだけであったとしても、Figmaのほぼすべてのチームをインデックス化する必要があります。このため、インデックスの作成とバックフィルにおけるコスト効率を維持することがさらに重要になりました。

インデックス作成の最適化

以前説明したパイプラインを見てみると、主要な計算コストは、埋め込みの生成ではなく、Figmaファイル内で有意なデザインを特定し、サムネイル化することによってもたらされていることがわかりました。当社ではそこに焦点を当て、以下のような最適化を実行しました。

• Ruby → C++: 当初のアプローチでは、Figmaファイル全体をJSONとしてシリアル化し、それをRubyで解析していましたが、これは非常に時間がかかり、メモリを大量に消費する作業でした。このロジックをRubyからC++に書き直し、間のシリアル化をなくすことで、作業時間と消費メモリを大幅に削減できました。
• ソフトウェアのレンダリング: サムネイルを、古いAWSインスタンスでのGPUベースのレンダリングから、新しいインスタンスタイプでのllvmpipeによるCPUベースのレンダリングに移行し、大幅なコスト削減を実現しました。CPUインスタンスタイプのマシンははるかに安価で新しく、作業をスピードアップできます。
• インデックスの鮮度を落とす: 当社の識別とインデックス作成パイプラインは、もともとファイル内の変更によってトリガーされるように作成されていました。しかしながら、Figmaユーザーがファイルを変更する場合、より長い編集セッションで多くの変更を加えることがよくあります。AIを活用したインデックスを変更のたびに素早く更新する必要はありませんが、インデックスを適度に最新の状態に保ちたいというのが本音です。データを確認したところ、インデックスのデバウンスを最大4時間ごとに行えば、処理が必要なデータは12%のみになることがわかりました。
• クラスターのオートスケール: Figmaは日中に使用されることが多いです。トラフィックが少ない時間にクラスターをスケールダウンすることで、不要な計算にお金を払う必要がなくなりました。

検索クラスターのスケーリング

インデックスサイズの縮小

2番目にお金がかかっていたのがOpenSearch。膨大なインデックスをメモリ内に維持するために大規模なクラスターが必要でした。これに対処するため、当社ではまず、インデックス化できるものを詳しく調べました。有意義なデザインとは何なのかに焦点を当て、ドラフトファイル、ファイル内の重複したデザイン、新たな変更なしでコピーされたファイルを削除したことで、インデックスが半分に削減されました。こういった努力によってファイル内に重複したデザインが表示されないようになり、ユーザー体験が向上したことで、製品の改善にもつながりました。

次に行ったのはベクトルの量子化です。OpenSearchのkNNプラグインは、デフォルトで埋め込みの各要素を4バイトの浮動小数点数として表現します。ベクトル量子化は、埋め込みのサイズを圧縮して保存と検索に必要なメモリを削減する手法ですが、その代わりに最近傍検索の精度がわずかに低下します。

OpenSearchのクセ

AWSでサポートされている最新リリースを実行しているにもかかわらず、OpenSearchのkNN検索では、興味深いクセやバグがいくつか見つかりました。

• 当社の検索機能をエンドツーエンドでテストしたところ、周期的に不確定性が生じることに気付きました。広範囲にわたるデバッグセッションの結果、OpenSearchのレプリカにルーティングされたクエリは、プライマリにルーティングされたクエリと比較して非決定的な結果を返していることが判明しました。興味深いことに、OpenSearchの奥深くで、これらのレプリカクエリにエラー(リーダーをSegmentReaderクラスのキャストできません)が発生していたのです。
  
  当社はAWSのOpenSearchチームと緊密に連携し、調査を行いました。何度も話し合った結果、OpenSearchチームは、削除パスで発生したこの種のキャスティングエラーがセグメントレプリケーションを利用しているクラスターのレプリカに影響を及ぼしていることを突き止め、ここで修正を提供しました。
• ストレージ容量とクエリのレイテンシを改善するため、検索インデックスの_sourceから埋め込みベクトル自体を削除します。_sourceフィールドは、検索インデクサーに渡された元のドキュメントの本文です。検索はできませんが、応答時にクライアントに返され、その埋め込みベクトルは膨大です。
  
  しかし、ドキュメントを更新する際、OpenSearchは各インデックスをその場で更新するのではなく、_sourceを使用して更新されたドキュメントを差分して書き込むことが分かりました。これは、_sourceから既存のフィールドを取得し、更新を適用して、ドキュメントを書き込むということであり、_sourceから埋め込みを削除すると、ファイル名が変更されたときなどにドキュメントを更新しようとするたびに、_sourceに埋め込みが存在しないため、誤ってドキュメントからも埋め込みを消去していました。この問題を解決するため、当社では、_sourceの最適化を維持したまま、更新時にDynamoDBから埋め込みデータを再取得します。

次のステップ

これらのAI検索機能は現在初期ベータ版で提供しており、今後数か月にわたってさらに多くのユーザーに展開される予定です。当社では、AIを活用した検索が皆様のワークフローをどのように変革するかを楽しみにしており、皆様からのフィードバックをお待ちしております。