Mhc Algorithm

mHC: Manifold-Constrained Hyper-Connections

Overview

mHC (Manifold-Constrained Hyper-Connections) stabilizes deep network training by constraining residual mixing matrices to be doubly stochastic. It provides:

• Stable Training: Lower gradient norm variance via doubly stochastic constraints
• Multiple Streams: Hyper-Connections with learnable mixing across residual streams
• Sinkhorn Projection: Log-space Sinkhorn-Knopp algorithm for doubly stochastic projection
• GPT Integration: Pattern for wrapping attention and MLP layers

Two components:

• HyperConnections Module: Core PyTorch module with Hres, Hpre, H_post matrices
• Sinkhorn-Knopp: Log-space projection to doubly stochastic manifold

Quick Reference

| Topic | Reference | |-------|-----------| | Core Concepts & Math | [Core Concepts](references/core-concepts.md) | | Sinkhorn Algorithm | [Sinkhorn-Knopp](references/sinkhorn-knopp.md) | | HyperConnections Module | [Module Implementation](references/module-implementation.md) | | GPT Integration | [GPT Integration](references/gpt-integration.md) | | Common Pitfalls | [Pitfalls](references/pitfalls.md) |

Installation

# Required packages
pip install torch einops numpy

Minimal Example

import torch
import torch.nn as nn
from einops import rearrange, einsum

def sinkhorn_knopp(logits, num_iters=20, tau=0.05):
    log_alpha = logits / tau
    for _ in range(num_iters):
        log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-1, keepdim=True)
        log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-2, keepdim=True)
    return torch.exp(log_alpha)

class HyperConnections(nn.Module):
    def __init__(self, num_streams, dim, branch=None, layer_idx=0):
        super().__init__()
        self.num_streams = num_streams
        self.branch = branch

        # Initialize H_res near identity (use small negative for gradient flow)
        init_h_res = torch.full((num_streams, num_streams), -0.1)
        init_h_res.fill_diagonal_(0.0)
        self.H_res_logits = nn.Parameter(init_h_res)

        # H_pre/H_post for depth connections
        init_h_pre = torch.full((1, num_streams), -0.1)
        init_h_pre[0, layer_idx % num_streams] = 0.0
        self.H_pre_logits = nn.Parameter(init_h_pre)
        self.H_post_logits = nn.Parameter(torch.zeros(1, num_streams))

    def forward(self, x):
        s = self.num_streams
        x = rearrange(x, "(b s) t d -> b t s d", s=s)

        h_res = sinkhorn_knopp(self.H_res_logits)
        x_mixed = einsum(h_res, x, "s t, b n s d -> b n t d")

        h_pre = self.H_pre_logits.softmax(dim=-1)
        branch_in = einsum(h_pre, x, "v s, b n s d -> b n v d").squeeze(-2)

        branch_out = self.branch(branch_in) if self.branch else branch_in

        h_post = self.H_post_logits.softmax(dim=-1)
        depth_out = einsum(branch_out, h_post, "b t d, v s -> b t s d")

        output = x_mixed + depth_out
        return rearrange(output, "b t s d -> (b s) t d")

Common Imports

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, einsum, repeat, reduce

When to Use What

| Scenario | Approach | |----------|----------| | Standard residual connection | No mHC needed | | Deep networks (>12 layers) with stability issues | Use mHC with numstreams=4 | | GPT/Transformer training | Wrap both attention and MLP with HyperConnections | | Custom Sinkhorn iterations | Adjust numiters (20 default) and tau (0.05 default) | | Memory-constrained training | Reduce num_streams or batch size |

External Resources

• mHC Paper: https://arxiv.org/abs/2512.24880
• Hyper-Connections: https://arxiv.org/abs/2409.19606
• Sinkhorn's Theorem: https://en.wikipedia.org/wiki/Sinkhorn%27s_theorem

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: benchflow-ai
• Source: benchflow-ai/skillsbench
• License: Apache-2.0
• Homepage: https://www.skillsbench.ai

Install and usage instructions live in the source repository linked above.