Manifold-Constrained Hyper-Connections (mHC)

mHC is a research implementation of the Manifold-Constrained Hyper-Connections (DeepSeek) as a drop-in variant of Hyper-Connections. It implements a layer update formula:

x_{l+1} = H_l^{res} x_l + H_l^{post,T} F(H_l^{pre} x_l, W_l)

Key mathematical constraints implemented include:

• H_res: A doubly stochastic matrix (Birkhoff polytope) where entries are $\ge 0$ and rows/columns sum to 1, achieved via Sinkhorn-Knopp.
• H_pre, H_post: Non-negative mixing maps.

This implementation uses static per-layer matrices where H_res_logits are projected via Sinkhorn, and H_pre_logits/H_post_logits are mapped to non-negative weights (e.g., via softmax).

The repository provides several configuration patterns for combining Value Residuals (vRes) with mHC:

• vRes only: Use configs matching train_fineweb10B_vres*.py.
• vRes + mHC: Use configs matching train_fineweb10B_vres_mhc*.py.
• vRes + mHC (Constrained): Use configs matching train_fineweb10B_cvres_mhc*.py. This variant enables v_residual_constrained=True, which uses convex mixing via softmax.

To run training examples, navigate to the examples/nanogpt/ directory. You can use various configuration files to test different model architectures (6-layer vs 48-layer) and research variants (HC, mHC, vRes, etc.).

Use torchrun for multi-GPU training, adjusting --nproc_per_node to match your available GPU count.

# Example: Running a 48-layer mHC config on 4 GPUs
torchrun --standalone --nproc_per_node=4 train.py config/train_fineweb10B_mhc_48l.py

To perform an ablation that keeps the residual routing close to the identity matrix, you can enable the identity-mix option in your config. This applies the formula H_res = (1-α) * I + α * S, where S is the projected matrix (Sinkhorn or orthostochastic) and α is a learned parameter.

Required config keys:

• mhc_residual_identity_mix = True
• mhc_residual_alpha = 0.01 (initial value)

mHC supports an alternative H_res projection using the Newton-Schulz method instead of the default Sinkhorn-Knopp.

To enable this, set mhc_h_res_proj = "orthostochastic" in your configuration file.

By default, the implementation uses fixed Newton-Schulz coefficients (ns_steps=5, ns_coeffs=(3.0, -3.2, 1.2)). For research purposes, you can provide a per-step schedule by passing a tuple of (a, b, c) triplets to ns_coeffs. In this case, set ns_steps to the length of the ns_coeffs tuple.