x-transformers 1.44.5__py3-none-any.whl → 1.44.7__py3-none-any.whl

x_transformers/x_transformers.py

@@ -882,6 +882,7 @@ class HyperConnection(Module):
         *,
         layer_index,
         num_residual_streams,
+        num_input_views = 1,
         tanh = True,
         **kwargs
     ):
@@ -900,13 +901,16 @@ class HyperConnection(Module):
 
         self.static_beta = nn.Parameter(torch.ones(num_residual_streams))
 
-        init_alpha0 = torch.zeros((num_residual_streams, 1))
-        init_alpha0[layer_index % num_residual_streams, 0] = 1.
+        init_alpha0 = torch.zeros((num_residual_streams, num_input_views))
+        init_alpha0[layer_index % num_residual_streams, :] = 1.
 
         self.static_alpha = nn.Parameter(torch.cat([init_alpha0, torch.eye(num_residual_streams)], dim = 1))
 
-        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + 1))
+        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + num_input_views))
         self.dynamic_alpha_scale = nn.Parameter(torch.ones(()) * 1e-2)
+
+        self.num_input_views = num_input_views
+
         self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim))
         self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
@@ -928,7 +932,13 @@ class HyperConnection(Module):
 
         mix_h = einsum('... s t, ... s d -> ... t d', alpha, residuals)
 
-        branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
+        views = self.num_input_views
+
+        if views == 1:
+            branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
+        else:
+            branch_input, residuals = mix_h[..., :views, :], mix_h[..., views:, :]
+            branch_input = rearrange(branch_input, '... v d -> v ... d')
 
         return branch_input, residuals, dict(beta = beta)
 
@@ -1200,6 +1210,7 @@ class Attention(Module):
         learned_value_residual_mix = False,
         laser = False, # https://arxiv.org/abs/2411.03493v1
         laser_softclamp_value = 15.,
+        qkv_receive_diff_residuals = False,
         onnxable = False,
         attend_sdp_kwargs: dict = dict(
             enable_flash = True,
@@ -1239,6 +1250,10 @@ class Attention(Module):
         assert not (shared_kv and value_dim_head != dim_head), 'key and value head dimensions must be equal for shared key / values'
         self.to_v = LinearNoBias(dim_kv, v_dim) if not shared_kv else None
 
+        # whether qkv receives different residual stream combinations from hyper connections
+
+        self.qkv_receive_diff_residuals = qkv_receive_diff_residuals
+
         # enhancing gradients to attention through exponentiated values
 
         self.laser = laser
@@ -1423,14 +1438,21 @@ class Attention(Module):
         cache: Intermediates | None = None,
         value_residual = None
     ):
-        b, n, h, kv_h, head_scale, num_mem_kv, device, has_context = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, self.num_mem_kv, x.device, exists(context)
+        b, n, h, kv_h, head_scale, num_mem_kv, device, has_context, qkv_receive_diff_residuals = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, self.num_mem_kv, x.device, exists(context), self.qkv_receive_diff_residuals
 
-        kv_input = default(context, x)
+        assert not (qkv_receive_diff_residuals and has_context), 'qkv receiving different sequences can only be used for self attention'
 
-        q_input = x
-        k_input = kv_input
-        v_input = kv_input
-        r_input = x
+        if qkv_receive_diff_residuals:
+            assert not exists(self.to_r)
+
+            q_input, k_input, v_input = x
+        else:
+            kv_input = default(context, x)
+
+            q_input = x
+            k_input = kv_input
+            v_input = kv_input
+            r_input = x
 
         if exists(mem):
             k_input, mem_packed_shape = pack([mem, k_input], 'b * d')
@@ -1735,7 +1757,9 @@ class AttentionLayers(Module):
         layerscale_init_value = 0.,
         unet_skips = False,
         num_residual_streams = 1,
+        qkv_receive_diff_residuals = False,
         reinject_input = False,              # seen first in DEQ paper https://arxiv.org/abs/1909.01377, but later used in a number of papers trying to achieve depthwise generalization https://arxiv.org/abs/2410.03020v1
+        learned_reinject_input_gate = False,
         add_value_residual = False,          # resformer from Zhou et al - https://arxiv.org/abs/2410.17897v1 - further corroboration by https://arxiv.org/abs/2412.15113 (faster emergence of ICL) - looks like this setting may becoming a necessity for every transformer soon
         learned_value_residual_mix = True,   # seeing big improvements when the value residual mix value is learned per token - credit goes to @faresobeid for taking the first step with learned scalar mix, then @Blinkdl for taking it a step further with data dependent. here we will use per token learned
         rel_pos_kwargs: dict = dict(),
@@ -1770,6 +1794,8 @@ class AttentionLayers(Module):
 
         assert not (num_residual_streams > 1 and gate_residual)
 
+        assert not (num_residual_streams == 1 and qkv_receive_diff_residuals)
+
         # positions related
 
         self.disable_abs_pos_emb = default(disable_abs_pos_emb, (rel_pos_bias or rotary_pos_emb))
@@ -1993,6 +2019,7 @@ class AttentionLayers(Module):
 
         self.reinject_input = reinject_input
         self.reinject_input_proj = nn.Linear(dim, dim, bias = False) if reinject_input else None
+        self.learned_reinject_input_gate = nn.Linear(dim, 1, bias = False) if learned_reinject_input_gate else None
 
         # add the value from the first self attention block to all latter projected self attention values as a residual
 
@@ -2018,7 +2045,7 @@ class AttentionLayers(Module):
 
             if layer_type == 'a':
                 self_attn_learned_value_residual = learned_value_residual_mix and not is_first_self_attn
-                layer = Attention(dim, heads = heads, causal = causal, learned_value_residual_mix = self_attn_learned_value_residual, **attn_kwargs)
+                layer = Attention(dim, heads = heads, causal = causal, qkv_receive_diff_residuals = qkv_receive_diff_residuals, learned_value_residual_mix = self_attn_learned_value_residual, **attn_kwargs)
                 is_first_self_attn = False
             elif layer_type == 'c':
                 layer = Attention(dim, heads = heads, **{**attn_kwargs, **cross_attn_kwargs})
@@ -2039,6 +2066,10 @@ class AttentionLayers(Module):
 
             if num_residual_streams > 1:
                 residual_fn = partial(HyperConnection, num_residual_streams = num_residual_streams)
+
+                if layer_type == 'a':
+                    residual_fn = partial(residual_fn, num_input_views = 3)
+
             elif gate_residual:
                 residual_fn = GRUGating
             else:
@@ -2224,7 +2255,9 @@ class AttentionLayers(Module):
         layer_variables = tuple(tuple(layer_variable[i] for i in layers_execute_order) for layer_variable in layer_variables)
 
         # derived input for reinjection if needed
+
         inp_inject = None
+
         if self.reinject_input:
             assert not exists(in_attn_cond)
             inp_inject = self.reinject_input_proj(x)
@@ -2233,6 +2266,10 @@ class AttentionLayers(Module):
             # handle in-attention conditioning, which serves the same purpose of having the network learn the residual
             inp_inject = in_attn_cond if in_attn_cond.ndim == 3 else rearrange(in_attn_cond, 'b d -> b 1 d')
 
+        if exists(inp_inject) and exists(self.learned_reinject_input_gate):
+            inp_inject_gate = self.learned_reinject_input_gate(x).sigmoid()
+            inp_inject = inp_inject * inp_inject_gate
+
         # store all hiddens for skips
 
         skip_hiddens = []

{x_transformers-1.44.5.dist-info → x_transformers-1.44.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: x-transformers
-Version: 1.44.5
+Version: 1.44.7
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.44.5.dist-info → x_transformers-1.44.7.dist-info}/RECORD

@@ -6,11 +6,11 @@ x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=-B4xutlod09EwaA1tE24GgD-5ioi4HAXIwmI5N4MYio,103124
+x_transformers/x_transformers.py,sha256=4jgGkk-OkPdZYPouPs768KHw5eFmskC3m61h6LPfkKY,104695
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.44.5.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.44.5.dist-info/METADATA,sha256=j_8nbdft6Qx1NtIjSd1NyERsrpu-W9L3-CuR0hug_S0,924
-x_transformers-1.44.5.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-x_transformers-1.44.5.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.44.5.dist-info/RECORD,,
+x_transformers-1.44.7.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.44.7.dist-info/METADATA,sha256=D9xg_-8z4eeOXCwP1w86SFthU9RanucoxlpBLkDCUJ0,924
+x_transformers-1.44.7.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+x_transformers-1.44.7.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.44.7.dist-info/RECORD,,