クラウドMacでiOSベクトル検索の再現性を検証

同じ検索語なのに、昨日は文書Aが先頭に出ていたのに、今日は文書Bが上位に来る。そんなときも、モデルを再学習する前に確認したいことがあります。候補の入力順、ゼロベクトル、浮動小数点数のオーバーフローによって、iOSのセマンティック検索で回帰が起きたように見える場合があるためです。固定ベクトルをクラウドMac上で実行し、まず検索層の受け入れテストを済ませてからモデル出力を調べれば、原因の範囲を大きく絞り込めます。

検索の契約を固定してからモデル出力につなぐ

今回の検証対象は、メモリ上で行う厳密なコサイン類似度検索だけです。近似インデックスは扱いません。生成済みのベクトルを入力とし、文書の識別子とスコアを出力します。モデル推論は別のテスト群に分け、1つの失敗で両方の実装を同時に疑う状況を避けます。

契約として、次の5つのルールを定めます。ベクトルは同じ次元数で空ではないこと、全要素が有限値であること、ゼロベクトルは拒否すること、スコアの降順かつ同点なら一意な識別子の昇順で並べること、しきい値と等しいスコアも含めることです。k がゼロなら空の結果を返し、負数ならエラーにします。候補数を超える場合は、条件を満たすすべての候補を返します。

ゼロベクトルは「クエリと無関係」という意味ではありません。方向を定義できないため、コサイン類似度がゼロになる直交ベクトルと同じ扱いにはできません。

この例では、入力の次元数の上限を16384に設定しています。これはローカル環境のリソース予算に基づく制限であり、モデルの仕様ではありません。実際のプロジェクトでは、使用するモデルに合わせて上限を絞り、ファイル読み込み時のバイト数にも制限を設けてください。

スケーリングによる正規化で数値オーバーフローを防ぐ

二乗和をそのまま計算すると、各要素が有限値でも、途中の計算結果が無限大になることがあります。先に絶対値の最大値で割ってからノルムを求めれば、大きな値を二乗する際のオーバーフローを避けられます。分母に安易に定数を足してはいけません。ノルムの小さいベクトルでは、方向を表す単位ベクトルへの正規化結果が変わってしまいます。

以下のコード全体を vector-contract.swift として保存してください。ツールチェーン付属のモジュールだけを使用しているため、クラウドMacのコマンドラインで実行できます。

import Foundation

enum VectorError: Error {
    case invalid
}

func unit(_ values: [Double]) throws -> [Double] {
    guard !values.isEmpty,
          values.count <= 16384,
          values.allSatisfy({ $0.isFinite }),
          let scale = values.map({ abs($0) }).max(),
          scale > 0 else {
        throw VectorError.invalid
    }
    let scaled = values.map { $0 / scale }
    let length = sqrt(scaled.reduce(0) { $0 + $1 * $1 })
    return scaled.map { $0 / length }
}

struct Candidate {
    let id: String
    let vector: [Double]
}

struct Hit {
    let id: String
    let score: Double
}

func search(
    query: [Double],
    candidates: [Candidate],
    k: Int,
    threshold: Double
) throws -> [Hit] {
    guard k >= 0, threshold.isFinite,
          (-1.0...1.0).contains(threshold),
          Set(candidates.map(\.id)).count == candidates.count,
          candidates.allSatisfy({ !$0.id.isEmpty }) else {
        throw VectorError.invalid
    }
    let q = try unit(query)
    var hits: [Hit] = []
    for candidate in candidates {
        guard candidate.vector.count == q.count else {
            throw VectorError.invalid
        }
        let v = try unit(candidate.vector)
        let raw = zip(q, v).reduce(0.0) {
            $0 + $1.0 * $1.1
        }
        guard raw.isFinite, abs(raw) <= 1 + 1e-12 else {
            throw VectorError.invalid
        }
        let score = min(1, max(-1, raw))
        if score >= threshold {
            hits.append(Hit(id: candidate.id, score: score))
        }
    }
    hits.sort {
        if $0.score == $1.score { return $0.id < $1.id }
        return $0.score > $1.score
    }
    return Array(hits.prefix(k))
}

let candidates = [
    Candidate(id: "b", vector: [2, 0]),
    Candidate(id: "a", vector: [1, 0]),
    Candidate(id: "c", vector: [0, 1]),
    Candidate(id: "d", vector: [-1, 0])
]
let hits = try search(
    query: [1, 0], candidates: candidates,
    k: 3, threshold: 0
)
precondition(hits.map(\.id) == ["a", "b", "c"])
precondition(hits.map(\.score) == [1, 1, 0])

let reversed = try search(
    query: [1, 0], candidates: Array(candidates.reversed()),
    k: 3, threshold: 0
)
precondition(reversed.map(\.id) == hits.map(\.id))

let huge = try unit([
    Double.greatestFiniteMagnitude,
    Double.greatestFiniteMagnitude
])
precondition(abs(huge.reduce(0) { $0 + $1 * $1 } - 1) < 1e-12)

func rejects(_ body: () throws -> Void) -> Bool {
    do {
        try body()
        return false
    } catch {
        return true
    }
}

precondition(rejects { _ = try unit([0, 0]) })
precondition(rejects { _ = try unit([Double.nan, 1]) })
precondition(rejects {
    _ = try search(
        query: [1], candidates: candidates,
        k: 1, threshold: 0
    )
})

この実装では、意図的にすべての候補を検証してから結果を切り出しています。そのため、k がゼロでも、不正なデータがあればエラーになります。本番APIで即座に空の結果を返したい場合は、その動作を別途契約として明記する必要があります。テストと本番で契約の解釈が分かれる状態にしてはいけません。

幾何学的なサンプルでランキングの境界条件を検証する

基本的なサンプルに実際のテキストは必要ありません。同じ方向ならスコアは1、逆方向なら−1、直交するなら0です。また、正の数を掛けても方向は変わらないはずです。コード内の候補aとbはどちらもクエリと同じ方向を向いています。入力順を逆にしても識別子順に並ぶことを確認すれば、走査順に依存した実装を検出できます。

さらに、次の項目も確認してください。

• k をそれぞれゼロ、1、候補数、候補数より大きい値にする。
• しきい値を−1、0、1に設定し、等しい値が結果に含まれることを確認する。
• 空配列、無限大、識別子の重複、次元数の不一致は、必ずエラーにする。
• 固定サンプルで、類似度の対称性と、正の数を掛けても類似度が変わらないことを検証する。

数値のアサーションには許容誤差を使えますが、ソートの比較処理に「差が許容誤差未満なら等しい」という判定を使ってはいけません。近似的な等しさは推移律を満たさず、ソートに必要な条件を崩すおそれがあります。業務要件としてスコアを段階に分けて並べる必要があるなら、先に量子化ルールを定義し、その後で整数の段階値を比較してください。

参照実装をクラウド上の受け入れテストに組み込む

コマンドライン開発ツールがインストールされた環境で、次を実行します。

xcrun swift vector-contract.swift

スクリプトは、成功すると正常終了し、アサーションの失敗や未処理のエラーがあれば失敗を示す終了ステータスを返します。CIでは、この終了ステータスを必ず保持してください。後続コマンドの成功で上書きしてはいけません。コンパイルして実行する場合は、通常のデバッグビルドを使い、事前条件チェックを除去する最適化オプションは指定しないでください。

この参照実装は、各ベクトルを個別に正規化し、結果全体をソートします。小規模サンプルの受け入れテストには向いていますが、大規模検索の性能を確保するための実装ではありません。本番で正規化済みベクトルをキャッシュする場合は、キャッシュと元のベクトルのバージョンが一致しているかもテストしてください。一括計算やインデックスを使う実装に切り替えた場合も、この参照実装の結果と出力を照合します。

再現可能な実行記録を残す

毎回、コミット識別子、ツールチェーンのバージョン、サンプルの要約、次元数、しきい値、期待する候補の順序を保存します。サンプルには人工的に作成したベクトルか、利用許可を得た匿名化済みデータだけを使い、ユーザーの元のテキストを公開ログに含めないでください。

アルゴリズムの回帰とモデルのドリフトを切り分ける

固定ベクトルのテストは通るのに実際の検索結果が変わった場合に、初めてモデルのバージョン、トークン化の方法、ベクトルの前処理を調べます。固定ベクトルでも失敗するなら、数値計算の経路、ソート、候補の識別情報を優先して確認してください。ベクトルを再生成して不具合を覆い隠してはいけません。

最後に、受け入れテストの範囲を明確にしておきましょう。厳密検索では結果の一致が必要です。近似インデックスには、別途、再現率の基準も必要になります。また、ハードウェアが異なれば、浮動小数点数の末尾の桁までビット単位で一致するとは限りません。まず入力の妥当性、明確な境界条件、同点時の安定した順序を確保し、その後でデータ規模を拡大します。そうすることで、回帰レポートが修正可能な問題を指し示すようになります。

よくある質問