RAPTOR要約ツリーの実装 — バイナリ重心クラスタリングとcollapsed-tree検索

課題: チャンクRAGは「ズームアウト」できない

チャンク分割+ベクトル検索のRAGは、答えが特定の段落に書いてある質問には強い。しかし「このプロジェクトの方針は?」「この50ファイルは全体として何を扱っている?」のような複数チャンクにまたがる抽象的な質問では、どのチャンクも部分しか持っておらず、検索が空振りします。

RAPTORはこの問題への定番アプローチで、チャンクをクラスタリング→各クラスタを要約→要約をさらにクラスタリング、と再帰して抽象度の階層を作ります。質問が具体的ならリーフに、抽象的なら上位の要約ノードにヒットする、という構造です。

ツリーのデータ構造

Marianの実装では、各ノードは埋め込みそのものではなくバイナリ量子化された符号(centroid)を持ちます。これは検索基盤全体がバイナリ粗探索を前提にしているためで、ツリー検索も同じHamming距離の土俵に乗ります。

// lib/marian-search/raptor-tree.ts
export interface RaptorNode {
  id: string
  level: number        // 0 = リーフ、上に行くほど抽象
  code: Uint8Array     // バイナリ重心(MSB-first)
  bits: number
  childIds: string[]   // リーフでは空
  size: number         // 配下のリーフ数(リーフは1)
  summary?: string     // LLM生成の要約(任意注入)
}

export interface RaptorTree {
  nodes: Map<string, RaptorNode>
  rootIds: string[]
  bits: number
  levels: number
}

親ノードの重心: 葉数で重み付けした多数決

親ノードのバイナリ重心は、子ノードのビットごとの多数決で決めます。ただし単純な子ノード数ではなく、配下のリーフ数(size)で重み付けします。

function weightedCentroid(children: readonly RaptorNode[], bits: number): Uint8Array {
  const code = new Uint8Array(Math.ceil(bits / 8))
  const total = children.reduce((sum, child) => sum + child.size, 0)
  for (let b = 0; b < bits; b++) {
    let ones = 0
    for (const child of children) if (bitAt(child.code, b)) ones += child.size
    if (ones * 2 >= total) setBit(code, b) // 配下リーフの過半数が立っていれば1
  }
  return code
}

重み付けしない多数決だと、リーフ1枚だけの小さな子と100枚を束ねる子が同じ1票になり、重心が小さい枝に引っ張られます。葉数加重により、重心は常に「配下の実データの分布」を反映します。

クラスタリング: 貪欲最近傍で十分という割り切り

オリジナルのRAPTOR論文はGMM(ガウス混合)+UMAPでソフトクラスタリングしますが、Marianの実装はシードを起点にHamming類似度の近い未割当ノードをbranching数(デフォルト4)まで取るだけの貪欲法です。

function nearestCluster(seed: number, level: RaptorNode[], assigned: boolean[],
                        branching: number, bits: number): number[] {
  const candidates = []
  for (let j = 0; j < level.length; j++) {
    if (assigned[j] || j === seed) continue
    candidates.push({ index: j, score: binarySimilarity(level[seed].code, level[j].code, bits) })
  }
  candidates.sort((a, b) => b.score - a.score || a.index - b.index)
  return [seed, ...candidates.slice(0, branching - 1).map((c) => c.index)]
}

この割り切りには2つの理由があります。第一に、クラスタリングの良し悪しは最終的にリランクで吸収されるため、粗くても検索品質への影響が限定的であること。第二に、決定的(deterministic)であること。同じ入力からは常に同じツリーが生まれ、テストもインクリメンタル再構築も単純になります。ソート時のタイブレークにインデックス順を使っているのも決定性のためです。

collapsed-tree検索: 階層を辿らない

検索時はツリーを上から辿りません。全階層の全ノードを一括でスコアリングし、上位nodeTopK件を取ってからリーフに展開します(collapsed-tree方式)。

export function collapsedSearch(tree: RaptorTree, queryCode: Uint8Array, options = {}) {
  const hits: RaptorHit[] = []
  for (const node of tree.nodes.values()) {
    hits.push({ id: node.id, level: node.level, size: node.size,
                score: binarySimilarity(queryCode, node.code, tree.bits) })
  }
  hits.sort((a, b) => b.score - a.score || b.size - a.size || a.id.localeCompare(b.id))
  return hits.slice(0, topK)
}

ルートから貪欲に降りる方式は、上位階層で1回判断を誤ると正解の枝に二度と戻れません。collapsed方式なら「具体的な質問はリーフに、抽象的な質問は要約ノードに」が検索スコアの自然な結果として起き、トラバーサルポリシーの設計が丸ごと不要になります。同点時にsizeの大きいノードを優先するのは、迷ったら広いカバレッジを取るという方針です。

永続化: ツリーもただのテーブル

ツリーはPostgresの1テーブルに、ノード1行で永続化します。バイナリ重心はembedding_binaryと同じMSB-firstビット文字列です。

create table if not exists public.marian_file_tree_nodes (
  user_id    uuid not null references auth.users(id) on delete cascade,
  node_id    text not null,
  level      int  not null,
  size       int  not null,
  bits       int  not null,
  child_ids  text[] not null default '{}',
  code       text not null,  -- MSB-firstビット文字列
  summary    text,
  is_root    boolean not null default false,
  updated_at timestamptz not null default now(),
  primary key (user_id, node_id)
);

要約文(summary)の生成はビルド関数にsummarizeコールバックとして任意注入します。LLMなしでもツリーは構築でき(重心とサイズだけで検索は機能する)、要約はあとから埋められる。LLM依存を構造から分離するこの設計のおかげで、ツリー構築はユニットテスト可能な純粋関数になっています。

まとめ

MarianのRAPTOR実装は、(1)バイナリ重心による安価なクラスタリング、(2)葉数加重多数決で分布を保つ重心、(3)決定的な貪欲クラスタリング、(4)トラバーサル不要のcollapsed-tree検索、(5)LLM要約の注入分離、で構成されます。ASK_RAPTOR=1でオプトインし、ツリーがなければ通常検索に自動フォールバックするため、既存のAskを壊さずに「広い質問」への回答能力を追加できます。