FM+RAG後日談:埋め込みベクトル化の精度改善

先日、Foundation Models で RAG を試みる内容を登壇したのだが、その時のスライドに添付したソースコードに誤りがあったので、以下ブログ記事に記載していたソースコードを修正した。

もともとはベクトル化対象のテキストを、トークン分割しつつ startIndex から endIndex まで手動で動かしながら畳み込みしていたものを、シンプルに enumerateTokenVectors(in:using:)  を使うようにしたら、↑記事で記載しているイマイチ精度が出ない問題を改善することができた。

以前の実装だと、何らかの条件で文字列最後までループが到達しないことが発生していたようだ。文頭の構文しかヒットしないという現象も、この原因を考えれば納得できる。


そもそも、ここで紹介している NLContextualEmbedding + mean pooling + L2 normalization で埋め込みベクトル化し、コサイン類似度を求める手法は、すでに以下のQiita記事で同じことが解説されていた。今後実装される方はこっちを参考にした方が幸せかもしれない。(もっと早く見つけたかった、、)

iOSに組み込まれたBERTでテキスト埋め込み・ベクトル検索をオンデバイス実行する #Mac – Qiita

登壇メモ:extension DC 2025 Day1 @DeNA

久々に登壇してきたので記録。

イベントページ:extension DC 2025 Day1@DeNA

夏から取り組み始めていたFoundation Models + RAG の集大成?を発表。結果的にFM側の挙動で綺麗な結果にはならなかったが、、RAGの一翼を担う自前の検索エンジンとしてはきちんと良い結果が出たので、その実装方法を中心にシェアした。

スライドは40枚作っていたが、何度練習しても5分に収まるかは一か八かだったので、会場でトピック丸ごと(この記事の内容)省略した。そのおかげで早口ながら完走はできたのでよかった。アップロードしたスライドには、スキップした内容も復活しておいた。

Apple で開催された Foundation Models のワークショップでたくさんサポートくださった武石さんともお会いでき、FMの挙動について具体で相談させていただき追加でアドバイス頂けたので、試してみたい。

参加メモ:新しいFoundation Modelフレームワークのアプリへの導入(ワークショップ)@Apple Japan


発表後、さまざまな方にお声がけいただき、中にはこのセッションのために来たとおっしゃってくれる方や、今回発表のアプローチを自社プロダクトへ実装検討されている方も何名かいらっしゃって、今回の内容が少しでもお役に立てれば何より。

発表内で紹介した検証は、パラグラフにも満たない短文と、30という限られたドキュメント量でしか試していないので、実運用するデータ規模によっては性能限界があるかもしれない。今回触れなかった文章の細分化や、プーリングのアルゴリズムを変更するなどチューニングの余地は多く残されている。

まだ道半ばなので、今後も試行錯誤を続けていくがその過程は都度「Foundation Models」でタグ付けしていく。

https://p0dee.com/blog/tag/foundation-models

最後に、ここまでの軌跡において救世主となった武石さんのポストに改めて感謝!


会場からの帰り道になぜかYouTuberに捕まって、「あなたの人生を語ってください」的なよくある企画に巻き込まれ、日が変わるまで沖縄料理屋で飲んでた。

SwiftでRAG実装 Part 2:クエリに類似するドキュメント検索の試み

SwiftでRAG実装 Part 1:テキストをコンテキストベクトル化 の続き。

前回はドキュメントデータから埋め込み生成。今回はそこから検索クエリに類似するドキュメントを抽出してみる。先に結論を言うと、結果はいまひとつなので調整が必要そう。

今回実装する計算ロジックは以下(現状の筆者の理解であることに注意)。

  1. ドキュメントからそれぞれ埋め込み生成し、ベクトル化する
  2. ドキュメントごとのベクトル表現(D次元)を、ドキュメントの数(M個)分並べて MxD 行列を生成する(ドキュメント行列)
  3. クエリ文字列をベクトル化する。1. と同じ手法で生成しベクトル次元は一致(D次元)。(クエリベクトル)
  4. ドキュメント行列・クエリベクトル(内積)を計算する (MxD・Dx1 = Mx1)
    • つまりドキュメントごとのベクトルとクエリベクトルとの内積が結果として求まる。内積ベクトルのノルム(コサイン類似度)が1に近いほど類似度高く、0、-1に近づくほど類似度が低いと判定できる
  5. 4. の結果で得られた類似度の配列(要素数M)を降順にソート、上から任意の件数を上位ヒットとする
var docTensors: [MLTensor] = ...

// docTensors に含むドキュメントごとのクエリとの類似度を検索し、上位任意件数(maxCount)をヒット結果として返す
// (返すのは該当文書テンソルの docTensors 配列内におけるインデックス番号)
func search(query: String, maxCount: Int) async -> [Int] {
    // ドキュメントを集積した MxD 行列
    let flatteneds = tensors.map { $0.flattened() }
    let docsMat = MLTensor(stacking: flatteneds)
    // クエリベクトル Mx1 (encode関数は前回記事参照)
    guard let queryVec = try? embedder.encode(text: query, asColumn: true /*列ベクトルで出力*/) else { return [] }
    
    // ドキュメント行列・クエリベクトル (MxD・Dx1 = Mx1)
    let mulResult = docsMat.matmul(queryTensor)
    // 類似度スコア配列に変換
    let calcScores = await mulResult.shapedArray(of: Float.self).scalars
    // ソートして、トップの結果をインデックス番号として抽出
    let map = zip(Array(0..<docTensors.count), calcScores)
    let sorted = map.sorted { $0.1 > $1.1 }.prefix(maxCount)
    return sorted.map { $0.0 }
}

これでサンプルの日記データを検索したところ、肌感としてあまり精度高くなかった。

ただし、クエリの何かしらの特徴は反映していそうだった。たとえば「I watched a movie」のクエリで検索すると、映画を見た日記エントリはヒットしないが、上位ヒットは軒並み「I (動詞過去形)」のパターンで始まっていたり、更に分かりやすい例を挙げると、次のように「at a cafe」クエリに対し「At a/the」で始まるエントリが抽出されたりした。まったくの出鱈目ではなさそうだが、肝心の文章的なコンテキストは明らかに落ちているように見える。

Screenshot

追記(2025/10/01):後日性能改善できたので後々投稿するが、取り急ぎこちらの資料にソースコードを添付した。(enumerateTokenVectors を使うようにしただけ)

SwiftでRAG実装 Part 1:テキストをコンテキストベクトル化

Foundation Models のツール呼び出しに使う RAG に Spotlight を利用する試みが絶たれたので、自前で RAG 実装することにした。(本記事の作業は↓のワークショップで取り組んだもの)

経緯:参加メモ:新しいFoundation Modelフレームワークのアプリへの導入(ワークショップ)@Apple Japan

このブログでは繰り返し何度も貼っているが、Apple 武石さんのこのポストを頼みの綱に頑張る。もはやバイブル。


といってもちんぷんかんぷんなので ChatGPT に聞いてみた。

NLContextualEmbedding を使ってRAGを構築するにはどう言った実装が必要になりますか?

NLContextualEmbedding を使った “自前RAG” の実装は、だいたい次の4層で考えると設計しやすいです。

  1. 前処理(分割・メタデータ付与)
  2. 埋め込み生成(NLContextualEmbedding)
  3. ベクタ格納と近傍検索(SQLite + BLOB / CoreData / なんでも)
  4. 検索→再ランキング→プロンプト組み立て→生成(Foundation Models)

いったん前処理は省いて、まずは短い文字列を元にした埋め込み生成(コンテキストベクトル化)を試みた。ChatGPT の実装をもとに動くよう修正し、最終的にそれっぽくベクトルを得ることができた。(コメントは筆者理解の補足なので間違いあるかも)(2025/10/5 実装修正)

import Accelerate
import NaturalLanguage

let embedding: NLContextualEmbedding

/// 文字列をコンテキストベクトル化(平均プーリング+L2正規化)
func encode(text: String) throws -> [Float] {
    let result = try embedding.embeddingResult(for: text, language: overrideLanguage)

    let dim = embedding.dimension
    var sum = [Float](repeating: 0, count: dim)
    var count = 0 // トークン数

    result.enumerateTokenVectors(in: text.startIndex..<text.endIndex) { vec_double, range in
        var vec_float = [Float](repeating: 0, count: dim)
        // double → float 変換
        vDSP_vdpsp(vec_double, 1, &vec_float, 1, vDSP_Length(dim))
        // vec_float を sum に足し合わせ
        vDSP_vadd(sum, 1, vec_float, 1, &sum, 1, vDSP_Length(dim))
        // トークン数をインクリメント
        count += 1
        return true
    }

    guard count > 0 else { return [Float](repeating: 0, count: dim) }

    // 平均プーリング(トークンベクトルの総和をトークン数で平均して畳み込み)
    var inv_n = 1.0 / Float(count)
    vDSP_vsmul(sum, 1, &inv_n, &sum, 1, vDSP_Length(dim))

    return l2Normalize(sum) // L2 normalization
}

// L2 正規化(ベクトル全体を二乗和平方根で割って正規化)
private func l2Normalize(_ v: [Float]) -> [Float] {
    var vec = v
    var norm: Float = 0
    vDSP_svesq(vec, 1, &norm, vDSP_Length(vec.count))
    norm = sqrtf(norm) + 1e-12
    vDSP_vsdiv(vec, 1, &norm, &vec, 1, vDSP_Length(vec.count))
    return vec
}
print(encode(text: "Hello, world.")

[-0.028537132, -0.014218736, -0.033890422, -0.024530113, 0.009770119, -0.01361734, 0.0034657633, 0.029605899, 0.013323085, -0.005046577, ..., 0.018509272, -0.026693422, -0.6423329, -0.03437927, 0.005926335, -0.022124525, 0.03561643, -0.056179043, 0.025543895, -0.00908023, 0.0050482955, 0.028503625]

ちなみに、今回 “Hello, world.” は Hello,, , world. の3トークンに分割された。


埋め込み生成の処理をおさらいすると

  1. 文字列をもとにベクトル埋め込みを生成(NaturalLanguage.NLContextualEmbedding.embeddingResult(for:language:)
  2. 文字列をトークン(サブテキスト単位)に分割
  3. トークンごとに、トークンベクトルを抽出(NaturalLanguage.NLContextualEmbeddingResult.tokenVector(at:)
  4. すべてのトークンベクトルの平均を計算 → コンテクストベクトル
  5. コンテクストベクトルを二乗和平方根で割って正規化

これ書きながら、2-3 のステップで頑張ってループ回しているところは enumerateTokenVectors(in:using:) 使ったほう良いかも、と思った。

参考:
埋め込み層 (Embedding Layer) [自然言語処理の文脈で] | CVMLエキスパートガイド
平均プーリング(Average Pooling) | CVMLエキスパートガイド
[iOS 17] 多言語BERT埋め込みモデルのサポート


Accelerate framework に馴染みがないので、ここで使われている関数を調べてみた。

  • vDSP_vdpsp: 倍精度のベクトルを単精度のベクトルに変換
  • vDSP_vadd: ベクトル同士の和 (stride 指定可能)
    • stride: 足し合わせ時の要素飛び石数 通常は1だが、オーディオバッファからLRチャンネルを分離して取得する時(stride=2)や、イメージバッファからRGBチャンネルを分離して取得する(stride=3)時に指定
  • vDSP_vsmul: ベクトル * スカラ値 の積算 (stride 指定可能)
  • vDSP_vsdiv: ベクトル / スカラ値 の除算 (stride 指定可能)
  • vDSP_svesq: ベクトルの二乗和 sum(a^2) 、結果はスカラ値