一个句子怎样进入 TreeHeap

系列定位:TreeHeap 论文特别篇(2/4)。

本篇只解释当前代码真实执行的数学过程,不把内部节点命名为主语、谓语、摘要或世界模型。

系列导航

  1. 问题、失败与设计演化
  2. 本篇:数学、参数与数据流
  3. 三种子 WMT 与双向 Dreams
  4. 边界、否证条件与复现

1. 先区分参数与状态

回归方程:

$

$ y=wx+b $

$

里面的

$w,b$

是长期学习参数,

$x$

是当前输入,

$y$

是当前输出。

TreeHeap 也必须做同样区分。

共享参数记作

$\theta$

$

$ \theta=\{embedding,Butterfly,FOLD,READ,GRU,output\} $

$

它们保存在 checkpoint 中,被全部训练样本共同使用。

一个具体句子

$x$

经过这些参数后,形成临时状态:

$

$ H_\theta(x)=\left(root_x,details_x,masks_x\right) $

$

$H_\theta(x)$

随句子改变,不是另一份模型参数。

可以这样理解:

theta:长期形成的读写规则
H(x):这套规则对当前句子的实例化

2. WRITE:把离散 token 写入 leaf

输入由三部分组成:

[direction token] + [SentencePiece pieces] + [EOS]

例如中译英方向的第一个 leaf 会写入 zh2en 特殊 token,后续 leaf 才是原句 pieces。

对有效位置:

$

$ x_i=E_{src}(w_i) $

$

当前实现没有额外的位置 embedding。位置差异来自:

  1. leaf 下标;
  2. Butterfly 中由地址决定的配对;
  3. 二叉 FOLD 中不同的左右路径。

不足容量的位置由 mask 关闭。WRITE 只完成离散到连续向量的映射,还没有自动产生高层语义。

3. Butterfly:固定容量内的长距离通信

设 TreeHeap 有

$N=2^D$

个 leaf。第

$s$

轮中,地址

$i$

与地址

$

$ j=i\oplus2^s $

$

通信。

一对状态

$(a,b)$

的可学习 kernel 为:

$

$ b'=b+\alpha_s\tanh(F_\theta(a)) $

$

$

$ a'=a+\alpha_s\tanh(G_\theta(b')) $

$

这里

$F_\theta,G_\theta$

是共享的小型非线性网络。所有地址使用同一套 kernel,不为每一对节点单独学习一张表。

逆运算按相反顺序计算:

$

$ a=a'-\alpha_s\tanh(G_\theta(b')) $

$

$

$ b=b'-\alpha_s\tanh(F_\theta(a)) $

$

所以 Butterfly 可以改变坐标,又不要求丢失输入。

每轮有

$N/2$

对,共有

$\log_2N$

轮:

$

$ \text{pair operations}=\frac{N}{2}\log_2N $

$

这是

$O(N\log N)$

的稀疏通信,不分配

$N\times N$

的稠密注意力矩阵。

4. FOLD:把两个 child 组织成 parent 与 detail

Butterfly 之后,TreeHeap 开始逐层 FOLD。对于左右状态

$(l,r)$

$

$ d=r-P_\theta(l) $

$

$

$ p=l+U_\theta(d) $

$

其中:

  • $P_\theta$预测右侧状态;
  • $d$保存没有被预测到的残差;
  • $U_\theta$决定残差怎样更新到 parent;
  • $p$进入更高一层。

当前 Update 为:

$

$ U_\theta(d)=0.5d+0.5\tanh(\widetilde U_\theta(d)) $

$

可学习部分从零初始化。训练开始时,它等价于稳定的

$0.5d$

更新;训练随后可以改变信息如何向 parent 上导。

5. UNFOLD:为什么它能恢复 child

已知 parent

$p$

与 detail

$d$

$

$ l=p-U_\theta(d) $

$

再计算:

$

$ r=d+P_\theta(l) $

$

就能恢复左右状态。

递归执行后,

$N$

个 leaf 被组织为:

1 个 root
+ 第 0 层 details
+ 第 1 层 details
+ ...
+ masks

detail 总数仍为

$N-1$

。因此当前多分辨率状态不是文件压缩:信息只是被重新组织,并没有自动减少存储量。

6. READ:Decoder 如何决定读取哪个深度

Decoder 每生成一个 token,都从 root 开始分配概率质量。

对节点

$n_i^{(k)}$

,计算停止概率:

$

$ p_{stop}(i,k,t)= \sigma\left(S_\theta\left[q(h_t),n_i^{(k)}+e_k\right]\right) $

$

如果不停止,剩余质量进入左右 child:

$

$ p(c\mid i,t)= \operatorname{softmax}_c \left(\frac{B_\theta(h_t)^\top n_c}{\sqrt m}\right) $

$

到达节点

$i$

的质量为

$m_i$

时:

$

$ m_i^{stop}+m_{left}+m_{right}=m_i $

$

所以概率没有在递归过程中凭空增加。所有停止节点的加权和形成当前上下文

$c_t$

7. Decoder 如何生成下一个 token

Decoder 把上一个目标 token、当前 TreeHeap context 和历史隐状态放进 GRU:

$

$ h_{t+1}=GRU([E_{tgt}(y_t),c_t],h_t) $

$

再输出词表概率:

$

$ p(y_{t+1})=softmax(W_o[h_{t+1},c_t]) $

$

训练时使用 teacher forcing:第

$t$

步输入真实的

$y_t$

。自由生成时,输入模型自己上一步选择的 token。

8. 梯度到底从哪里来

唯一语言目标是目标 token 的交叉熵:

$

$ \mathcal L=-\sum_t\log p_\theta(y_t\mid y_{\lt t},H_\theta(x)) $

$

如果正确 token 的概率太低,loss 就升高。反向传播依次经过:

词表输出
  -> GRU
  -> recursive READ
  -> UNFOLD levels
  -> root + details
  -> lifting FOLD
  -> Butterfly
  -> source embedding

这条链就是私有协议形成的物理通道。模型没有内部节点标签,也没有人工告诉它哪个 node 是“食物”或者“主语”。

9. 这套结构哪里是确定的,哪里是学习的

类型 内容
确定性结构 heap 地址、XOR 配对轮次、FOLD/UNFOLD 计算顺序、概率质量守恒
可学习参数 embedding、Butterfly kernel、Predictor、Update、stop、branch、GRU、输出层
样本临时状态 leaf、root、details、READ context
最终监督 目标 token 交叉熵

这一区分很重要。Butterfly 地址图是人工给定的归纳偏置;具体传递什么信息由梯度学习。FOLD 的逆式由数学定义保证;不同深度最后承担什么任务,必须由实验观察。

10. 当前计算代价

Butterfly 是

$O(N\log N)$

,FOLD/UNFOLD 是

$O(N)$

。当前 recursive READ 每个输出时间步会访问总计小于

$2N$

个节点,所以生成约为:

$

$ O(TN) $

$

其中

$T$

是输出长度。

因此,“稀疏”目前只是结构性质,不等于工程上已经更快。是否节省 GPU 小时,必须用吞吐、显存和训练成本实际测量。

11. 本篇结论

一个句子进入 TreeHeap 后,没有被压成一个神奇 root。它经历的是:

WRITE
-> Butterfly communication
-> reversible FOLD
-> H(root, details, masks)
-> recursive READ
-> token distribution
-> cross-entropy gradient

下一篇将检查这条链到底产生了什么证据,以及哪些漂亮数字其实不能被解释得太强。

下一篇:三种子 WMT 与双向 Dreams


License: GPLv3。