chore: track claude skills, tools, templates, reference code and research-wiki
- Add all claude skills (brainstorming, commit, debugging, TDD, etc.) - Add claude hooks (pre-commit-guard, post-edit-quality) - Add research templates (experiment plan, research brief, etc.) - Add claude tools (arxiv/semantic_scholar/openalex fetch, wiki, exa) - Add TRM4 reference implementation as algorithm fidelity baseline - Add research-wiki content (plans, index, graph, query_pack) - Update .gitignore to exclude .graphify_version runtime state
This commit is contained in:
@@ -0,0 +1,330 @@
|
||||
# Video-Tree-TRM 实验计划
|
||||
|
||||
> **目标**: 验证结合 TRM 多层推理探索能力（Cross-Attention Selector + ACT Halt）
|
||||
> 与 PageIndex 树状检索能力的 Video-Tree-TRM 在长文本和长视频问答任务上的效果。
|
||||
>
|
||||
> **两条并行实验线路**：
|
||||
> - **线路 A** — 长文本检索（LongBench / NarrativeQA）
|
||||
> - **线路 B** — 长视频检索（VideoMME）
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
|
||||
- [评测体系](#1-评测体系)
|
||||
- [实验环境准备](#2-实验环境准备)
|
||||
- [线路 A：长文本检索](#3-线路-a长文本检索)
|
||||
- [线路 B：长视频检索](#4-线路-b长视频检索)
|
||||
- [基线方法](#5-基线方法)
|
||||
- [消融实验](#6-消融实验)
|
||||
- [里程碑检查表](#7-里程碑检查表)
|
||||
- [结果记录表](#8-结果记录表)
|
||||
|
||||
---
|
||||
|
||||
## 1. 评测体系
|
||||
|
||||
### 1.1 质量指标
|
||||
|
||||
| 指标 | 适用模态 | 计算方式 |
|
||||
|------|---------|---------|
|
||||
| **EM (Exact Match)** | 文本 QA | 标准化后精确字符串匹配，0/1 |
|
||||
| **F1** | 文本 QA | token 级 precision/recall，`token_f1()` 已实现于 `answer_generator.py` |
|
||||
| **Accuracy** | 视频 QA（多选） | 模型选项与标准答案选项完全匹配的比率 |
|
||||
|
||||
### 1.2 效率指标
|
||||
|
||||
| 指标 | 说明 |
|
||||
|------|------|
|
||||
| **Avg Rounds** | 所有样本的平均检索轮次（衡量 ACT Halt 效果，越少越好） |
|
||||
| **Max Rounds Hit Rate** | 触达 `max_rounds` 上限而非主动停止的样本比率 |
|
||||
|
||||
### 1.3 诊断指标
|
||||
|
||||
| 指标 | 说明 |
|
||||
|------|------|
|
||||
| **Nav Accuracy @L1** | 第一轮检索选中正确 L1 节点的比率 |
|
||||
| **Nav Accuracy @L2** | 在正确 L1 下选中正确 L2 节点的比率 |
|
||||
| **Nav Accuracy @L3** | 在正确 L2 下选中正确 L3 节点的比率 |
|
||||
|
||||
---
|
||||
|
||||
## 2. 实验环境准备
|
||||
|
||||
### 2.1 安装依赖
|
||||
|
||||
```bash
|
||||
conda activate Video-Tree-TRM
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 2.2 配置文件
|
||||
|
||||
复制并填写环境变量：
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# 填写 LLM_API_KEY / VLM_API_KEY / EMBED_API_KEY（若使用远程嵌入）
|
||||
```
|
||||
|
||||
默认配置文件 `config/default.yaml` 已预置：
|
||||
|
||||
```yaml
|
||||
embed.model_name: "BAAI/bge-base-zh-v1.5" # 嵌入模型
|
||||
retriever.embed_dim: 2560
|
||||
retriever.max_rounds: 5
|
||||
train.max_epochs_phase1: 30
|
||||
train.max_epochs_phase2: 20
|
||||
```
|
||||
|
||||
> 消融实验时通过 `--set key=value` 覆盖，无需修改 yaml 文件。
|
||||
|
||||
---
|
||||
|
||||
## 3. 线路 A：长文本检索
|
||||
|
||||
> **优先级**: P0（LongBench）→ P1（NarrativeQA）
|
||||
> **并行方式**: 与线路 B 同步推进，互不阻塞
|
||||
|
||||
### A.1 数据集准备
|
||||
|
||||
| 数据集 | 样本量 | 格式 | 优先级 |
|
||||
|--------|--------|------|--------|
|
||||
| **LongBench** | ~5K | JSONL `{"query", "answer", "source_path", "modality": "text"}` | P0 |
|
||||
| **NarrativeQA** | ~30K | 同上 | P1 |
|
||||
|
||||
```bash
|
||||
# 下载 LongBench（示例）
|
||||
mkdir -p data/longbench data/narrativeqa
|
||||
# 将转换后的 JSONL 文件放置于对应目录
|
||||
```
|
||||
|
||||
配置覆盖：
|
||||
|
||||
```bash
|
||||
# LongBench
|
||||
--set train.dataset=longbench --set train.dataset_path=data/longbench
|
||||
|
||||
# NarrativeQA（P1，LongBench 完成后）
|
||||
--set train.dataset=narrativeqa --set train.dataset_path=data/narrativeqa
|
||||
```
|
||||
|
||||
### A.2 索引构建
|
||||
|
||||
```bash
|
||||
# 对所有文本文件批量构建 TreeIndex（含磁盘缓存）
|
||||
conda run -n Video-Tree-TRM python main.py index \
|
||||
--source data/longbench/doc.txt \
|
||||
--modality text \
|
||||
--config config/default.yaml
|
||||
```
|
||||
|
||||
> 缓存命中后重复运行零开销，构建结果保存至 `cache/trees/`。
|
||||
|
||||
### A.3 两阶段训练
|
||||
|
||||
```bash
|
||||
# Phase 1：导航训练（单轮，~30 epoch）
|
||||
conda run -n Video-Tree-TRM python train.py \
|
||||
--config config/default.yaml \
|
||||
--set train.dataset=longbench \
|
||||
--set train.dataset_path=data/longbench
|
||||
|
||||
# Phase 2 权重加载自 Phase 1 最佳检查点
|
||||
# 在 config/default.yaml 中设置:
|
||||
# retriever.checkpoint: checkpoints/phase1_epoch30.pt
|
||||
# train.max_epochs_phase2: 20
|
||||
conda run -n Video-Tree-TRM python train.py \
|
||||
--config config/default.yaml \
|
||||
--set retriever.checkpoint=checkpoints/phase1_epoch30.pt
|
||||
```
|
||||
|
||||
训练检查点保存至 `checkpoints/phase1_epoch{N}.pt` / `phase2_epoch{N}.pt`。
|
||||
|
||||
### A.4 推理评测
|
||||
|
||||
```bash
|
||||
# 单样本问答验证
|
||||
conda run -n Video-Tree-TRM python main.py query \
|
||||
--source data/longbench/doc.txt \
|
||||
--modality text \
|
||||
--question "文档的核心观点是什么？" \
|
||||
--config config/default.yaml
|
||||
```
|
||||
|
||||
批量评测（自行实现评测脚本）：
|
||||
- 遍历测试集 JSONL
|
||||
- 调用 `Pipeline.query()` 获取预测答案
|
||||
- 计算 EM / F1，汇总 Avg Rounds
|
||||
|
||||
### A.5 基线对比
|
||||
|
||||
见 [第 5 节](#5-基线方法)。
|
||||
|
||||
---
|
||||
|
||||
## 4. 线路 B：长视频检索
|
||||
|
||||
> **优先级**: P2（可与线路 A 并行推进）
|
||||
|
||||
### B.1 数据集准备
|
||||
|
||||
| 数据集 | 样本量 | 格式 | 评测指标 |
|
||||
|--------|--------|------|---------|
|
||||
| **VideoMME** | ~2K | JSONL `{"query", "answer", "source_path", "modality": "video", "timestamp": float}` | Accuracy（多选） |
|
||||
|
||||
```bash
|
||||
mkdir -p data/videomme
|
||||
# 下载 VideoMME 并转换为上述 JSONL 格式
|
||||
```
|
||||
|
||||
### B.2 索引构建
|
||||
|
||||
```bash
|
||||
conda run -n Video-Tree-TRM python main.py index \
|
||||
--source data/videomme/video.mp4 \
|
||||
--modality video \
|
||||
--config config/default.yaml
|
||||
```
|
||||
|
||||
关键配置：
|
||||
|
||||
```yaml
|
||||
tree:
|
||||
l1_segment_duration: 600.0 # L1 段时长（秒）
|
||||
l2_clip_duration: 60.0 # L2 clip 时长（秒）
|
||||
l3_fps: 1.0 # L3 帧提取频率
|
||||
l2_representative_frames: 10 # VLM 描述用代表帧数
|
||||
```
|
||||
|
||||
### B.3 两阶段训练
|
||||
|
||||
```bash
|
||||
# Phase 1（视频模态）
|
||||
conda run -n Video-Tree-TRM python train.py \
|
||||
--config config/default.yaml \
|
||||
--set train.dataset=videomme \
|
||||
--set train.dataset_path=data/videomme
|
||||
|
||||
# Phase 2（视频模态，加载 Phase 1 检查点）
|
||||
conda run -n Video-Tree-TRM python train.py \
|
||||
--config config/default.yaml \
|
||||
--set train.dataset=videomme \
|
||||
--set train.dataset_path=data/videomme \
|
||||
--set retriever.checkpoint=checkpoints/phase1_epoch30.pt
|
||||
```
|
||||
|
||||
### B.4 推理评测
|
||||
|
||||
```bash
|
||||
conda run -n Video-Tree-TRM python main.py query \
|
||||
--source data/videomme/video.mp4 \
|
||||
--modality video \
|
||||
--question "视频中发生了什么事？" \
|
||||
--config config/default.yaml
|
||||
```
|
||||
|
||||
批量评测：遍历测试集 → `Pipeline.query()` → 与多选选项比对 → 计算 Accuracy。
|
||||
|
||||
### B.5 基线对比
|
||||
|
||||
见 [第 5 节](#5-基线方法)。
|
||||
|
||||
---
|
||||
|
||||
## 5. 基线方法
|
||||
|
||||
| 方法 | 描述 | 实现方式 |
|
||||
|------|------|---------|
|
||||
| **BM25 + LLM** | 传统稀疏检索 | `rank_bm25` 库检索 top-k 段落 → 拼接上下文 → `LLMClient.chat()` |
|
||||
| **Dense Retrieval + LLM** | BGE 向量检索 + rerank | `EmbeddingModel.embed_tensor()` 全量检索 top-k → rerank → 生成 |
|
||||
| **PageIndex（无 TRM）** | 树状导航，cosine 路由，无推理模块 | 替换 `CrossAttentionSelector` 为 cosine 相似度选节点 |
|
||||
| **Tree-TRM（原论文）** | 原始实现 | 参考 `Reference/Tree-TRM/` 目录 |
|
||||
| **Video-Tree-TRM（ours）** | 本项目实现 | `Pipeline.query()` |
|
||||
|
||||
评测时各方法使用相同数据集和评测脚本，确保公平对比。
|
||||
|
||||
---
|
||||
|
||||
## 6. 消融实验
|
||||
|
||||
在主实验（LongBench）最优配置基础上，逐一变更单一变量：
|
||||
|
||||
| 编号 | 变量 | 候选值 | 配置覆盖 | 预期观察 |
|
||||
|------|------|--------|---------|---------|
|
||||
| **A1** | 选择器类型 | Cross-Attention vs Cosine | 替换 `CrossAttentionSelector` 实现 | CA 路由是否带来 F1 提升 |
|
||||
| **A2** | 推理深度 | L_cycles ∈ {1, 2, 4, 8} | `--set retriever.L_cycles=N` | 质量-计算量权衡 |
|
||||
| **A3** | 推理模块层数 | L_layers ∈ {1, 2, 4} | `--set retriever.L_layers=N` | 网络深度的边际收益 |
|
||||
| **A4** | 多轮检索上限 | max_rounds ∈ {1, 3, 5} | `--set retriever.max_rounds=N` | ACT 多轮边际收益 |
|
||||
| **A5** | ACT Halt 机制 | 有 / 无 | `act_loss_weight=0.0` 禁用 | ACT 对效率和质量的贡献 |
|
||||
| **A6** | 注意力头数 | num_heads ∈ {1, 4, 8} | `--set retriever.num_heads=N` | 多头注意力的容量影响 |
|
||||
|
||||
每组消融实验保持其余超参数为默认值（`config/default.yaml`）。
|
||||
|
||||
---
|
||||
|
||||
## 7. 里程碑检查表
|
||||
|
||||
### 线路 A（文本）
|
||||
|
||||
- [ ] 环境配置完成，`python main.py --help` 正常输出
|
||||
- [ ] LongBench 数据集下载 & 转换为 JSONL 格式
|
||||
- [ ] LongBench 索引构建完成（`cache/trees/` 缓存生成）
|
||||
- [ ] Phase 1 训练完成（文本，LongBench）
|
||||
- [ ] Phase 2 训练完成（文本，LongBench）
|
||||
- [ ] LongBench 批量评测完成（EM / F1 / Avg Rounds）
|
||||
- [ ] 4 条基线方法评测完成
|
||||
- [ ] NarrativeQA 实验（P1，可选）
|
||||
|
||||
### 线路 B（视频）
|
||||
|
||||
- [ ] VideoMME 数据集下载 & 转换为 JSONL 格式
|
||||
- [ ] VideoMME 索引构建完成（视频帧提取 + VLM 描述生成）
|
||||
- [ ] Phase 1 训练完成（视频，VideoMME）
|
||||
- [ ] Phase 2 训练完成（视频，VideoMME）
|
||||
- [ ] VideoMME 批量评测完成（Accuracy / Avg Rounds）
|
||||
|
||||
### 消融实验
|
||||
|
||||
- [ ] A1: Cross-Attention vs Cosine 路由
|
||||
- [ ] A2: L_cycles 扫描（1/2/4/8）
|
||||
- [ ] A3: L_layers 扫描（1/2/4）
|
||||
- [ ] A4: max_rounds 扫描（1/3/5）
|
||||
- [ ] A5: 有/无 ACT Halt
|
||||
- [ ] A6: num_heads 扫描（1/4/8）
|
||||
|
||||
---
|
||||
|
||||
## 8. 结果记录表
|
||||
|
||||
### 8.1 主实验结果
|
||||
|
||||
| 方法 | LongBench EM | LongBench F1 | VideoMME Acc | Avg Rounds |
|
||||
|------|-------------|-------------|-------------|-----------|
|
||||
| BM25 + LLM | | | | — |
|
||||
| Dense Retrieval + LLM | | | | — |
|
||||
| PageIndex（无 TRM） | | | | |
|
||||
| Tree-TRM（原论文） | | | | |
|
||||
| **Video-Tree-TRM（ours）** | | | | |
|
||||
|
||||
### 8.2 消融实验结果（LongBench F1）
|
||||
|
||||
| 变量 | 值 | F1 | Avg Rounds | 备注 |
|
||||
|------|----|----|-----------|------|
|
||||
| 选择器 | Cross-Attention | | | 默认 |
|
||||
| 选择器 | Cosine | | | A1 |
|
||||
| L_cycles | 1 | | | A2 |
|
||||
| L_cycles | 2 | | | A2 |
|
||||
| L_cycles | 4 | | | A2 默认 |
|
||||
| L_cycles | 8 | | | A2 |
|
||||
| L_layers | 1 | | | A3 |
|
||||
| L_layers | 2 | | | A3 默认 |
|
||||
| L_layers | 4 | | | A3 |
|
||||
| max_rounds | 1 | | | A4 |
|
||||
| max_rounds | 3 | | | A4 |
|
||||
| max_rounds | 5 | | | A4 默认 |
|
||||
| ACT | 有 | | | A5 默认 |
|
||||
| ACT | 无 | | | A5 |
|
||||
| num_heads | 1 | | | A6 |
|
||||
| num_heads | 4 | | | A6 默认 |
|
||||
| num_heads | 8 | | | A6 |
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,557 @@
|
||||
# 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，首尾均包含）：
|
||||
|
||||
```python
|
||||
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 输出）：
|
||||
|
||||
```python
|
||||
prompt = (
|
||||
'该片段的整体内容: "{l2_description}"\n'
|
||||
"以下是该片段中连续的 {n} 帧画面。\n"
|
||||
"对每帧用一到两句话描述其具体画面内容。\n"
|
||||
"重点关注: 动作、物体变化、文字信息、人物表情。\n"
|
||||
"不要重复片段整体描述，聚焦每帧的区分性信息。\n"
|
||||
'只返回 JSON 数组，格式: ["帧1描述", "帧2描述", ...]，不要其他内容。'
|
||||
)
|
||||
descs = VLM(frames_batch, prompt) # 主路径：一次调用，JSON 解析
|
||||
# 降级：json.loads 失败 → 逐帧调用
|
||||
```
|
||||
|
||||
**L1 摘要 prompt**：
|
||||
|
||||
```python
|
||||
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 算法伪代码
|
||||
|
||||
```python
|
||||
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 训练流程
|
||||
|
||||
```python
|
||||
# 训练伪代码
|
||||
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. 检索结果与答案生成的接口
|
||||
|
||||
```python
|
||||
@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 推理/演示入口 ✅ 已实现
|
||||
```
|
||||
Reference in New Issue
Block a user