hyper-connections 0.3.15__py3-none-any.whl → 0.4.0__py3-none-any.whl

hyper_connections/manifold_constrained_hyper_connections.py

@@ -21,6 +21,7 @@ s - residual streams
 t - residual streams + num branch inputs
 f - number of fractions (division of feature dimension space)
 v - number of views for branch input
+p - proposals
 """
 
 # helper functions
@@ -67,7 +68,7 @@ def log_domain_sinkhorn_knopps(log_alpha, iters = 20):
         log_alpha = log_alpha - log_alpha.logsumexp(dim = -2, keepdim = True)
         log_alpha = log_alpha - log_alpha.logsumexp(dim = -1, keepdim = True)
 
-    return log_alpha.exp()
+    return log_alpha.exp().to(dtype)
 
 # main functions
 
@@ -213,6 +214,7 @@ class ManifoldConstrainedHyperConnections(Module):
         sinkhorn_iters = 20,
         log_domain_sinkhorn = False,
         forward_method_names: tuple[str, ...] = (),
+        num_dynamic_alpha_proposals = 1
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -256,6 +258,11 @@ class ManifoldConstrainedHyperConnections(Module):
         assert num_input_views >= 1
         self.num_input_views = num_input_views
 
+        # number of dynamic alpha proposals, for averaging Hres across proposals
+
+        self.has_dynamic_alpha_proposals = num_dynamic_alpha_proposals > 1
+        self.num_dynamic_alpha_proposals = num_dynamic_alpha_proposals
+
         # width connection
 
         init_alpha0 = torch.zeros((num_residual_streams_fracs, num_input_views_fracs))
@@ -263,7 +270,7 @@ class ManifoldConstrainedHyperConnections(Module):
 
         self.static_alpha = nn.Parameter(cat((init_alpha0, torch.eye(num_residual_streams_fracs)), dim = 1))
 
-        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams_fracs + num_input_views_fracs))
+        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(num_dynamic_alpha_proposals, dim, num_residual_streams_fracs + num_input_views_fracs))
 
         self.pre_branch_scale = nn.Parameter(torch.ones(1) * 1e-2)
         self.residual_scale = nn.Parameter(torch.ones(1) * 1e-2)
@@ -346,7 +353,7 @@ class ManifoldConstrainedHyperConnections(Module):
 
         normed = normed.float()
 
-        wc_weight = normed @ self.dynamic_alpha_fn.float()
+        wc_weight = einsum(normed, self.dynamic_alpha_fn.float(), '... d, p d e -> p ... e')
 
         pre_branch_scale = repeat(self.pre_branch_scale.float(), '1 -> s', s = self.num_fracs)
         residual_scale = repeat(self.residual_scale.float(), '1 -> s', s = self.num_fracs * streams)
@@ -372,6 +379,11 @@ class ManifoldConstrainedHyperConnections(Module):
 
         alpha = cat((alpha_pre, alpha_residual), dim = -1)
 
+        if self.has_dynamic_alpha_proposals:
+            alpha = reduce(alpha, 'p ... -> ...', 'mean')
+        else:
+            alpha = rearrange(alpha, '1 ... -> ...')
+
         alpha = self.split_fracs(alpha) # (batch, seq, fracs1, streams, fracs2, input + residual streams)
 
         # beta for weights from branch output back to residual streams

hyper_connections/vit.py

@@ -66,12 +66,12 @@ class Attention(Module):
         return self.to_out(out)
 
 class Transformer(Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., num_residual_streams = 4):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., num_residual_streams = 4, num_dynamic_alpha_proposals = 1):
         super().__init__()
         self.norm = nn.LayerNorm(dim)
         self.layers = ModuleList([])
 
-        init_hyper_conn, self.expand_streams, self.reduce_streams = mHC.get_init_and_expand_reduce_stream_functions(num_residual_streams)
+        init_hyper_conn, self.expand_streams, self.reduce_streams = mHC.get_init_and_expand_reduce_stream_functions(num_residual_streams, num_dynamic_alpha_proposals = num_dynamic_alpha_proposals)
 
         for _ in range(depth):
             self.layers.append(ModuleList([
@@ -92,7 +92,7 @@ class Transformer(Module):
         return self.norm(x)
 
 class ViT(Module):
-    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., num_residual_streams = 4):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., num_residual_streams = 4, num_dynamic_alpha_proposals = 1):
         super().__init__()
         image_height, image_width = pair(image_size)
         patch_height, patch_width = pair(patch_size)
@@ -117,7 +117,7 @@ class ViT(Module):
 
         self.dropout = nn.Dropout(emb_dropout)
 
-        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, num_residual_streams, num_dynamic_alpha_proposals)
 
         self.pool = pool
         self.to_latent = nn.Identity()

{hyper_connections-0.3.15.dist-info → hyper_connections-0.4.0.dist-info}/METADATA

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

{hyper_connections-0.3.15.dist-info → hyper_connections-0.4.0.dist-info}/RECORD

@@ -3,10 +3,10 @@ hyper_connections/hyper_connections.py,sha256=rqFJj3U0LF3uDKNKNPBpRrmf0oa2BGWVbD
 hyper_connections/hyper_connections_channel_first.py,sha256=Mh_hzhTi96ZoOPmhSKwUaF4TbHpNqhs83wNe5hNuL7o,6532
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=6BXKdSwyx6wdQVseebKG2EQkhVaVLrrepOlL8lLnex4,7855
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=ueT3CJPHrt5hRU7q1bFF0rANWJh_pXqclt6HiUu1gBY,11331
-hyper_connections/manifold_constrained_hyper_connections.py,sha256=iS0I6Ha5iAkWdrH-dEYpyLkHRRF8yarFqHJwlTV2qLI,17689
+hyper_connections/manifold_constrained_hyper_connections.py,sha256=-gWJbFR7PT-ji1pfPIMUGWdsfXlcUWVkyJvDmfM7BqM,18216
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
-hyper_connections/vit.py,sha256=fTC8hAYkD4qm-KURAj8SJ66C6ZWtsBdHf_kS-4rJZGQ,5049
-hyper_connections-0.3.15.dist-info/METADATA,sha256=PF7nUEiHAtWHgnRQfni_eNwywX0z7_5vKyj_bNm9oi0,6705
-hyper_connections-0.3.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.3.15.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.3.15.dist-info/RECORD,,
+hyper_connections/vit.py,sha256=BOWVfCAIzDQdnTq8OBzNUyiKGGILYZkIQ6mr1GKJVB0,5225
+hyper_connections-0.4.0.dist-info/METADATA,sha256=jEzPsFZ-71ZJ-WbkuhzcogDbvhWk4rKI3oINGQ2nMe0,6704
+hyper_connections-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.4.0.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.4.0.dist-info/RECORD,,