技术摘要：加速流形约束超连接的 Birkhoff 投影

1. 问题陈述

流形约束超连接（mHC）已成为超连接（HC）的一种原则性扩展，旨在稳定深度神经网络的训练。虽然 HC 通过扩展残差流宽度来实现多路径信息流，但其不受约束的残差混合矩阵可能导致范数爆炸或消失，从而引起训练不稳定。mHC 通过将残差映射投影到 Birkhoff 多胞形（即双随机矩阵集合）来解决这一问题，从而保持信号强度并确保稳定的梯度传播。

然而，目前的 mHC 实现依赖于 Sinkhorn–Knopp 算法来进行迭代投影。这种方法存在三个关键局限性：

1. 计算开销： Sinkhorn–Knopp 的迭代性质引入了显著的延迟，尤其是在数百万个训练步骤中针对每个 token 执行时。
2. 内存成本： 自动微分框架中的反向传播通常需要展开整个迭代过程，这会导致巨大的内存消耗。
3. 准确性与稳定性： 在具有挑战性的输入（例如大模量矩阵）下，该算法可能收敛缓慢。使用固定的较小迭代次数（例如 20 次）往往会产生不准确的投影，无法满足双随机约束，从而削弱了对 mHC 稳定性至关重要的范数控制保证。

本研究专注于在实践中非常重要的 $4 \times 4$ Birkhoff 投影设置（对应于扩展率 $n=4$），旨在开发一个端到端的加速框架，以同时提高前向投影的准确性和反向微分的效率。

2. 方法论

所提出的框架利用了 Birkhoff 投影问题与 熵正则化最优传输（OT） 之间的数学等价性。通过利用 $4 \times 4$ 情况下的小规模结构，作者开发了一种结合隐式微分和硬件感知实现的二阶优化方法。

2.1 前向传递：通过对偶形式的牛顿法

作者将 Birkhoff 投影重新表述为一个在 $\mathbb{R}^3$ 上的无约束凸优化问题。

• 降维： 熵正则化 OT 形式的对偶问题通常涉及 $n+m$ 个变量。然而，由于对偶变量中的冗余性（在标量平移下的不变性），该问题可以简化为 $m-1$ 个自由变量。对于 $4 \times 4$ 的情况（$n=m=4$），这使得问题简化为 3 维无约束凸最小化。
• 二阶求解器： 作者推导了目标函数的梯度和 Hessian 矩阵的闭式表达式。这使得使用 牛顿法 成为可能，牛顿法具有二次收敛特性。与需要多次迭代才能收敛的 Sinkhorn–Knopp 不同，牛顿法通常只需极少的迭代即可收敛，即使对于大模量输入也能提供高精度。
• 数值稳定性： 为了防止指数计算过程中的溢出，实现采用了最大值减法技术（在应用 Softmax 函数之前，通过将输入矩阵行减去其最大值进行平移）。

2.2 反向传递：隐式微分

为了消除展开迭代求解器带来的内存开销，作者使用 隐式微分 取代了标准的反向传播。

• 解析推导： 利用隐函数定理，推导出了投影相对于输入矩阵的导数。这种方法可以直接计算精确梯度，而无需存储前向传递过程中的中间状态。
  Closer
• Hessian 复用： 对于 $4 \times 4$ 的情况，在前向传递的最后一次牛顿迭代中计算出的 Hessian 矩阵与隐式微分步骤中所需的矩阵是完全一致的。这允许复用已计算的值，进一步降低了计算成本。

2.3 GPU 高效实现

该框架使用专门设计的 Warp 级 CUDA 核函数 实现，旨在最大化并行度并最小化内存 I/O。

• 寄存器级原语： 实现方案在单个 GPU Warp（32 个线程）内同时处理两个 $4 \times 4$ 矩阵。
• Warp 内洗牌（Intra-Warp Shuffling）： 该算法并非使用缓慢的全局或共享内存进行通信，而是利用 __shfl_xor_sync 指令进行寄存器级的数据洗牌。这使得线程能够交换数据并计算列和（梯度计算中的关键步骤），且每矩阵仅需四个基本指令，完全避免了内存 I/O。
• 吞吐量： 这种设计实现了高吞吐量和极低的开销，特别适用于 mHC 应用中常见的批处理小矩阵场景。

3. 核心贡献

1. 用于前向传递的牛顿法： 将 Birkhoff 投影重新表述为 3D 无约束凸问题，并通过带有闭式梯度和 Hessian 的牛顿法进行求解，实现了比迭代 Sinkhorn 方法更优的二次收敛性和更高的准确度。
2. 用于反向传递的隐式微分： 利用隐函数定理推导出梯度的解析表达式，实现了在不存储中间迭代值的情况下进行精确梯度计算，从而显著降低了内存使用量。
3. Warp 级 CUDA 核函数： 特化的 GPU 实现，通过仅使用寄存器级原语和 Warp 内洗牌处理每个 Warp 中的两个 $4 \times 4$ 矩阵，消除了全局和共享内存的 I/O，以实现最高效率。

4. 实验结果

作者将提出的求解器（记作 mHC-proj）与七个具有代表性的开源基准测试进行了对比，包括纯 PyTorch Sinkhorn、基于 Triton 的 Sinkhorn 以及现有的 TileLang 实现。实验是在 NVIDIA RTX 6000 Ada Generation GPU 上进行的。

• 准确度：
  • 小模量输入： 对于小模量输入（如 $N(0,1)$），所有方法表现都相对较好，尽管所提方法显示的中间误差更低。
  • 大模量输入： 对于大模量输入（如 $N(0, 10^2)$），由于在固定迭代限制内收敛不足，基于 Sinkhorn 的方法表现出显著误差（中值误差高出多达 3 个数量级）。相比之下，mHC-proj 保持了极高的准确度，其误差比前者小几个数量级，确保了可靠的双随机投影。
• 速度：
  • 前向传递： 所提求解器相对于基准测试实现了显著的加速。
  • 端到端（前向 + 反向）： 引入反向传递过程凸显了隐式微分的效率。在 Batch Size 为 128K 时，mHC-proj 的 CUDA 实现相比于 DeepSeek TileKernels 实现（一个具有代表性的基于 Sinkhorn 的基准）实现了 超过 20 倍的加速。
  • 可扩展性： 随着 Batch Size 的增大，性能差距进一步扩大，证明了该框架有效利用大规模并行性的能力。

5. 重要性与主张

论文声称，准确、二阶且具备硬件感知能力的投影求解器是扩展类 mHC 架构的关键赋能因素。通过解决标准 Sinkhorn–Knopp 方法中的计算和准确度瓶颈，所提出的框架：

• 恢复了稳定性保证： 它确保了残差混合矩阵即使在挑战性输入下也是真正的双随机矩阵，从而维护了对稳定训练至关重要的范数控制特性。
• 实现了高效性： 它大幅降低了训练和推理过程中频繁进行流形投影带来的计算和内存开销。
• 提供了切实可行的解决方案： 通过专注于 $4 \times 4$ 设置（这是当前 mHC 实现的标准扩展率），这项工作提供了一个立即可用的解决方案，提升了最先进超连接架构的可靠性与速度。