Video-Tree-TRM5/reference/docs/architecture.md

Video-Tree-TRM 统一架构设计

1. 核心定位

结合 TRM 递归推理 + PageIndex 树状导航的 模态无关 RAG 系统。

设计原则:
  - 检索阶段模态无关（全文本嵌入空间）
  - 模态差异封装在两端（预处理 + 答案生成）
  - TRM ACT 机制控制动态检索深度
  - 树深度固定 3 层，检索轮次动态

---

2. 系统总览

 ┌──────────────────── 预处理（离线，一次性） ────────────────────┐
 │                                                               │
 │  原始输入                   统一表示                           │
 │  ┌──────────┐              ┌──────────────────────┐           │
 │  │ 长文本    │──TextTree──→│                      │           │
 │  └──────────┘   Builder    │     TreeIndex         │           │
 │  ┌──────────┐              │  (全文本嵌入, [D])    │           │
 │  │ 长视频    │──VideoTree─→│                      │           │
 │  └──────────┘   Builder    └──────────┬───────────┘           │
 │                                       │                       │
 └───────────────────────────────────────│───────────────────────┘
                                         │
 ┌──────────────────── 检索（在线） ─────│───────────────────────┐
 │                                       ↓                       │
 │  query ──text_embed──→ q ──→ RecursiveRetriever               │
 │                              │  TRM 双循环推理                │
 │                              │  ACT halt 动态停止             │
 │                              │  多轮 root-to-leaf 遍历        │
 │                              └──→ List[NodePath]              │
 │                                     │                         │
 └─────────────────────────────────────│─────────────────────────┘
                                       │
 ┌──────────────────── 生成（在线） ────│─────────────────────────┐
 │                                      ↓                        │
 │  文本模式: LLM(query, raw_text_chunks)  → answer              │
 │  视频模式: VLM(query, frame_images)     → answer              │
 │                                                               │
 └───────────────────────────────────────────────────────────────┘

---

3. TreeIndex：统一数据结构

三层固定深度，文本/视频同构。

TreeIndex
├── metadata: IndexMeta          # 来源、嵌入器版本、创建时间
├── embed_dim: int               # 嵌入维度 D
└── roots: List[L1Node]          # L1 节点列表（树的根层）

L1Node (段/章)
├── id: str
├── summary: str                 # 聚合自子 L2 描述 (2-3句)
├── embedding: ndarray [D]       # text_embed(summary)
├── time_range: (float, float)   # 仅视频模式
└── children: List[L2Node]

L2Node (片段/小节)
├── id: str
├── description: str             # 直接从原始内容生成 (1-2句)
├── embedding: ndarray [D]       # text_embed(description)
├── time_range: (float, float)   # 仅视频模式
└── children: List[L3Node]

L3Node (帧/文本块)
├── id: str
├── description: str             # 视频=VLM帧描述(继承L2上下文), 文本=原始段落文本
├── embedding: ndarray [D]       # text_embed(description)
├── raw_content: str             # 原始文本块（文本模式）
├── frame_path: Optional[str]    # 帧图像路径（视频模式）
└── timestamp: Optional[float]   # 帧时间戳（视频模式）

关键设计：embedding 字段全部来自同一个 text_embed() 函数，不存在跨模态嵌入空间问题。raw_content / frame_path 仅供答案生成阶段使用，检索器不关心。

---

4. 预处理管线

4.1 摘要粒度规范

各层摘要的唯一职责是生成好的路由嵌入——让 (q+z)·M^T/√D 能选对节点。

层级	路由功能	摘要目标	粒度
L1	粗路由：选主题区域	"这个区域关于什么"	2-3句：主题 + 覆盖范围 + 关键实体
L2	细聚焦：选具体片段	"这个片段发生了什么"	1-2句：具体事件/内容 + 区分性细节
L3(文本)	精确定位：选文本块	段落原文	原始段落文本，不做摘要
L3(视频)	精确定位：选帧	"这帧画面的具体内容"	1-2句：继承L2上下文，聚焦区分性视觉信息

关键原则:
  - L1 要宽: 涵盖其下所有 L2 的语义范围，让相关 query 能"进对门"
  - L2 要窄: 只描述自己，与同级兄弟节点形成区分
  - L3 要具体: 提供精确定位的细节信息

文本示例：

L1: "第三章讨论了用户认证系统的设计，涵盖OAuth流程、JWT令牌管理和权限控制。"
  L2: "OAuth 2.0授权码流程的实现，包括重定向和回调处理。"
  L2: "JWT令牌的生成、验证和刷新机制。"
  L2: "基于角色的权限控制模型（RBAC）。"

视频示例：

L1: "厨师在厨房制作意大利面，从准备食材到最终装盘。"
  L2: "切洋葱、蒜末和番茄，准备酱料食材。"
    L3: "[准备酱料食材中] 厨师左手按住洋葱，右手快速切丁，砧板上已有切好的蒜末。"
    L3: "[准备酱料食材中] 厨师将番茄放入沸水中烫皮，锅中冒出蒸汽。"
  L2: "煮面条并制作番茄肉酱。"
  L2: "装盘摆盘并撒上帕尔马干酪。"

4.2 构建顺序：L2 轴心策略

两种模态共享同一构建原则：L2 为轴心，向下展开 L3，向上聚合 L1。

构建依赖关系:

  L3 需要 L2 上下文才能生成高质量描述（视频尤为关键）
  L1 摘要应从 L2 聚合，确保覆盖面完整
  → 循环依赖 → 解法: L2 不依赖 L3，而是从原始内容独立生成

构建顺序:

  Step 1: 结构切分 → 确定 L1/L2 边界
  Step 2: L2 先行 → 从原始内容直接生成 L2 描述
  Step 3: L3 向下 → 注入 L2 上下文，生成细粒度描述
  Step 4: L1 向上 → 聚合 L2 描述，生成粗粒度摘要

         Step 4: 聚合
            ↑
  ┌─────── L1 ───────┐      "第三章讨论了认证系统..."
  │                   │       ← LLM("总结这些片段: " + L2_descriptions)
  │   Step 2: 先行    │
  ├── L2 ── L2 ── L2 ┤      "OAuth流程的实现..."
  │    │    │    │     │      ← 文本: LLM 摘要段落组
  │    ↓    ↓    ↓     │        视频: VLM 看少量代表帧
  │   L3s  L3s  L3s    │
  │   Step 3: 向下     │      "厨师左手按住洋葱..."
  └────────────────────┘      ← 注入 L2 上下文后生成

4.3 TextTreeBuilder

输入: 长文本文档
输出: TreeIndex（全部 embedding=None）

Pipeline:
  Step 1 — 结构切分
     有 ToC → 解析章节层级 → L1/L2 边界
     无 ToC → LLM 语义分段 (一次性调用)

  Step 2 — L2 先行
     L2: 段落组 → LLM 生成摘要 (1-2句)
         摘要目标: 与兄弟 L2 形成区分
         （batch_chat 并发生成所有 L2 摘要）

  Step 3 — L3 向下
     L3: 原始段落文本直接复用
         存储 raw_content = description = 原始文本
         （文本模式 L3 不需要 L2 上下文，不调用 LLM）

  Step 4 — L1 向上聚合
     L1: LLM("总结这些小节: " + 所有子 L2 摘要) → 摘要
         摘要目标: 覆盖下属所有 L2 的语义范围 (2-3句)

  Step 5 — 序列化 → TreeIndex（JSON，无 embedding）

  ⚠ 延迟 Embedding: 所有节点 embedding=None
     首次检索时由 Pipeline._embed_tree() → tree.embed_all() 统一填充

4.4 VideoTreeBuilder

状态: ✅ 已实现（video_tree_trm/video_tree_builder.py）

输入: 长视频文件路径 或 YouTube URL
输出: TreeIndex（全部 embedding=None）

Pipeline（ThreadPoolExecutor 异步事件循环）:
  Step 0 — 输入类型判断
     本地文件: 直接使用 OpenCV 读取
     YouTube URL: yt-dlp -g 获取 CDN 直链 + yt-dlp --dump-json 获取时长

  Step 1 — 时间切分（固定步长）
     本地: cv2 读取总时长
     HTTP 流: 使用 yt-dlp 元数据时长（duration_hint，避免 cv2 流上不可靠）
     固定步长 l1_segment_duration=600s → L1 区间列表
     每个 L1 区间 → 等分 l2_clip_duration=60s → L2 clips

  Step 2 — 预计算任务图
     收集所有 L2 任务 + 记录每个 L1 的 L2 数量（l2_counts）

  Step 3 — 一次性提交所有 L2 任务（非阻塞）
     ThreadPoolExecutor(max_workers=concurrency)
     每个 L2 clip: 均匀 seek l2_representative_frames=10 帧 → VLM → 1-2句描述

  Step 4 — 事件循环（cfwait FIRST_COMPLETED）
     L2 完成 → 立即提交 L3 任务（_build_l3_task，线程安全）
       L3: 按 l3_fps 提取全量帧（持久化缓存），注入 L2 上下文，VLM 批量帧描述（JSON）
           降级路径: JSON 解析失败 → 逐帧 VLM 调用
     L3 完成 → 检查该 L1 的所有 L2 是否就绪 → 触发 L1 任务
       L1: 拼接所有 L2 描述 → vlm.chat() → 2-3句摘要
     主线程单线程操作 l1_l2_buckets，无竞争，无需 Lock

  Step 5 — 有序重建 l1_nodes，组装 TreeIndex（写 IndexMeta + 日志）

  ⚠ 延迟 Embedding: 所有节点 embedding=None
     首次检索时由 Pipeline._embed_tree() → tree.embed_all() 统一填充

L2 代表帧采样（均匀 seek，首尾均包含）：

step = (end_sec - start_sec) / (n_rep - 1)
timestamps = [start_sec + i * step for i in range(n_rep)]
# → 直接 cap.set(CAP_PROP_POS_MSEC, ts * 1000) 提取，缓存为 l2_{ts:.3f}.jpg

L3 帧描述 prompt（继承 L2 上下文 + 要求 JSON 输出）：

prompt = (
    '该片段的整体内容: "{l2_description}"\n'
    "以下是该片段中连续的 {n} 帧画面。\n"
    "对每帧用一到两句话描述其具体画面内容。\n"
    "重点关注: 动作、物体变化、文字信息、人物表情。\n"
    "不要重复片段整体描述，聚焦每帧的区分性信息。\n"
    '只返回 JSON 数组，格式: ["帧1描述", "帧2描述", ...]，不要其他内容。'
)
descs = VLM(frames_batch, prompt)  # 主路径：一次调用，JSON 解析
# 降级：json.loads 失败 → 逐帧调用

L1 摘要 prompt：

l2_texts = "\n".join(f"- {node.description}" for node in l2_children)
prompt = f"以下是一个视频段落中各片段的描述:\n{l2_texts}\n用2-3句话总结该段落的整体内容，涵盖所有片段的主题。"
l1_summary = vlm.chat(prompt)  # 纯文本调用

---

5. RecursiveRetriever：TRM 式递归检索

5.1 核心算法

节点选择使用 Cross-Attention（学习 W_q/W_k/W_v/W_o 投影），替代简单 cosine 路由。 L_level 推理模块使用 MLP-based（RMSNorm + SwiGLU），操作对象为向量 [B, D]。 三个可学习组件（CrossAttentionSelector, ReasoningModule, q_head）跨层级共享权重。

┌─────────────────────────────────────────────────────────────┐
│  RecursiveRetriever                                          │
│                                                              │
│  可训练组件（共享权重）:                                       │
│    selector = CrossAttentionSelector(W_q, W_k, W_v, W_o)    │
│    L_level  = ReasoningModule(RMSNorm + SwiGLU, L_layers层)  │
│    q_head   = Linear(D, 1)                                   │
│                                                              │
│  输入: query (str), tree (TreeIndex)                         │
│  输出: List[RetrievalPath]                                   │
│                                                              │
│  q = text_embed(query)          # 查询嵌入 [D]，冻结         │
│  z = q.clone()                  # 初始潜在状态               │
│                                                              │
│  for round in range(max_rounds):  ← ACT halt 控制            │
│  │                                                           │
│  │  ┌─── Phase 1: L1 粗粒度路由 (Cross-Attention) ──┐       │
│  │  │ s1, w1, k1* = selector(q+z, M_L1)             │       │
│  │  │   Q=W_q(q+z), K=W_k(M_L1), V=W_v(M_L1)       │       │
│  │  │   s1 = W_o(MultiHeadAttn(Q,K,V))   # 软选择    │       │
│  │  │   k1* = argmax(mean_head_scores)    # 硬索引    │       │
│  │  │ z = z + s1                                     │       │
│  │  │ for _ in L_cycles:                             │       │
│  │  │   z = L_level(z, s1 + q)                       │       │
│  │  └────────────────────────────────────────────────┘       │
│  │                                                           │
│  │  ┌─── Phase 2: L2 细粒度聚焦 (Cross-Attention) ──┐       │
│  │  │ M_L2 = children_embeds(L1[k1*])                │       │
│  │  │ s2, w2, k2* = selector(q+z, M_L2)             │       │
│  │  │ z = z + s2                                     │       │
│  │  │ for _ in L_cycles:                             │       │
│  │  │   z = L_level(z, s2 + q)                       │       │
│  │  └────────────────────────────────────────────────┘       │
│  │                                                           │
│  │  ┌─── Phase 3: L3 精确定位 (Cross-Attention) ────┐       │
│  │  │ M_L3 = children_embeds(L2[k2*])                │       │
│  │  │ s3, w3, k3* = selector(q+z, M_L3)             │       │
│  │  │ z = z + s3                                     │       │
│  │  └────────────────────────────────────────────────┘       │
│  │                                                           │
│  │  collected.append(Path(k1*, k2*, k3*))                    │
│  │                                                           │
│  │  ┌─── ACT Halt Decision ───────────────────┐             │
│  │  │ halt_logit = q_head(z)                   │             │
│  │  │ if halt_logit > 0 and round_idx > 0:     │             │
│  │  │     break  # 至少跑 1 轮                  │             │
│  │  └──────────────────────────────────────────┘             │
│  │                                                           │
│  └── z 状态保留到下一轮（累积已检索信息）                     │
│                                                              │
│  return collected                                            │
└─────────────────────────────────────────────────────────────┘

5.2 算法伪代码

class RecursiveRetriever:
    """TRM 式递归检索器（Cross-Attention 版本）。"""

    def __init__(self, embed_dim: int, num_heads: int, L_layers: int, L_cycles: int, max_rounds: int):
        self.selector = CrossAttentionSelector(embed_dim, num_heads)  # 共享，跨层级
        self.L_level = ReasoningModule(embed_dim, L_layers)           # 共享，跨层级
        self.q_head = Linear(embed_dim, 1)                            # ACT halt head
        self.L_cycles = L_cycles
        self.max_rounds = max_rounds

    def forward(
        self, q: Tensor, tree: TreeIndex, return_internals: bool = False
    ) -> Dict[str, Any]:
        z = q.clone()                   # [B, D]，初始潜在状态 = 查询嵌入
        paths = []

        for round_idx in range(self.max_rounds):
            # ── 三阶段树遍历（一次完整 root-to-leaf） ──
            path, z, step_attns = self._traverse_one_path(q, z, tree)
            paths.append(path)

            # ── ACT halt（推理模式，至少走 1 轮） ──
            halt_logit = self.q_head(z)  # [B, 1]
            if not self.training and halt_logit.item() > 0 and round_idx > 0:
                break

        return {"paths": paths, "num_rounds": len(paths), "z_final": z}

    def _traverse_one_path(self, q, z, tree):
        """单次 root → L1 → L2 → L3 遍历。"""

        # Phase 1
        k1, z = self._select_and_reason(q, z, tree.l1_embeddings())
        # Phase 2
        k2, z = self._select_and_reason(q, z, tree.l2_embeddings_of(k1))
        # Phase 3
        k3, z = self._select_and_reason(q, z, tree.l3_embeddings_of(k1, k2))

        return Path(k1, k2, k3), z

    def _select_and_reason(self, q, z, M):
        """单层: Cross-Attention 选择 + L_cycles 内循环推理。"""

        # Navigate: Cross-Attention
        state = q + z                                      # [B, D]
        selected_info, attn_w, k_star = self.selector(state, M)
        #   Q = W_q(state)   →  [B, 1, D]
        #   K = W_k(M)       →  [B, N, D]
        #   V = W_v(M)       →  [B, N, D]
        #   selected_info = W_o(MultiHeadAttn(Q, K, V))    [B, D]  ← 可微
        #   k_star = argmax(avg_head_scores)                        ← 路径索引

        # Update: z += attention 加权信息
        z = z + selected_info

        # Reason: TRM L-level MLP 内循环
        for _ in range(self.L_cycles):
            z = self.L_level(z, selected_info + q)

        return k_star, z

5.3 多轮检索的 z 状态流

Round 1:                                    Round 2:
z₀ = q                                     z₃ 来自 Round 1（包含已检索信息）
  │                                           │
  ├─ Phase1: CA(q+z₀, M_L1) → s1            ├─ Phase1: CA(q+z₃, M_L1) → s4
  │  z₁ = z₀ + s1 → L_level×L_cycles        │  z₄ = z₃ + s4 → L_level×L_cycles
  ├─ Phase2: CA(q+z₁, M_L2[k1]) → s2       ├─ Phase2: CA(q+z₄, M_L2[k4]) → s5
  │  z₂ = z₁ + s2 → L_level×L_cycles        │  z₅ = z₄ + s5 → L_level×L_cycles
  ├─ Phase3: CA(q+z₂, M_L3[k2]) → s3       ├─ Phase3: CA(q+z₅, M_L3[k5]) → s6
  │  z₃ = z₂ + s3                            │  z₆ = z₅ + s6
  │                                           │
  ACT: q_head(z₃) < 0 → 继续                ACT: q_head(z₆) > 0 → 停止
                                              → 返回 [Path(k1,k2,k3), Path(k4,k5,k6)]

关键: s = W_o(MultiHeadAttn(W_q(q+z), W_k(M), W_v(M)))
      比简单 cosine 路由更强:
        - 学习的投影决定"关注什么"进行选择
        - Multi-head 捕获多种相关性信号
        - V 投影让"更新信息"与"匹配键"解耦
      z₃ 累积了 Round 1 的 attention 信息
        → (q + z₃) 的投影方向自动偏离已选区域

---

6. ACT Halt 训练方案

6.1 训练目标

ACT halt head 学习: "已收集的信息是否足以回答 query"

                     ┌──────────────────────────────┐
                     │  reward 定义                   │
                     │                                │
                     │  R = answer_quality - λ·rounds │
                     │                                │
                     │  answer_quality:               │
                     │    文本QA: EM / F1 score        │
                     │    视频QA: 选项匹配准确率       │
                     │                                │
                     │  λ: 步数惩罚系数               │
                     │    鼓励用更少轮次达到同样质量   │
                     └──────────────────────────────┘

6.2 训练数据

数据集	模态	样本量	reward 信号
LongBench	文本	~5K	ground truth → F1
NarrativeQA	文本	~30K	ground truth → ROUGE
VideoMME	视频	~2K	选项匹配 → 0/1

6.3 训练流程

# 训练伪代码
for query, tree, ground_truth in dataset:
    q = text_embed(query)
    z = q.copy()
    total_loss = 0

    for round_idx in range(max_rounds):
        path, z = retriever.traverse_one_path(q, z, tree)

        # halt 决策
        halt_logit = q_head(z)

        # 用当前已收集的 context 生成答案，计算质量
        context = collect_raw_content(paths_so_far)
        answer = generator(query, context)
        quality = compute_score(answer, ground_truth)  # EM/F1

        # Q-learning target
        # 如果现在停 → reward = quality
        # 如果继续   → reward = γ · future_quality - λ
        target_q = quality if is_last else γ * next_quality - λ
        loss += mse(sigmoid(halt_logit), target_q)

    total_loss.backward()
    optimizer.step()

6.4 可训练组件 vs 冻结组件

冻结 (不训练):
  ✗ text_embed()        # 预训练嵌入器，冻结
  ✗ TreeIndex embeddings # 预计算的节点嵌入，冻结

可训练:
  ✓ CrossAttentionSelector (W_q, W_k, W_v, W_o)  # 节点选择投影
  ✓ L_level (ReasoningModule: RMSNorm + SwiGLU)   # MLP 推理模块
  ✓ q_head (ACT halt)                              # 停止决策头

---

7. 检索结果与答案生成的接口

@dataclass
class RetrievalPath:
    """一条 root-to-leaf 路径。"""
    k1: int                    # L1 索引
    k2: int                    # L2 索引
    k3: int                    # L3 索引
    l1_summary: str            # L1 摘要
    l2_description: str        # L2 描述
    l3_description: str        # L3 描述
    raw_content: Optional[str] # 原始文本（文本模式）
    frame_path: Optional[str]  # 帧路径（视频模式）
    timestamp: Optional[float] # 时间戳（视频模式）


@dataclass
class RetrievalResult:
    """检索器输出。"""
    query: str
    paths: List[RetrievalPath]  # 多轮收集的路径
    num_rounds: int             # 实际检索轮次
    z_final: ndarray            # 最终潜在状态


# 答案生成器接口
def generate_answer(query: str, result: RetrievalResult) -> str:
    if is_text_mode(result):
        context = "\n".join(p.raw_content for p in result.paths)
        return LLM(query=query, context=context)
    else:
        frames = [load_image(p.frame_path) for p in result.paths]
        captions = [p.l3_description for p in result.paths]
        return VLM(query=query, images=frames, captions=captions)

---

8. 与现有系统的关系

参考代码                            新架构                    变化
─────────────────────────────────────────────────────────────────
video_pyramid.py (HSP)            → TreeIndex               重构为统一格式
video_tree_trm.py (cosine路由)    → RecursiveRetriever      Cross-Attention+ACT
  select_next_node (CrossAttn)    → CrossAttentionSelector   保留CA思路，简化为向量级
  L_level (Transformer blocks)    → ReasoningModule          MLP-based (向量非序列)
visual_projection.py              → 删除                    L3 全文本化
video_indexer.py (CLIP encode)    → embeddings.py            统一 text_embed()
pipeline.py                       → pipeline.py             ✅ 已实现（含延迟 embed 策略）
answer_generator.py               → answer_generator.py     ✅ 已实现
config.py                         → config.py               全面重构

新增:
  + tree_index.py                 统一数据结构              ✅ 已实现
  + embeddings.py                 嵌入服务封装              ✅ 已实现
  + llm_client.py                 LLM/VLM 客户端            ✅ 已实现
  + text_tree_builder.py          文本模式预处理            ✅ 已实现
  + video_tree_builder.py         视频模式预处理            ✅ 已实现
  + recursive_retriever.py        TRM 递归检索器 (CA+MLP+ACT) ✅ 已实现
  + losses.py                     NavigationLoss + ACTLoss  ✅ 已实现
  + train.py                      两阶段训练入口            ✅ 已实现
  + main.py                       推理/演示入口             ✅ 已实现