我构建了 ALHR：一种基于树的稀疏注意力系统，在保持准确率的同时实现次平方级推理。[P]

Reddit r/MachineLearning 论文
sparse-attention efficient-inference transformers kv-cache-compression hierarchical-routing memory-efficiency

摘要

作者介绍了 ALHR（Adaptive Learnable Hierarchical Routing，自适应可学习分层路由），一种基于静态二叉树与可学习函数的稀疏注意力系统。在 MQAR 测试中，每个查询平均仅读取约 30 个键（相比之下稠密注意力需要读取 512 个），实现了约 35.3 倍的 KV 压缩，同时保持接近的准确率，并将推理复杂度降至 NlogN，显存占用线性扩展。

ALHR——Adaptive Learnable Hierarchical Routing（自适应可学习分层路由），使用静态二叉树和可学习函数来最小化需要读取的键的数量。不过在训练的第一阶段确实使用了一个稠密教师模型。 **MQAR 测试（1024 个 token）** - 稠密注意力每个查询平均读取的键数：512 个 - ALHR 每个查询平均读取的键数：30 个 - 稠密注意力的 Top-1 准确率：94.9% - ALHR 的 Top-1 准确率：92.1% - 稠密注意力的 KV 压缩率：1x（读取 100%） - ALHR 的 KV 压缩率：35.3x（仅读取 2.83%） - 稠密注意力的峰值显存：57 MB（随序列长度平方级增长） - ALHR 的峰值显存：422 MB（线性扩展） - ALHR 的缓存压缩率：100% 实际运行日志和所用的 Kaggle notebook 位于仓库的 logs 文件夹中。 **局限性：** - 全规模测试尚未完成 - 该模型的训练仍然是平方级复杂度，但推理复杂度为 NlogN（如仓库日志中所示） 很希望听到大家的意见！ ALHR 仓库：https://github.com/vdev-ctrl/Adaptive-Learnable-Hierarchical-Routing-
查看原文

相似文章

学习跳跃块：自我发现的超度量路由用于硬件加速稀疏注意力

Reddit r/artificial

本文介绍了动态超度量注意力（Dynamic Ultrametric Attention），这是一个框架，其中Transformer在训练期间学习每头块稀疏路由拓扑，然后在推理时将这些拓扑卸载到自定义的Triton块稀疏内核上，与密集注意力相比，实现了高达28倍的加速和98.4%的内存减少。