Tree-of-Attention — 把长上下文拆成树再分支注意

待复核

是什么

Tree-of-Attention（ToA，注意力树）是一类把长上下文先建成树、再沿分支做注意力的做法：不是让每个 token 一次性盯住整篇文档，而是先在局部枝上算清楚，再把摘要往上汇、把问题往下路由。日常类比：查百科不是从第一页翻到最后一页，而是目录 → 章节 → 小节 → 段落，只打开相关那几支。

标准 self-attention 像全员大会：每个人都要和所有人说话，长度一长就又贵又吵（复杂度约 O(n²)，中间信息还容易丢）。ToA 把文档切成块，块再挂到父节点上，形成一棵上下文树；查询时只展开有希望的分支，其它枝可以剪掉或只看摘要。

它和 tree-of-thoughts-2023 名字像、对象不同：ToT 是推理步骤长成树；ToA 是注意力作用的上下文长成树。

一句话记：ToA 用「目录式路由」代替「整本书互看」，让长上下文推理先找对枝，再在枝内精读。

为什么重要

不理解 ToA，下面这些事都很难讲清：

为什么「上下文窗口开到 128K」仍会答错中间段落——扁平注意力会把注意力摊薄，树形路由能先定位再细读
为什么长文档问答、多跳检索常先切块再聚合——那是在手工搭一棵浅树
为什么 speculative decoding / 多候选生成要共享前缀 KV——分支共享的是同一棵前缀树
为什么只堆 flash-attention 不够：算得快仍是「全员互看」，未必解决「该看哪一枝」

核心要点

ToA 可以拆成三步记：

建树（chunk → node）：把长输入切成叶子块（段落、小节、固定长度窗口），再按标题层级或相似度挂到父节点。类比：先把书拆成页，再装订成章。
分支注意（branch attend）：查询先和父节点摘要比相关度，只进入 top-k 子枝；枝内再做局部 attention。类比：先看目录选章，再精读那一章。
汇合 / 剪枝（reduce & prune）：子枝结果向上合并；明显无关的枝提前停。类比：几条线索汇到案情板，死胡同不再派人。

三步合起来：用树拓扑换「全序列互看」，把算力花在相关分支上。工程上常见两种落地：稀疏 mask（一次前向带树约束）或两段式（先选枝，再对选中叶子跑标准 attention）。

实践案例

案例 1：长文档问答——先路由再细读

假设文档 32 段，问题是「第三节的实验设置用了什么学习率」。

DocTree:
  root
   ├─ ch1 (摘要: 背景)
   ├─ ch2 (摘要: 方法)
   └─ ch3 (摘要: 实验) ← 路由命中
        ├─ §3.1 数据
        └─ §3.2 超参 ← 再命中 → 局部 attention 读出 lr=1e-4

逐部分解释：

root 不直接读 32 段全文，只比较各章摘要与问题的相关分
只展开 ch3，其它章的 KV 可以不加载或只保留摘要向量
叶子上的局部 attention 才真正「盯 token」，成本远小于全文档 O(n²)

案例 2：用掩码表达「只能看祖先与同枝」

# nodes: 先序遍历的 token 段；parent[i] = 父节点下标
def tree_mask(i, j, parent, depth):
    # i 可以看 j：当 j 是 i 的祖先，或同枝近邻
    return is_ancestor(j, i, parent) or same_branch(i, j, parent)

逐部分解释：

is_ancestor：子节点总能看到父摘要（像目录始终可见）
same_branch：同章内允许局部全注意，跨章默认屏蔽
实现上常落成稀疏 attention mask，或「先选枝再对选中叶子跑标准 attention」两段式

案例 3：多跳问题——沿树走两步

Q: 「方法章引用的数据集，在实验章报告了多少准确率？」
step1: 路由 → 方法章叶子 → 抽出数据集名 D
step2: 带着 D 再路由 → 实验章叶子 → 读出 acc=91.2%

逐部分解释：

单次扁平 attention 容易在「方法」和「实验」之间注意力对不齐
树路由把多跳拆成两次「选枝 + 局部读」
每步都可以记录命中节点，方便事后解释「答案从哪一枝来」

踩过的坑

树切坏了：按固定 512 token 硬切，会把一句话劈成两叶，路由再准也读残——优先按标题 / 空行切。
摘要丢关键实体：父节点只存「本章讲实验」却丢掉数据集名，第二跳就路由失败——摘要里要保留专有名词。
枝数开太大：每层 top-k=全部，又退化成近似全注意，又贵又吵——先从小 k（如 2–4）做起。
和 ToT 混用却不共享前缀：推理树每个候选都重算同一段上下文 KV，显存爆炸——分支候选应复用同一 DocTree 的 KV。

补充一句：路由分数若只用表面词重叠，同义改写的问题会选错章——相关分最好带一点语义向量，而不是纯关键词。

适用 vs 不适用

适用：

文档有天然层级（论文、手册、法规），长度 ≫ 模型舒适窗口（如 >8K–32K token）
问题只需少数章节即可回答，愿意用路由换算力
需要可解释「答案来自哪一枝」的检索式问答

不适用：

短文本（<2K）且必须全局依赖——树路由开销大于收益
强全局依赖任务（全文修辞、逐句对齐翻译）——剪枝会伤正确性
完全无结构乱文本又拒绝建摘要——树退化成随机切块

历史小故事（可跳过）

2017：transformer-2017 把 self-attention 做成标配，但也把 O(n²) 写进了默认假设。
2020：longformer-2020 等用滑窗 / 全局 token 稀疏化，先证明「不必人人互看」。
2023：tree-of-thoughts-2023 让推理长成树；同年前后 Tree Cross Attention 等把检索做成对数级树搜。
2024–2026：Tree Attention / Flash Tree-Attention 等把「共享前缀的分支解码」做成系统能力；ToA 把同一直觉用到长上下文推理路由上。

学到什么

长上下文的瓶颈常常是「看谁」而不是「能不能塞进窗口」
树 = 路由 + 局部注意 + 向上汇合，三件套缺一就容易退回暴力全注意
ToA 管上下文结构，ToT 管推理搜索，名字像但层不同，可组合不可混称
切分与摘要质量决定上限——算法再巧，坏树救不回来

关联

tree-of-thoughts-2023 —— 推理树 vs 注意力树
transformer-2017 —— 全员互看的默认注意力
longformer-2020 —— 用稀疏 pattern 做长文档
flash-attention —— IO 友好的精确注意力核
paged-attention —— 推理时 KV 分页管理
attention —— Attention / Transformer 总入口

反向链接

（暂无反向链接）

Tree-of-Attention — 把长上下文拆成树再分支注意

是什么

为什么重要

核心要点

实践案例

案例 1：长文档问答——先路由再细读

案例 2：用掩码表达「只能看祖先与同枝」

案例 3：多跳问题——沿树走两步

踩过的坑

适用 vs 不适用

历史小故事（可跳过）

学到什么

延伸阅读

关联

反向链接