Gated Tree Cross-Attention for Checkpoint-Compatible Syntax Injection in Decoder-Only LLMs¶

会议: ACL 2026
arXiv: 2602.15846
代码: https://github.com/Pineandgrass/GatedTreeCrossAttention
领域: LLM 架构 / 句法注入 / 检查点兼容
关键词: GTCA、句法注入、checkpoint-compatible、constituency chunk memory、token update mask

一句话总结¶

作者给冻结的 decoder-only LLM（Qwen-2.5-7B、Llama-3-8B）外挂一个 Gated Tree Cross-Attention 侧支路 —— 离线 Berkeley parser 预算 constituency 树并按高度索引成 chunk memory，token 隐状态通过 head-wise 门控 cross-attention 读它得到残差更新，再配合 token update mask + 三阶段训练防止干扰；BLiMP 准确率从 78.58/79.95 提升到 83.12/84.61，同时 MCQA、HellaSwag、WinoGrande 完全不退步。

研究背景与动机¶

领域现状：decoder-only LLM 在聚合 benchmark 上分数很高，但在细粒度语法压力测试（BLiMP、HANS、CoLA）上常常翻车 —— 用户体验是"同一意思、不同写法、得到完全相反的答案"，这种脆性还会级联到下游推理。probing 工作已经反复证明 LLM 内部隐状态能恢复出依存几何（Hewitt & Manning 2019），即句法是"被编码"的。

现有痛点：① "encoding ≠ usage"，可恢复不等于真在用 —— BLiMP 上 GPT-2 离人类水平仍远；② 主流的句法注入方法（修改 attention bias、tree-RNN、加 dependency-aware attention）通常需要重写架构或全量再训练，对已经训好的 LLM 不友好，naive 注入还会触发灾难性遗忘；③ LoRA/QLoRA 之类参数高效方法不改变 attention 结构，没法引入"显式树结构"这种 inductive bias。

核心矛盾：想给一个已训好的 checkpoint 引入显式句法层级信号，但又不能动 backbone 也不能干扰 pretrained competence；同时不能影响 likelihood-based MCQA 评分（如果改了选项 token 的隐状态，选项之间的相对似然就被污染了）。

本文目标：构造一条"可装可卸、随时旁路"的句法注入路径，让模型自己学会什么时候信、信多少，并能在保持 backbone 不变的前提下，在句法 benchmark 上稳定提升，同时不伤其他能力。

切入角度：作者把 constituency parse 树离线算好并按 hash 缓存（消除训练时 parser 开销），按 tree height 切成 chunk memory，并用 layer-aligned 方式喂给对应 transformer 层 —— 高层喂高 chunk、低层喂叶 chunk，让层级 inductive bias 与 transformer 自然分层对齐。

核心 idea：把树结构作为一份外部缓存 + 门控注意力源塞给 decoder-only LLM —— 像 RAG 一样，但 retrieve 的是句法 chunk 而不是文档 —— 然后用 head-wise gate 把"用还是不用"完全交给模型自学。

方法详解¶

整体框架¶

GTCA 是一条挂在 decoder-only transformer 上的 forward wrapper 侧支路。对每个输入：① 离线用 Berkeley Neural Parser 算 constituency 树并缓存（hash 索引）；② parse-tree encoder 把树按 height 切，mean-pool 各 chunk 的 span token embedding 再过 height-specific 投影 + LayerNorm 得到 chunk memory \(C^\ell\)；③ 在 transformer 第 \(\ell\) 层，pre-update 隐状态 \(H_{\text{pre}}^\ell\) 作 query 通过 head-wise gated cross-attention 读 \(C^\ell\)，输出残差 \(\Delta H^\ell\)；④ 用 token update mask 把 \(\Delta H^\ell\) 仅施加到 question + answer field（option token 强制 mask=0），得到 \(H_{\text{post}}^\ell\) 送下一层。所有 backbone 参数可保持冻结或低秩微调，GTCA 分支随时可关。

关键设计¶

Height-aligned Chunk Memory（高度对齐的 chunk 缓存）:
- 功能：把 constituency 树按高度切成多层 chunk，与 transformer 层数一一对应，让低层喂局部、高层喂全句结构。
- 核心思路：定义 chunk 高度 \(h(u)=D-\text{depth}(u)\)，叶 token 高度为 0、根节点高度最大；对每个 chunk 用 span mean-pool \(p_u=\text{MeanPool}(E^i, i\in S(u))\)，再过 height-specific 投影矩阵 \(W_{h(u)}\in\mathbb{R}^{d\times d}\) 加 LayerNorm 得到 \(c_u\)。第 \(\ell\) 层的 chunk memory \(C^\ell\) 只包含 \(h(u)=\min(\ell, D)\) 的 chunk，按左到右 BFS 顺序最多保留 \(K=64\) 个。
- 设计动机：transformer 已有"低层抓局部句法、高层抓全句语义"的归纳偏置（多篇 probing 工作的共识），让 chunk memory 沿同样的层级对齐能让模型最容易把外部树信息接进去，也避免高层 token 反复被低层 chunk 噪声打扰。
Head-wise Gated Cross-Attention（头级门控的旁路注意力）:
- 功能：用 cross-attention 让 token 状态读 chunk memory，但门控决定每个 head 该信任多少。
- 核心思路：标准 cross-attention \(Q=H_{\text{pre}}^\ell W_Q^\ell\), \(K=C^\ell W_K^\ell\), \(V=C^\ell W_V^\ell\)，外加一个 head-wise 门 logit \(G^\ell = H_{\text{pre}}^\ell W_G^\ell\)；attention 输出乘上 sigmoid 门 \(\text{Gated\_Attn}^\ell = \text{Attn}^\ell \odot \sigma(G^\ell)\)，每个 head 一个标量门（不是 element-wise），参数更少、训练更稳。再加 causal mask 屏蔽右边界超过当前 token 的 chunk，保证自回归性。最后 \(\Delta H^\ell = \text{Merge}(\text{Gated\_Attn}^\ell)W_O^{ca,\ell}\) 作为残差注入。
- 设计动机：硬注入（无门控）等于强制 backbone 在每层都吸收 chunk 信号，容易破坏 pretrained 表征；head-wise gate 让每个 head 学到"我现在的查询需要句法吗"，等价于让模型自学一个 sparse routing —— 这把"显式树是 hard constraint"变成了"显式树是可选 prior"。
Token Update Mask + 三阶段训练（双安全机制）:
- 功能：限制干扰范围（空间）+ 控制干扰时机（时间），共同防止 catastrophic forgetting。
- 核心思路：① 空间维度：定义二值 mask \(m_{\text{tok}}\in\{0,1\}^n\)，让 \(H_{\text{post}}^\ell \leftarrow H_{\text{pre}}^\ell + \alpha_{\text{struct}}(m_{\text{tok}} \odot \Delta H^\ell)\)；MCQA 输入里 option token 强制 \(m_{\text{tok}}=0\)，因为 likelihood-based 评分依赖 option token 的 logprob，改动它们等于直接改答案概率分布。② 时间维度：三阶段训练 schedule 先训 GTCA 分支自身、再联合微调 backbone 中受影响的子模块、最后逐步开放，避免冷启动时大 \(\Delta H\) 把已学好的 token 状态推飞。
- 设计动机：解决"checkpoint-compatible"这个核心约束 —— 任何对 hidden state 的改动都可能毁掉 pretrained 能力；mask 锁住空间作用域，stage schedule 锁住时间作用域，两者结合让"加了 GTCA 后 MCQA 不退步"成为可能。

损失函数 / 训练策略¶

继续训练用语言建模 loss + MCQA-friendly 格式；三阶段 schedule 第一阶段冻结 backbone 只训 GTCA 投影与 gate，第二阶段开放与 GTCA 直接交互的子模块，第三阶段全量但用低学习率收敛。chunk 数量上限 \(K=64\)；scaling factor \(\alpha_{\text{struct}}\) 控制残差幅值；offline parse 用 Berkeley Neural Parser，按 token ID hash 缓存，训练时零 parser 开销。

实验关键数据¶

主实验（BLiMP 句法能力）¶

模型	Baseline BLiMP	+ GTCA	\(\Delta\)
Qwen-2.5-7B	78.58	83.12	+4.54
Llama-3-8B	79.95	84.61	+4.66

类别	任务	Baseline	+ GTCA	说明
句法	BLiMP	78.58-79.95	83.12-84.61	提升 4-5 pp
句法	CoLA (GLUE)	—	一致提升	语法接受度判断
MCQA	CLOTH	—	持平或微升	完形填空，不退
MCQA	MMLU	—	持平或微升	知识 QA，不退
常识	HellaSwag	—	持平	常识续写
常识	WinoGrande	—	持平	共指消解

消融实验¶

配置	关键指标	解读
Full GTCA (height-aligned + gated + mask + stage)	BLiMP 83.12	完整模型
w/o head-wise gate (硬注入)	显著掉点	backbone 表征被污染
w/o token update mask（也改 option token）	MCQA 退步	option likelihood 漂移
w/o 三阶段 staged training	训练不稳 / 性能掉	冷启动大 \(\Delta H\) 毁 pretrained
一次性投影（替代 height-specific \(W_{h(u)}\)）	BLiMP 略降	层级耦合是必要的

关键发现¶

门控是注入成功的关键：head-wise gate 让模型自己决定每个 head 该不该信任 chunk memory，相比硬注入既保留 pretrained 能力又能选择性使用结构 —— 这把"显式句法 vs 隐式 LLM 能力"的 trade-off 从架构选择问题变成了自学习问题。
option token 必须 read-only：MCQA 任务下，把句法更新施加到 option token 上会改变它们的 logprob，从而漂移答案选择；这条工程经验对所有想在 likelihood-based MCQA 上做继续训练的方法都有警示意义。
layer-aligned chunk memory 给出可解释的层级利用：UUAS probe 显示注入 GTCA 后内部隐状态的 unlabeled undirected attachment 一致性增强，且高层依赖更高 chunk、低层依赖叶 chunk —— 与 transformer 自身的句法分层共识吻合。
句法提升不以广义能力为代价：所有 MCQA / 常识任务持平或微升，说明三道安全机制（gate + mask + staged）确实成功隔离了干扰范围。

亮点与洞察¶

"句法当作外部检索源"的范式：把 parse tree 视为可缓存、可旁路、可门控的外部记忆，使句法注入和 RAG 在工程哲学上对齐 —— LLM 永远不被改写，外部知识源永远是 hot-swappable。这套范式可以原样套到"形态学外挂"、"逻辑表达式外挂"、"领域本体外挂"。
checkpoint-compatible 是工业最关心的属性：现代 LLM 训练成本极高，任何要求"重写 backbone 或重新预训练"的方法在实际部署里几乎都过不了立项；GTCA 用一条 forward wrapper + 三道安全机制让句法增强进可攻退可守，这种"加挂式"思路值得所有后续可微外挂模块借鉴。
token update mask 是被忽视的关键 detail：很多继续训练方法都没区分"哪些 token 可以被改、哪些不能"，本文用 MCQA option token 这个具体例子展示了 hidden state 改动如何泄漏到 likelihood-based 评分，是工程实践的好教材。
height-specific 投影 + layer alignment：把"tree height ↔ transformer layer"做硬绑定是一个很自然但少有人系统验证的选择；本文给出消融证据并配合 UUAS probe 解释，比直觉断言更扎实。

局限与展望¶

只在两个 ~7-8B 的 decoder-only 上做：14B/70B、MoE、混合架构未测；不同 backbone 的 attention 几何不同，GTCA 的 height-alignment 是否还适用未知。
依赖外部 constituency parser：parser 的错误会被缓存并永久注入，作者没有讨论 parser 噪声鲁棒性；对低资源语言或非英语，parser 质量本身就是瓶颈。
训练阶段 chunk memory 缓存对存储和工程链路有要求（hash 索引、span alignment）；对在线 / streaming 场景如何重算缓存未给方案。
BLiMP 是英语句法压力测试，提升 4-5 pp 看上去显著但仍远低于 95+ 人类水平；GTCA 对 long-range syntactic dependency（如 island、binding）是否有提升、提升多少，未在文中拆解。
没有与 LoRA / Adapter 等 PEFT 方法做"参数预算等价对比"，GTCA 加的可训参数量是多少、与同等预算的 LoRA 对句法增益谁更高，缺少对照。

评分¶

新颖性: ⭐⭐⭐⭐ 把树结构外挂 + 头级门控 + 双安全机制组合是新颖且工程上扎实
实验充分度: ⭐⭐⭐⭐ 两个 backbone + 6 个 benchmark + UUAS probe + 消融，但 backbone 规模偏小
写作质量: ⭐⭐⭐⭐ 方法部分公式与符号清晰，三道安全机制讲得很透
价值: ⭐⭐⭐⭐ checkpoint-compatible 这条线对工业部署很实用；句法外挂范式有迁移潜力