中文

块级可微Sinkhorn注意力:带间隙感知粘土桥的尾部精炼梯度

机器学习 2026-05-21 v2 计算与语言

摘要

我们研究在 TPU 硬件上进行长上下文平衡型熵最优传输(OT)注意力问题,采用已停止基座、固定深度尾部精炼替代方案。经过 T 步的已停止 Sinkhorn 求解后,我们展开短暂的精炼尾部并精确地对该替代方案进行微分。在报告的 R=2 TPU 路径中,反向传播包含四个阶梯计划因子。我们证明了一个确切的单参考瓷砖时刻表:R=2 的 score cotangent 是单个参考计划瓷砖乘以一个由向量 cotangent 和双重差异构建的显式修饰场。这导致块级成本 O((T+R)LW)O((T+R)LW)O(Ld)O(Ld) 输入存储,以及针对固定头维数 dd 和带宽 WWO(L)O(L) 额外 HBM 用量。我们还将当前 \texttt{dustbin\_block} 路径形式化为该单位目标替代方案在增强支持上的相同单活性粘土桥;该桥是代数形式的,不声称建立一般 KL 不平衡或任意容量间隙模型。我们提供局部替代方案偏差界、事后偏差证书以及严格正活跃块的射影收缩证书。在合成掩码问题上,优化的内核与相同中心替代方案的精确自动微分相符,误差在 10510^{-5}--101010^{-10} 范围内。在 TPU v6e-8 上,四配置 Pfam 屏幕完成端到端,提升的平衡 R=2 运行在三小时预算内维持约每秒 8.5 个示例,达到步骤 1437。持留 Pfam 测试碎片相对于步骤 0 的重建从 5.57 提升至 2.05,稀疏 CE 从 5.53 提升至 5.30(CE 诊断性地记录而非直接优化);目标-中心对齐指标在实际效果中未见实质性提升,确定性对角线参考在这些指标上仍然更强。

关键词

引用

@article{arxiv.2605.08123,
  title  = {Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge},
  author = {Dylan Forde},
  journal= {arXiv preprint arXiv:2605.08123},
  year   = {2026}
}