MHC
一、关键概念与背景铺垫
1. 残差连接（Residual Connection）
由 ResNet（2016）提出的深度学习核心架构范式，公式为：
$$ x_{l+1} = x_l + \mathcal{F}(x_l, \mathcal{W}_l) $$
$x_l$ 为第 $l$ 层输入，$\mathcal{F}$ 为残差函数。
其核心优势是恒等映射特性（Identity Mapping）：浅层信号可未经修改直接传递到深层（即公式中 $x_l$ 直接保留），这能避免深层网络训练时的梯度消失 / 爆炸问题，是大规模模型（如 Transformer、大语言模型 LLM）稳定训练的基础。
注： 反向传播用链式法则 $$ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} $$
而： $$ \frac{\partial y}{\partial x} = \frac{\partial(x + F(x))}{\partial x} = I + \frac{\partial F}{\partial x} $$
$I = 恒等矩阵$
- 这是 residual 的“直通梯度”
- 不管 $\frac{\partial F}{\partial x}$ 多小或多大，梯度都会直接加上一份原始梯度
梯度公式完整写出来： $$ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \left( I + \frac{\partial F}{\partial x} \right) = \frac{\partial L}{\partial y} + \frac{\partial L}{\partial y} \cdot \frac{\partial F}{\partial x} $$
- 第一项: $\frac{\partial L}{\partial y}$ → 直通梯度，不会消失
- 第二项: $\frac{\partial L}{\partial y} \cdot \frac{\partial F}{\partial x}$ → $F(x)$ 的梯度
2. 超连接（Hyper-Connections, HC）
近年提出的残差连接扩展方法（以 Zhu 等人 2024 年研究为代表），核心改进是拓宽残差流宽度 + 丰富连接模式：
- 传统残差连接是 “单流” 传递（特征维度为 $C$），HC 将其扩展为 “多流”（维度变为 $n×C$， $n$ 为扩展系数，论文中设为 4）；
- 引入 3 个可学习映射矩阵 ${H}_l^{pre}$，${H}_l^{post}$，${H}_l^{res}$，动态调节不同流、不同层间的特征交互强度，在不增加单元计算量（FLOPs）的前提下提升网络拓扑复杂度。
单层的扩展公式为：
$$ x_{l+1} = {H}_l^{res}x_l + {H}_l^{post{T}} \mathcal{F}({H}_l^{pre}x_l, \mathcal{W}_l) $$
其中：
- $\mathcal{H}_l^{res} \in \mathbb{R}^{n \times n}$：残差流内特征混合的可学习矩阵
- $\mathcal{H}_l^{pre} \in \mathbb{R}^{1 \times n}$：将 $nC$维的流压缩为 $C$ 维层的输入
- $\mathcal{H}_l^{post} \in \mathbb{R}^{1 \times n}$：将 $C$ 维层输出映射回 $nC $维流
二、HC 的优势与缺陷
1. HC 的优势：性能提升的来源
HC 通过 “拓宽残差流 + 多样化连接”，打破了传统残差连接 “单流线性传递” 的局限：
- 多流结构能同时保留更多不同维度的特征信息，减少信息在传递中的损耗；
- 可学习映射矩阵让特征在层间、流间的交互更灵活，能捕捉更复杂的依赖关系（如 LLM 中的长距离语义关联）
- 实证表明，这种设计能带来显著性能提升（论文后续实验中，HC 在 27B 模型的 MMLU、GSM8K 等基准测试中，较传统残差连接提升 4-6 个百分点）
2. HC 的致命缺陷：稳定性与效率的双重瓶颈
尽管 HC 能提升性能，但它的 “无约束设计” 彻底破坏了残差连接的核心 ——恒等映射特性，进而引发两大关键问题：
问题 1：训练稳定性差，可扩展性受限
传统残差连接中，多层传递后信号满足$x_L = x_l + \sum_{i=l}^{L-1} \mathcal{F}(x_i, \mathcal{W}_i)$，浅层的信号 $x_l$ 直接保留，保证信息强度“守恒”。
而 HC 的多层传递公式为：
$$ x_{L} = \left(\prod_{i=1}^{L-l} \mathcal{H}_{L-i}^\text{res}\right) x_l + \sum_{i=l}^{L-1}\left(\prod_{j=1}^{L-1-i} H_{L-j}^{res}\right)H_{i}^{\text{post}\top} \mathcal{F}(\mathcal{H}_i^{\text{pre}} x_i, \mathcal{W}_i) $$
$\prod_{i=1}^{L-l} \mathcal{H}_{L-i}^\text{res}$ 是可学习矩阵的乘积（无约束）。这种乘积会导致信号要么 “爆炸”（强度无限增大）、要么 “消失”（强度趋近于 0），论文实验显示 HC 在 27B 模型训练到 12k 步时，损失会突然飙升，梯度 norm 波动幅度是 mHC 的 10 倍以上，无法支持更大规模（如 100B + 参数）模型的训练。
问题 2：系统开销大
残差流宽度扩展导致每个 Token 的内存访问成本增至 $(5n + 1)C + n^2 + 2n$，GPU 内存占用显著上升。
TODO: Deepseek mhc 论文的 Table 2 截图
三、本篇论文的核心方法 Manifold-Constrained HC
mHC 的核心是将 HC 的残差连接空间投影到特定流形(Manifold)，恢复恒等映射特性，同时通过基础设施优化降低开销。
3.1 流形约束设计：双随机矩阵(Doubly Stochastic Matrices)与 Birkhoff 多面体(Brikhoff polytope)
将 $H_l^{res}$ 约束到双随机矩阵流形（Birkhoff 多面体），满足： $$ \mathcal{P}_{M^{res}}(\mathcal{H}_l^{res}) := {\mathcal{H}_l^{res} \in \mathbb{R}^{n \times n} | \mathcal{H}_l^{res}1_n = 1_n, 1_n^{\top}\mathcal{H}_l^{res} = 1_n^{\top}, \mathcal{H}_l^{res} >= 0 } $$
TODO: 做映射的步骤，对应图 7 和 图 8
人话描述：MHC 的核心是解决了 HC 中的 x 的因数 $H_l^{res}$ 的不稳定问题。我们希望把它变成双随机矩阵。怎么做呢，把它投影到 Birkhoff 多面体来实现，怎么投影呢？利用 Sinkhorn-knopp 算法来实现。这样的话，就可以利用双随机矩阵的多次相乘仍然保持稳定的特性，解决了HC多层嵌套后可能出现的爆炸和消失问题。
3.2 高效基础设施优化
MHC 通过 3 项优化将 n = 4 时的额外时间开销控制在 6.7%，内存空间还较 HC 有优化。
3.2.1 Kernel Fusion -- 减少内存搬运
- 重排 RMSNorm 顺序（先矩阵乘法后归一化），保持数学等价性
- 用 TileLang 框架实现混合精度核，融合多操作（如线性投影、偏置加法），减少内存带宽瓶颈
3.2.2 选择性重计算 -- 减少内存占用
- 前向传播后丢弃 mHC 中间激活值，反向传播时重新计算（仅需存储块首输入 $x_{l0}$）
- 最优块大小 $L_n \approx \sqrt{\frac{nL}{n + 2}}$($L$为总层数)，平衡内存占用与重计算开销
3.2.3 DualPipe 通信重叠 -- 减少通信开销
- 用高优先级流执行 MLP 的 $\mathcal{F}_{post,res}$，避免通信阻塞；
- 解耦重计算与管道通信依赖，利用本地缓存的 $x_{l0}$，减少延迟
四、核心实验结果
4.1 实验设置
- 模型架构：基于 DeepSeek-V3 的 MoE 架构，含 3B/9B/27B 参数模型（n=4）
- 训练配置：序列长度 4096，AdamW 优化器（β=(0.9,0.95)），学习率按步衰减；27B 模型训练 262B tokens，3B 模型额外训练 1.05T tokens；
- 评估任务：8 项下游基准（BBH、DROP、GSM8K、HellaSwag 等），覆盖推理、阅读理解、数学等场景。
4.2 核心实验结果
4.2.1 训练稳定性与效果
- 稳定性：27B 模型中，mHC 无损失骤升（图 5a），梯度范数与基线持平（图 5b），Amax 增益从 HC 的 3000 降至 1.6（图 7）
- 训练效果：mHC 全面超越基线与 HC（表 4），如 BBH（51.0% vs HC 48.9%）、DROP（53.9% vs HC 51.6%），推理能力显著提升。
4.2.2 缩放性分析
- 计算缩放：3B→9B→27B 模型中，mHC 的性能优势持续保持（图 6a），无衰减
- token 缩放：3B 模型训练 1.05T tokens 时，mHC 的相对损失比基线低 1%+（图 6b），验证大规模数据下的有效性。
五、结论与未来方向
5.1 核心结论
- mHC 通过双随机矩阵流形约束，解决了 HC 破坏恒等映射的问题，恢复训练稳定性；
- 基础设施优化实现低开销（n=4时 + 6.7% 时间），支持大规模模型训练；
- 实验证明 mHC 在性能、稳定性、缩放性上全面优于 HC 与基线。
人话总结：看下HC的公式，F 处理的是 1xC 的单流，而输入x则是 nxC的多流。$H_l^{pre}$ 把 x 从 nxC 降到 1xC 给 F 做计算，输出的结果是 1xC，接着再用 $H_l^{post}$ 的转置 nx1 再把结果扩展到 nxC，再与 $H_l^{res}$ x 相加。虽然这样做虽然可以扩展残差流的效果，有点类似于单头Attention和多头 Attention，而因为在多层嵌套场景下 $H_l^{res}$ 的效果不稳定，容易出现爆炸和消失，所以导致训练的稳定性不强。基于这个考虑，MHC 使用 sinkhorn-knopp 算法把 $H_l^{res}$ 转换成双随机矩阵，既保留了多流表达性强的优点，又利用双随机矩阵的特性，即多次相乘后还是双随机矩阵，保证所有元素非负，而且行/列之和都为 1，解决了多层嵌套场景下 $H_l^{res}$ 易爆炸或消失的问题
5.2 未来方向
- 探索其他流形约束（如正交矩阵流形），优化 “可塑性 - 稳定性” 权衡
- 推动宏观架构设计研究，深化拓扑结构对优化与表征学习的影响。
六、代码的实现细节
这里记录下 tilelang仓库中的 mhc 实现 的分析：
a[test_example_mhc.py] --> b[example_mhc_pre.py] a --> c[example_mhc_post.py]
对于 example_mhc_pre.py 来说，它的调用逻辑是这样的：
a[main] --> b[test] b -->|1| c[generate_test_data] b -->|2| d[mhc_pre\n//使用 tilelang 去实现] b -->|3| e[mhc_pre_ref\n// 使用纯 pytorch 实现]
%% tilelang 的核心逻辑 d -->|1| d1[mhc_pre_gemm_sqrsum_tilelang] d -->|2| d2[mhc_pre_big_fuse_tilelang]
%% pytorch 调用的 sinkhorn_knopp 函数 e -->|1| e1[sinkhorn_normalize_ref]
我们的核心研究点就在 d1，d2 和 e1
对于 example_mhc_pre.py 来说，它的调用逻辑是这样的：
a[main] --> b[test] b -->|1| c[generate_test_data] b -->|2| d[mhc_post\n// tilelang 的实现] b -->|3| e[mhc_post_ref\n // pytorch 的实现]
d --> f[mhc_post_tilelang]