vdev-ctrl/Adaptive-Learnable-Hierarchical-Routing-

来源：https://github.com/vdev-ctrl/Adaptive-Learnable-Hierarchical-Routing-

ALHR（自适应可学习分层路由）：在不显著损失准确率的情况下实现次二次复杂度的树注意力，以稠密参考模型为基准分三个阶段训练（测试版）

一种利用静态二叉树和可学习函数，在不显著损失准确率的情况下实现次二次复杂度推理的注意力机制。 我使用静态二叉树来组织键矩阵，其中每个分支都以其下方所有键向量的总和作为索引。

简要概述：

这是一种自注意力的替代方案，使用 O(NlogN) 次运算，而不是 O(N^2) 次运算，且不会造成任何显著的准确率损失，同时使用更少的内存。 以下是基本的 MQAR 推理结果：

注意： 16k 时的准确率下降，并不是由某种导致准确率降幅随词元数量增长而扩大的架构缺陷造成的，正如我们所观察到的，1k 和 4k 词元时的准确率降幅相同。

原因是我们的模型太小，而我没有足够的资源来扩大模型规模，因此你可以看到，稠密模型和 ALHR 在 16K 时都受到了影响。

另外，请观察这里的显存增长情况

对于稠密模型，词元数量翻倍时，峰值显存变为原来的 x1.96，而树模型则为 1.16

当词元数量增加到四倍时，稠密模型的峰值显存变为原来的 x2.9，而我们的模型则为 x1.56

ALHR 峰值显存随规模增长的效率高于稠密模型。

你可以在 logs 文件夹中核查所使用的 Kaggle 脚本和原始结果。

推理阶段

输入被拆分为查询和键之后，会基于键矩阵构建并存储静态二叉树。每层只执行一次这个过程。

此时，我们的可训练函数——预算预测器——就派上用场了。预算预测器是一个可训练函数，它接收当前查询向量和所有向量的全局总和，输出该查询所需的键数量（它为查询输出一个叶节点组，用来表示该查询所需的键数量）。

获得这些信息后，我们首先将查询与它自身的键向量及其邻居进行匹配搜索（这一系统类似于 Treeformer）。 为此，它会使用先前与自身邻居匹配成功的查询的数据。

邻居搜索完成后，如果查询的叶节点组仍未得到满足（意味着还有更多键需要搜索），查询就会沿树向下遍历，将自身与各分支进行比较，并沿相似度更高的分支下降。它向下遍历的分支数量由该查询还需找到多少个键决定（从叶节点组对应的数量中减去已找到的邻居键数量）。

每个查询找到各自的键后，就会执行标准自注意力的点积与求和运算。

将沿树向下遍历的过程计入后，推理的运算复杂度约为 O(NlogN)

训练阶段

这是该机制最重要的部分，也是开销最大的部分。它分为 3 个阶段。

阶段 1：在这个阶段，树模块完全不参与，而是使用数据训练一个稠密模型和一个 FFN（前馈网络）。正是这一步使训练具有二次复杂度。这个稠密模型充当教师。

阶段 2：提取训练好的稠密模型的注意力矩阵，并标注每个查询的 top-k 个键。仅使用这些数据，不使用任何外部数据，来训练我们的树注意力模块。稠密模型就在这个阶段教导树模型。这里不会训练到完全收敛，因为稠密矩阵与稀疏矩阵存在根本差异，而且我们也不希望进行精确拟合。

阶段 3：将树模型与 FFN 一起训练。不是一个新的 FFN，而是阶段 1 中连接到稠密模型的同一个 FFN。由于我们在这个阶段不再使用稠密教师模型，而是直接使用数据训练，树模型实际上可以学会超越其教师（在下面的 MQAR 测试中，它确实展示了这一点）。

在阶段 2 和阶段 3 中，树模块内接受训练的函数是查询/键拆分器和预算预测器。 通过选择树中仅包含过去键向量的一个或多个分支，并将其视为整棵树，可以保持因果性。

结果与变量

以下内容也存放在仓库的 logs 目录下。

理想情况下，我会在 1、4、16k 词元下运行 MQAR 和 Tinystories 测试来进一步验证，但就目前情况而言，我只有 MQAR 测试的结果。

完整日志和可复现的 Kaggle 脚本位于 logs 文件夹中

以下是结果：

1024 词元标准 MQAR 测试

现在你可以看到，我们只读取少量键，就能达到接近稠密模型的准确率。

事实上，在我之前的一些测试中，我们的树模型超过了它自己的教师。不过，我目前没有足够的资源在这一规模下复现该结果，但理论上是可能的。

你会注意到，所使用的显存明显更多，但随着词元数量增加，稠密模型的峰值显存会急剧增加，而我们的模型则呈线性增长。

我们也可以从这里验证推理具有次二次复杂度，约为 Nlog N，日志中也能看到这一点。

再说一次，我想在最多 16k 词元的范围内进行这项多随机种子测试，但我没有足够的资源。（iPhone 16e 和 Kaggle 😭）

这些是专有结果。

你可以在 logs 文件夹中核查我为获得这些结果所使用的数据集和 Kaggle 提示词。

变量：影响训练运行的可变因素包括：阶段 1 的目标值，阶段 1、2、3 的训练步数，top-K 词元数量（目前 k=16 的结果最好），以及预算预测器的上限。

我没有时间，也没有资源来测试所有组合，但这些因素确实能够提高准确率。

这些变量是

阶段 1——目标：0.95。训练达到 0.95 时停止 Top k 词元——16，在阶段 2 中，我们使用自注意力矩阵中排名前 16 的词元来训练 ALHR 最后，最大键数 = 96，最大预算类别 = 32

注意：如果你查看日志，会发现我们的树注意力系统耗时几乎是稠密模型的 10 倍，这是因为它还没有进行内核优化，完全使用 PyTorch 运行。

论文

我目前正在撰写论文，并准备将其上传到 Zenodo，但由于资源不足，还要兼顾大学生活，这个过程一直很艰难。

贡献者

这是一个个人项目（我只是个 18 岁的大学生，哈哈）。只有我、我的 iPhone 16e 和 Kaggle 笔记本。我使用 Claude sonnet 5.5 生成了代码。

捐赠与赞助

请捐款支持进一步开发，以及我脑海中更多的想法。 （我已经厌倦了在一部小手机上做这些事情 😭）

UPI：vdhiroshnin@okhdfcbank

请我喝杯咖啡：https://buymeacoffee.com/vdhiroshnif

捐赠者可以查看开发日志和开发过程，以及一些我不得不暂时搁置的想法。

联系我

电子邮箱：[email protected]

引用

TreeFormer