hyper-connections 0.1.8__tar.gz → 0.1.10__tar.gz

{hyper_connections-0.1.8 → hyper_connections-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.1.8
+Version: 0.1.10
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

{hyper_connections-0.1.8 → hyper_connections-0.1.10}/hyper_connections/hyper_connections.py

@@ -34,22 +34,31 @@ def identity(t):
 
 # main functions
 
-def get_expand_reduce_stream_functions(num_streams, disable = False):
+def get_expand_reduce_stream_functions(num_streams, add_stream_embed = False, dim = None, disable = False):
 
     if num_streams == 1 or disable:
         return (nn.Identity(), nn.Identity())
 
-    expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+    if add_stream_embed:
+        assert exists(dim), '`dim` must be passed into get_init_and_expand_reduce_stream_functions for returning an expansion function with stream embeddings added'
+
+        expand_fn = StreamEmbed(num_streams, dim, expand_to_streams = True)
+    else:
+        expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+
     reduce_fn = Reduce(pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
     return expand_fn, reduce_fn
 
-def get_init_and_expand_reduce_stream_functions(num_streams, disable = False):
+def get_init_and_expand_reduce_stream_functions(num_streams, dim = None, add_stream_embed = False, disable = False):
 
     hyper_conn_klass = HyperConnections if not disable else Residual
 
     init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, disable = disable)
+    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, add_stream_embed = add_stream_embed, dim = dim, disable = disable)
+
+    if exists(dim):
+        init_hyper_conn_fn = partial(init_hyper_conn_fn, dim = dim)
 
     return (init_hyper_conn_fn, *expand_reduce_fns)
 
@@ -132,6 +141,7 @@ class HyperConnections(Module):
         channel_first = False,
         dropout = 0.,
         residual_transform: Module | None = None, # to support resnet blocks where dimension in not equal to dimension out - usually a residual conv
+        add_branch_out_to_residual = True # will disable depth connections (weighted residual sum with beta) if set False
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -151,7 +161,7 @@ class HyperConnections(Module):
         self.num_residual_streams = num_residual_streams
         init_residual_index = default(layer_index, randrange(num_residual_streams)) % num_residual_streams # just choose one random residual stream if layer index not given
 
-        self.static_beta = nn.Parameter(torch.ones(num_residual_streams))
+        # width connection
 
         init_alpha0 = torch.zeros((num_residual_streams, 1))
         init_alpha0[init_residual_index, 0] = 1.
@@ -160,8 +170,15 @@ class HyperConnections(Module):
 
         self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + 1))
         self.dynamic_alpha_scale = nn.Parameter(torch.ones(()) * 1e-2)
-        self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim))
-        self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
+
+        # depth connection related (beta)
+
+        self.add_branch_out_to_residual = add_branch_out_to_residual
+
+        if add_branch_out_to_residual:
+            self.static_beta = nn.Parameter(torch.ones(num_residual_streams))
+            self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim))
+            self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
         # dropouts
 
@@ -196,9 +213,12 @@ class HyperConnections(Module):
 
         # beta for weights from branch output back to residual streams
 
-        dc_weight = self.act(normed @ self.dynamic_beta_fn)
-        dynamic_beta = dc_weight * self.dynamic_beta_scale
-        beta = dynamic_beta + self.static_beta
+        beta = None
+
+        if self.add_branch_out_to_residual:
+            dc_weight = self.act(normed @ self.dynamic_beta_fn)
+            dynamic_beta = dc_weight * self.dynamic_beta_scale
+            beta = dynamic_beta + self.static_beta
 
         mix_h = einsum(alpha, residuals, '... s t, ... s d -> ... t d')
 
@@ -210,6 +230,8 @@ class HyperConnections(Module):
         return branch_input, maybe_transformed_residuals, dict(beta = beta)
 
     def depth_connection(self, branch_output, residuals, *, beta):
+        assert self.add_branch_out_to_residual
+
         # 'depth' connection
 
         if self.channel_first:
@@ -244,6 +266,10 @@ class HyperConnections(Module):
         branch_input, residuals, residual_kwargs = self.width_connection(residuals)
 
         def add_residual_fn(branch_out):
+
+            if not self.add_branch_out_to_residual:
+                return branch_out
+
             (branch_out, *rest), tree_spec = tree_flatten(branch_out)
 
             branch_out = self.depth_connection(branch_out, residuals, **residual_kwargs)

{hyper_connections-0.1.8 → hyper_connections-0.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.1.8"
+version = "0.1.10"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }