Hyper-Connections

Hyper-Connections allow you to replace standard residual streams with multiple parallel streams as proposed in the Hyper-Connections paper.

To implement this, follow these four steps:

1. Initialize and Wrap: Use get_init_and_expand_reduce_stream_functions to get helper functions, then wrap your existing branch layer (e.g., nn.Linear) using the returned init_hyper_conn function.
2. Expand: Before your model's trunk (the loop containing your layers), call expand_stream on your residual tensor to expand it into multiple streams.
3. Forward: Pass the expanded residual through your wrapped branch function(s).
4. Reduce: After the trunk loop, call reduce_stream to sum the multiple streams back into a single residual tensor.

Note: For Transformers, the exact placement of the reduce_stream (before or after the final norm) is a research consideration.

import torch
from torch import nn
from hyper_connections import get_init_and_expand_reduce_stream_functions

# Setup
dim = 512
branch = nn.Linear(dim, dim)
residual = torch.randn(2, 1024, dim)
num_streams = 4

# Get helper functions
init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(num_streams)

# 1. Wrap the branch
hyper_conn_branch = init_hyper_conn(dim = dim, branch = branch)

# 2. Expand to multiple streams
residual = expand_stream(residual)

# 3. Forward through wrapped branch
residual = hyper_conn_branch(residual)

# 4. Reduce streams with summation
residual = reduce_stream(residual)

The Residual class is a base implementation for simple residual connections. It provides a decorate_branch method that wraps a standard branch function into a forward_and_add_residual function. This ensures that the output of the branch is automatically added back to the residual stream.

It supports:

• branch: The module to apply.
• residual_transform: A module (like a projection) applied to the residual before the connection.
• width_connection: Defines how residuals are prepared for the branch.
• depth_connection: Defines how the branch output is merged back into the residuals.

from hyper_connections.manifold_constrained_hyper_connections import Residual
import torch.nn as nn

class MyResidualBlock(Residual):
    def __init__(self, dim, **kwargs):
        super().__init__(dim = dim, **kwargs)

    def forward(self, residual, *args, **kwargs):
        # This uses the decorated branch logic
        return self.forward_and_add_residual(residual, *args, **kwargs)

# Usage
block = MyResidualBlock(dim=512, branch=nn.Linear(512, 512))
res = torch.randn(8, 512)
new_res = block(res)

Hyper-connections operate through two primary phases: Width Connection and Depth Connection.

1. Width Connection (Mixing Inputs)

Before the branch module is executed, the multiple residual streams are mixed into a single input stream for the branch.

• Dynamic Alpha: The streams are normalized and passed through a projection to calculate alpha weights. These weights determine how much each residual stream contributes to the branch_input.
• Additional Inputs: If additional_input_paths are provided, these extra streams are also projected and mixed into the process.
• Result: A single branch_input is produced, and the original residuals are transformed into a mixed state.

2. Depth Connection (Distributing Outputs)

After the branch module produces an output, it must be distributed back to the individual residual streams.

• Dynamic Beta: The branch output and the mixed residuals are used to calculate beta weights.
• Distribution: The branch output is scaled by beta and added back to each of the individual residual streams.
• Result: The output is a set of updated residual streams, typically expanded back to the (b*s, ..., d) format.

The hyper-connection mechanism follows these steps:

1. Width Connection: The input residual streams are mixed using dynamic and static weights ($\alpha$) to create num_branch_inputs. This creates a set of inputs that are then expanded to be processed by the branches.
2. Branch Processing: The branch modules (or functions) process these inputs. If multiple branch modules are provided, they process subsets of the branch inputs.
3. Depth Connection: The outputs from the branches are mixed using weights ($\beta$) and added back to the original residual streams. This ensures the connection is residual in nature.

Mental Model for Implementation:

• Before the trunk: Expand the single residual stream into $N$ streams (using expand_fn).
• The Trunk: Apply your branch functions to these $N$ streams.
• After the trunk: Reduce the $N$ streams back into a single residual stream (using reduce_fn).

When using the HyperConnections class directly, it encapsulates these steps via forward().

mHCv2 manages the flow of information between multiple residual streams and a central branch module using two types of connections:

1. Width Connection (Residuals $\to$ Branch):

   • It takes multiple residual streams and produces a single (or multiple views of) branch input.
   • It uses a dynamic weight $\alpha$ generated from the residual streams.
   • The $\alpha$ weights are constrained using the Sinkhorn-Knopp algorithm to ensure they lie on a specific manifold (doubly stochastic or similar), ensuring stable and controlled mixing.

2. Depth Connection (Branch $\to$ Residuals):

   • It takes the output of the branch module and distributes it back into the residual streams.
   • It uses a dynamic weight $\beta$ to control how much of the branch output is added to each stream.

This mechanism allows the model to dynamically decide which parts of the residual memory are most relevant for the current computation and how the computation should update that memory.

If you need more control than the wrapped branch approach, you can implement the connection manually. This allows you to access the learned betas used for the residual addition.

1. Instantiate the hyper connection with init_hyper_conn(dim=...).
2. Expand the residual using expand_stream.
3. Call the hyper connection object on the expanded residual. It returns two components: branch_input (the input to your branch layer) and add_residual (a function that performs the learned residual addition).
4. Pass branch_input through your branch layer, then pass the result to add_residual.
5. Reduce the streams using reduce_stream after your trunk loop.

Alternatively, you can use the shorthand: residual = hyper_conn.decorate_branch(branch)(residual).

from hyper_connections import get_init_and_expand_reduce_stream_functions

# ... setup dim, branch, residual, num_streams ...
init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4)

# 1. Instantiate
hyper_conn = init_hyper_conn(dim = 512)

# 2. Expand
residual = expand_stream(residual)

# 3. Manual forward with learned betas
branch_input, add_residual = hyper_conn(residual)
branch_output = branch(branch_input)
residual = add_residual(branch_output)

# 4. Reduce
residual = reduce_stream(residual)

The get_init_and_expand_reduce_stream_functions function is the primary entry point for configuring Hyper-Connections. It accepts the following arguments:

• num_streams (int): The number of parallel residual streams to create.
• disable (bool, default: False): If set to True, the returned functions will behave like standard plain residuals, allowing you to compare Hyper-Connections against a baseline without changing your model architecture code.
• num_fracs (int, default: 1): Used to implement Fractionated Feature Dimensions. Setting this to a value greater than 1 allows you to mix streams and fractions of the feature dimension, as proposed in the Frac-Connections paper.

# Disable hyper-connections to use plain residuals for comparison
init, expand, reduce = get_init_and_expand_reduce_stream_functions(4, disable = True)

# Use fractionated feature dimensions
init, expand, reduce = get_init_and_expand_reduce_stream_functions(1, num_fracs = 4)