x-transformers 1.44.8__py3-none-any.whl → 2.0.1__py3-none-any.whl

x_transformers/x_transformers.py

@@ -9,7 +9,7 @@ from packaging import version
 import torch
 from torch.amp import autocast
 import torch.nn.functional as F
-from torch import nn, einsum, Tensor
+from torch import nn, einsum, Tensor, cat, stack, arange
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torch.nn import Module, ModuleList, ModuleDict
 
@@ -18,14 +18,22 @@ from collections import namedtuple
 from contextlib import nullcontext
 from dataclasses import dataclass
 
+from loguru import logger
+
+from x_transformers.attend import Attend, Intermediates
+from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
+
 import einx
 from einops.layers.torch import Rearrange
 from einops import rearrange, repeat, reduce, pack, unpack
 
-from loguru import logger
+# einstein notation
 
-from x_transformers.attend import Attend, Intermediates
-from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
+# b - batch
+# n - sequence
+# d - feature dimension
+# h - attention heads
+# i, j - sequence (source, target)
 
 # constants
 
@@ -220,7 +228,7 @@ def dropout_seq(seq, mask, dropout):
     num_keep = max(1,  int(keep_prob * n))
     keep_indices = logits.topk(num_keep, dim = 1).indices
 
-    batch_indices = torch.arange(b, device = device)
+    batch_indices = arange(b, device = device)
     batch_indices = rearrange(batch_indices, 'b -> b 1')
 
     seq = seq[batch_indices, keep_indices]
@@ -228,7 +236,7 @@ def dropout_seq(seq, mask, dropout):
     if exists(mask):
         seq_counts = mask.sum(dim = -1)
         seq_keep_counts = torch.ceil(seq_counts * keep_prob).int()
-        keep_mask = torch.arange(num_keep, device = device) < rearrange(seq_keep_counts, 'b -> b 1')
+        keep_mask = arange(num_keep, device = device) < rearrange(seq_keep_counts, 'b -> b 1')
 
         mask = mask[batch_indices, keep_indices] & keep_mask
 
@@ -274,7 +282,7 @@ class AbsolutePositionalEmbedding(Module):
         assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
 
         if not exists(pos):
-            pos = torch.arange(seq_len, device = device)
+            pos = arange(seq_len, device = device)
 
         if exists(seq_start_pos):
             pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
@@ -290,7 +298,7 @@ class ScaledSinusoidalEmbedding(Module):
         self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5)
 
         half_dim = dim // 2
-        freq_seq = torch.arange(half_dim).float() / half_dim
+        freq_seq = arange(half_dim).float() / half_dim
         inv_freq = theta ** -freq_seq
         self.register_buffer('inv_freq', inv_freq, persistent = False)
 
@@ -298,13 +306,13 @@ class ScaledSinusoidalEmbedding(Module):
         seq_len, device = x.shape[1], x.device
 
         if not exists(pos):
-            pos = torch.arange(seq_len, device = device)
+            pos = arange(seq_len, device = device)
 
         if exists(seq_start_pos):
             pos = pos - seq_start_pos[..., None]
 
         emb = einsum('i, j -> i j', pos, self.inv_freq)
-        emb = torch.cat((emb.sin(), emb.cos()), dim = -1)
+        emb = cat((emb.sin(), emb.cos()), dim = -1)
         return emb * self.scale
 
 class RelativePositionBias(Module):
@@ -344,8 +352,8 @@ class RelativePositionBias(Module):
 
     def forward(self, i, j):
         device = self.device
-        q_pos = torch.arange(j - i, j, dtype = torch.long, device = device)
-        k_pos = torch.arange(j, dtype = torch.long, device = device)
+        q_pos = arange(j - i, j, dtype = torch.long, device = device)
+        k_pos = arange(j, dtype = torch.long, device = device)
         rel_pos = einx.subtract('j, i -> i j', k_pos, q_pos)
         rp_bucket = self._relative_position_bucket(rel_pos, causal = self.causal, num_buckets = self.num_buckets, max_distance = self.max_distance)
         values = self.relative_attention_bias(rp_bucket)
@@ -376,7 +384,7 @@ class CoPE(Module):
         if not soft_onehot:
             return
 
-        self.register_buffer('positions', torch.arange(max_pos))
+        self.register_buffer('positions', arange(max_pos))
 
     def forward(self, query, attn_logits):
 
@@ -445,13 +453,13 @@ class DynamicPositionBias(Module):
         n, device = j, self.device
 
         # get the (n x n) matrix of distances
-        seq_arange = torch.arange(n, device = device)
-        context_arange = torch.arange(n, device = device)
+        seq_arange = arange(n, device = device)
+        context_arange = arange(n, device = device)
         indices = einx.subtract('i, j -> i j', seq_arange, context_arange)
         indices += (n - 1)
 
         # input to continuous positions MLP
-        pos = torch.arange(-n + 1, n, device = device).float()
+        pos = arange(-n + 1, n, device = device).float()
         pos = rearrange(pos, '... -> ... 1')
 
         if self.log_distance:
@@ -525,8 +533,8 @@ class AlibiPositionalBias(Module):
         if exists(self.bias) and self.bias.shape[-1] >= j and self.bias.shape[-2] >= i:
             return self.bias[..., -i:, -j:]
 
-        seq_arange = torch.arange(j - i, j, device = device)
-        context_arange = torch.arange(j, device = device)
+        seq_arange = arange(j - i, j, device = device)
+        context_arange = arange(j, device = device)
         bias = -einx.subtract('j, i -> 1 i j', context_arange, seq_arange).abs()
 
         bias = bias * self.slopes
@@ -642,7 +650,7 @@ class RotaryEmbedding(Module):
         # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
         base *= base_rescale_factor ** (dim / (dim - 2))
 
-        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
+        inv_freq = 1. / (base ** (arange(0, dim, 2).float() / dim))
         self.register_buffer('inv_freq', inv_freq)
 
         assert interpolation_factor >= 1.
@@ -652,7 +660,7 @@ class RotaryEmbedding(Module):
             self.register_buffer('scale', None)
             return
 
-        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+        scale = (arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
 
         self.scale_base = scale_base
         self.register_buffer('scale', scale)
@@ -660,7 +668,7 @@ class RotaryEmbedding(Module):
     def forward_from_seq_len(self, seq_len):
         device = self.inv_freq.device
 
-        t = torch.arange(seq_len, device = device)
+        t = arange(seq_len, device = device)
         return self.forward(t)
 
     @autocast('cuda', enabled = False)
@@ -671,7 +679,7 @@ class RotaryEmbedding(Module):
             t = rearrange(t, 'n -> 1 n')
 
         freqs = torch.einsum('b i , j -> b i j', t.type_as(self.inv_freq), self.inv_freq) / self.interpolation_factor
-        freqs = torch.stack((freqs, freqs), dim = -1)
+        freqs = stack((freqs, freqs), dim = -1)
         freqs = rearrange(freqs, '... d r -> ... (d r)')
 
         if not exists(self.scale):
@@ -679,7 +687,7 @@ class RotaryEmbedding(Module):
 
         power = (t - (max_pos // 2)) / self.scale_base
         scale = self.scale ** rearrange(power, '... n -> ... n 1')
-        scale = torch.stack((scale, scale), dim = -1)
+        scale = stack((scale, scale), dim = -1)
         scale = rearrange(scale, '... d r -> ... (d r)')
 
         return freqs, scale
@@ -687,7 +695,7 @@ class RotaryEmbedding(Module):
 def rotate_half(x):
     x = rearrange(x, '... (d r) -> ... d r', r = 2)
     x1, x2 = x.unbind(dim = -1)
-    x = torch.stack((-x2, x1), dim = -1)
+    x = stack((-x2, x1), dim = -1)
     return rearrange(x, '... d r -> ... (d r)')
 
 @autocast('cuda', enabled = False)
@@ -703,7 +711,7 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
     # partial rotary embeddings, Wang et al. GPT-J
     t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
     t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
-    out = torch.cat((t, t_unrotated), dim = -1)
+    out = cat((t, t_unrotated), dim = -1)
 
     return out.type(orig_dtype)
 
@@ -833,6 +841,15 @@ class SimpleRMSNorm(Module):
     def forward(self, x):
         return F.normalize(x, dim = -1) * self.scale
 
+class MultiheadRMSNorm(Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.rmsnorm = SimpleRMSNorm(dim)
+        self.gamma = nn.Parameter(torch.zeros(heads, 1, dim))
+
+    def forward(self, x):
+        return self.rmsnorm(x) * (self.gamma + 1.)
+
 # residual and residual gates
 
 class Residual(Module):
@@ -904,7 +921,7 @@ class HyperConnection(Module):
         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.static_alpha = nn.Parameter(cat([init_alpha0, torch.eye(num_residual_streams)], dim = 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)
@@ -973,7 +990,7 @@ class ShiftTokens(Module):
         splitted = x.split(feats_per_shift, dim = -1)
         segments_to_shift, rest = splitted[:segments], splitted[segments:]
         segments_to_shift = [shift(*args, mask = mask) for args in zip(segments_to_shift, shifts)]
-        x = torch.cat((*segments_to_shift, *rest), dim = -1)
+        x = cat((*segments_to_shift, *rest), dim = -1)
         return self.fn(x, **kwargs)
 
 class FoldAxially(Module):
@@ -1080,7 +1097,7 @@ class ConcatCombine(Module):
 
     def forward(self, x, prev_layers: list[Tensor]):
         skip = prev_layers[self.prev_layer_ind]
-        concatted_skip = torch.cat((skip, x), dim = -1)
+        concatted_skip = cat((skip, x), dim = -1)
         return self.combine(concatted_skip)
 
 # feedforward
@@ -1189,12 +1206,10 @@ class Attention(Module):
         hybrid_fold_axial_dim: int | None = None,
         one_kv_head = False,
         kv_heads = None,
-        shared_kv = False,
         value_dim_head = None,
         dim_out = None,
-        tensor_product = False,      # https://arxiv.org/abs/2208.06061
         add_zero_kv = False,         # same as add_zero_attn in pytorch
-        rotary_embed_values = False,
+        rotate_num_heads = None,
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
         data_dependent_alibi_per_row_dim_head = 8,
@@ -1205,12 +1220,15 @@ class Attention(Module):
         cope_talking_heads = False,
         softclamp_logits = False,
         logit_softclamp_value = 50.,
-        neutreno_value_residual = False, # Nguyen et al. https://arxiv.org/abs/2312.00751
-        neutreno_alpha = 0.4,
         learned_value_residual_mix = False,
-        laser = False, # https://arxiv.org/abs/2411.03493v1
+        laser = False,                # https://arxiv.org/abs/2411.03493v1
         laser_softclamp_value = 15.,
         qkv_receive_diff_residuals = False,
+        use_latent_q = False,
+        dim_latent_q = None,
+        use_latent_kv = False,
+        dim_latent_kv = None,
+        latent_rope_subheads = None,
         onnxable = False,
         attend_sdp_kwargs: dict = dict(
             enable_flash = True,
@@ -1242,13 +1260,51 @@ class Attention(Module):
         v_dim = value_dim_head * kv_heads
         out_dim = value_dim_head * heads
 
-        self.to_q = LinearNoBias(dim, q_dim)
-        self.to_k = LinearNoBias(dim_kv, k_dim)
+        # determine input dimensions to qkv based on whether intermediate latent q and kv are being used
+        # for eventually supporting multi-latent attention (MLA)
+
+        self.to_latent_q = None
+        self.to_latent_kv = None
+        self.to_rotateable_k = None # for their "decoupled rope", subheads of keys that comes directly from base sequence (does not go through latents)
+
+        dim_q_input = dim
+        dim_kv_input = dim_kv
+
+        if use_latent_q:
+            assert exists(dim_latent_q)
+            self.to_latent_q = LinearNoBias(dim, dim_latent_q)
+            dim_q_input = dim_latent_q
+
+        if use_latent_kv:
+            assert exists(dim_latent_kv)
+            self.to_latent_kv = LinearNoBias(dim, dim_latent_kv)
+            dim_kv_input = dim_latent_kv
+
+        if exists(latent_rope_subheads):
+            assert not exists(rotate_num_heads)
+            rotate_num_heads = latent_rope_subheads
+
+            k_dim = dim_head * (kv_heads - latent_rope_subheads)
 
-        # shared key / values, for further memory savings during inference
+            self.to_rotateable_k = LinearNoBias(dim, dim_head * latent_rope_subheads)
+            self.split_rotateable_k_heads = Rearrange('b n (h d) -> b h n d', h = latent_rope_subheads)
 
-        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
+        self.use_latent_q = use_latent_q
+        self.use_latent_kv = use_latent_kv
+
+        # query key projection
+
+        self.to_q = LinearNoBias(dim_q_input, q_dim)
+        self.to_k = LinearNoBias(dim_kv_input, k_dim)
+        self.to_v = LinearNoBias(dim_kv_input, v_dim)
+
+        # split and merge of attention heads
+
+        self.split_q_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.split_k_heads = Rearrange('b n (h d) -> b h n d', d = dim_head)
+        self.split_v_heads = Rearrange('b n (h d) -> b h n d', d = value_dim_head)
+
+        self.merge_heads = Rearrange('b h n d -> b n (h d)')
 
         # whether qkv receives different residual stream combinations from hyper connections
 
@@ -1259,15 +1315,6 @@ class Attention(Module):
         self.laser = laser
         self.laser_softclamp_value = laser_softclamp_value
 
-        # relations projection from tp-attention
-
-        self.to_r = LinearNoBias(dim, v_dim) if tensor_product else None
-
-        # the value residual used by Nguyen et al. in https://arxiv.org/abs/2312.00751 for countering oversmoothing
-
-        self.neutreno_value_residual = neutreno_value_residual
-        self.neutreno_alpha = neutreno_alpha
-
         # add GLU gating for aggregated values, from alphafold2
 
         self.to_v_gate = None
@@ -1393,12 +1440,22 @@ class Attention(Module):
 
         # hybrid module, in same vein as hymba https://www.arxiv.org/abs/2411.13676
 
+        hybrid_mix = None
+        hybrid_norms = None
         hybrid_module = maybe(deepcopy)(hybrid_module)
 
         if exists(hybrid_module) and exists(hybrid_fold_axial_dim):
             hybrid_module = FoldAxially(axial_dim = hybrid_fold_axial_dim, fn = hybrid_module)
+            hybrid_mix = LinearNoBias(dim, heads)
+
+            hybrid_norms = ModuleList([
+                MultiheadRMSNorm(dim_head, heads = heads),
+                MultiheadRMSNorm(dim_head, heads = heads)
+            ])
 
         self.hybrid_module = hybrid_module
+        self.hybrid_norms = hybrid_norms
+        self.hybrid_mix = hybrid_mix
         self.hybrid_mask_kwarg = hybrid_mask_kwarg # for bidirectional, can forward `mask` into the hybrid module and let it handle variable lengths
 
         # output dimension by default same as input, but can be overridden
@@ -1406,9 +1463,15 @@ class Attention(Module):
         dim_out = default(dim_out, dim)
         self.to_out = nn.Sequential(LinearNoBias(out_dim, dim_out * 2), nn.GLU()) if on_attn else LinearNoBias(out_dim, dim_out)
 
-        # whether to rotate positions into values, for absolute positions in addition to relative
+        # the number of attention heads to rotate, for decoupled rope in multi-latent attention
+
+        rotate_num_heads = default(rotate_num_heads, heads)
 
-        self.rotary_embed_values = rotary_embed_values
+        assert 0 < rotate_num_heads <= heads
+        is_partial_rotate_heads = rotate_num_heads < heads
+        assert not (is_partial_rotate_heads and kv_heads < heads), 'grouped query attention not compatible with partial rotate heads (decoupled rope for multi-latent attention), yet'
+
+        self.rotate_num_heads = rotate_num_heads
 
         # whether parent can kv cache
 
@@ -1438,47 +1501,79 @@ class Attention(Module):
         cache: Intermediates | None = None,
         value_residual = None
     ):
-        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
+        b, n, h, kv_h, head_scale, num_mem_kv, device, has_context, qkv_receive_diff_residuals, is_multi_latent_attn = 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, self.use_latent_kv
+
+        # an interesting possibility with hyper connections
+        # having queries, keys, values be routed from different layers
 
         assert not (qkv_receive_diff_residuals and has_context), 'qkv receiving different sequences can only be used for self attention'
 
         if qkv_receive_diff_residuals:
-            assert not exists(self.to_r)
+            assert x.ndim == 4 and x.shape[0] == 3
 
             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
+            q_input, k_input, v_input = x, kv_input, kv_input
 
         if exists(mem):
             k_input, mem_packed_shape = pack([mem, k_input], 'b * d')
             v_input, _ = pack([mem, v_input], 'b * d')
 
+        # multi-latent attention logic
+        # https://arxiv.org/abs/2405.04434 - Deepseek-AI team
+
+        k_sub_heads = None # the rotateable subheads of keys derived from base sequence
+
+        if self.use_latent_q:
+            q_input = self.to_latent_q(q_input)
+
+        if is_multi_latent_attn:
+            assert not qkv_receive_diff_residuals
+            needs_k_sub_heads = exists(self.to_rotateable_k)
+
+            latent_kv_input = self.to_latent_kv(k_input)
+
+            if needs_k_sub_heads:
+                rotateable_k = self.to_rotateable_k(k_input)
+                k_sub_heads = self.split_rotateable_k_heads(rotateable_k)
+
+            if exists(cache):
+                cached_latent_kv, maybe_cached_k_sub_heads = cache.cached_kv
+                latent_kv_input = cat((cached_latent_kv, latent_kv_input), dim = -2)
+
+                if exists(maybe_cached_k_sub_heads):
+                    k_sub_heads = cat((maybe_cached_k_sub_heads, k_sub_heads), dim = -2)
+
+            if return_intermediates:
+                cached_kv = (latent_kv_input, k_sub_heads)
+
+            k_input = v_input = latent_kv_input
+
+        # query, key, value projection
+
         q = self.to_q(q_input)
         k = self.to_k(k_input)
-        v = self.to_v(v_input) if exists(self.to_v) else k
-        r = self.to_r(r_input) if exists(self.to_r) else None
+        v = self.to_v(v_input)
+
+        q = self.split_q_heads(q)
+        k = self.split_k_heads(k)
+        v = self.split_v_heads(v)
 
-        q = rearrange(q, 'b n (h d) -> b h n d', h = h)
+        # take care of decoupled rope from multi-latent attention
 
-        k, v, r = tuple(maybe(rearrange)(t, 'b n (h d) -> b h n d', h = kv_h) for t in (k, v, r))
+        if exists(k_sub_heads):
+            k = cat((k, k_sub_heads), dim = 1)
 
-        # if previous values passed in for residual, either invoke resformer or neutreno
+        # if previous values passed in for residual, either invoke resformer
 
         orig_values = v
 
+        # https://arxiv.org/abs/2410.17897v1
+
         if exists(value_residual):
-            if self.neutreno_value_residual:
-                diff_values = (value_residual - v) * self.neutreno_alpha
-                diff_values = repeat(diff_values, 'b h n d -> b (r h) n d', r = h // kv_h)
-            else:
-                # https://arxiv.org/abs/2410.17897v1
-                value_residual_mix = self.to_value_residual_mix(q_input)
-                v = v * value_residual_mix + value_residual * (1. - value_residual_mix)
+            value_residual_mix = self.to_value_residual_mix(q_input)
+            v = value_residual.lerp(v, value_residual_mix)
 
         # qk normalization
 
@@ -1492,28 +1587,36 @@ class Attention(Module):
 
         # take care of caching
 
-        if exists(cache):
-            ck, cv = cache.cached_kv
+        if not is_multi_latent_attn:
+            if exists(cache):
+                ck, cv = cache.cached_kv
 
-            if exists(mem):
-                mk, k = unpack(k, mem_packed_shape, 'b h * d')
-                mv, v = unpack(v, mem_packed_shape, 'b h * d')
+                if exists(mem):
+                    mk, k = unpack(k, mem_packed_shape, 'b h * d')
+                    mv, v = unpack(v, mem_packed_shape, 'b h * d')
 
-            k = torch.cat((ck, k), dim = -2)
-            v = torch.cat((cv, v), dim = -2)
+                k = cat((ck, k), dim = -2)
+                v = cat((cv, v), dim = -2)
 
-            if exists(mem):
-                k = torch.cat((mk, k), dim = -2)
-                v = torch.cat((mv, v), dim = -2)
+                if exists(mem):
+                    k = cat((mk, k), dim = -2)
+                    v = cat((mv, v), dim = -2)
 
-        if return_intermediates:
-            mem_len = mem.shape[-2] if exists(mem) else 0
-            cached_kv = (k[..., mem_len:, :], v[..., mem_len:, :])
+            if return_intermediates:
+                mem_len = mem.shape[-2] if exists(mem) else 0
+                cached_kv = (k[..., mem_len:, :], v[..., mem_len:, :])
 
         if exists(rotary_pos_emb):
+            rotate_num_heads = self.rotate_num_heads
+            partial_rotate_heads = rotate_num_heads < h
+
             freqs, xpos_scale = rotary_pos_emb
             q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if exists(xpos_scale) else (1., 1.)
 
+            if partial_rotate_heads:
+                q_rest, q = q[:, :-rotate_num_heads], q[:, -rotate_num_heads:]
+                k_rest, k = k[:, :-rotate_num_heads], k[:, -rotate_num_heads:]
+
             q = apply_rotary_pos_emb(q, freqs, q_xpos_scale)
 
             if has_context:
@@ -1524,8 +1627,9 @@ class Attention(Module):
 
             k = apply_rotary_pos_emb(k, freqs, k_xpos_scale)
 
-            if self.rotary_embed_values:
-                v = apply_rotary_pos_emb(v, freqs, k_xpos_scale)
+            if partial_rotate_heads:
+                q = cat((q_rest, q), dim = 1)
+                k = cat((k_rest, k), dim = 1)
 
         input_mask = context_mask
 
@@ -1540,7 +1644,7 @@ class Attention(Module):
                 elif not exists(input_mask):
                     input_mask = pad_at_dim(mem_mask, (0, seq_len), dim = -1, value = True)
                 else:
-                    input_mask = torch.cat((mem_mask, input_mask), dim = -1)
+                    input_mask = cat((mem_mask, input_mask), dim = -1)
 
         # i, j determined for relative positional bias, excluding memory key / values
 
@@ -1555,8 +1659,8 @@ class Attention(Module):
                 mem_k = l2norm(mem_k)
                 mem_k = mem_k * self.qk_norm_k_scale
 
-            k = torch.cat((mem_k, k), dim = -2)
-            v = torch.cat((mem_v, v), dim = -2)
+            k = cat((mem_k, k), dim = -2)
+            v = cat((mem_v, v), dim = -2)
 
             if exists(input_mask):
                 input_mask = pad_at_dim(input_mask, (self.num_mem_kv, 0), dim = -1, value = True)
@@ -1580,8 +1684,8 @@ class Attention(Module):
             masks.append(~attn_mask)
 
         if exists(self.max_attend_past):
-            range_q = torch.arange(j - i, j, device = device)
-            range_k = torch.arange(j, device = device)
+            range_q = arange(j - i, j, device = device)
+            range_k = arange(j, device = device)
             dist = einx.subtract('i, j -> 1 1 i j', range_q, range_k)
             max_attend_past_mask = dist > self.max_attend_past
             max_attend_past_mask = pad_at_dim(max_attend_past_mask, (num_mem_kv, 0), value = False, dim = -1) # handle memory key / values
@@ -1629,18 +1733,10 @@ class Attention(Module):
         if self.laser:
             out = log(out)
 
-        # store the values for resformer or Neutreno
+        # store the values for resformer
 
         intermediates.values = orig_values
 
-        if exists(value_residual) and self.neutreno_value_residual:
-            out = out + diff_values
-
-        # https://arxiv.org/abs/2208.06061 proposes to add a residual for better gradients
-
-        if exists(r):
-            out = out * r + out
-
         # normformer scaling of heads
 
         if head_scale:
@@ -1652,11 +1748,9 @@ class Attention(Module):
             head_gate = self.to_v_head_gate(x)
             out = einx.multiply('b n h, b h n d ->b h n d', head_gate.sigmoid(), out)
 
-        # merge heads
-
-        out = rearrange(out, 'b h n d -> b n (h d)')
+        # if exists hybrid module, must do a normalization
 
-        # hybrid module
+         # hybrid module
 
         if exists(self.hybrid_module):
 
@@ -1674,8 +1768,23 @@ class Attention(Module):
             # handle hybrid out
 
             (hybrid_out, *rest_hybrid_outs), _ = tree_flatten(hybrid_outputs)
+
+            # handle variable hybrid output and multi rmsnorm before summing to main attention output (also normed)
+
+            if hybrid_out.ndim == 3:
+                hybrid_out = rearrange(hybrid_out, 'b n (h d) -> b h n d', h = h)
+
+            out_norm, hybrid_out_norm = self.hybrid_norms
+
+            out = out_norm(out)
+            hybrid_out = hybrid_out_norm(hybrid_out)
+
             out = 0.5 * (out + hybrid_out)
 
+        # merge heads
+
+        out = self.merge_heads(out)
+
         # alphafold2 styled gating of the values
 
         if exists(self.to_v_gate):
@@ -1747,8 +1856,6 @@ class AttentionLayers(Module):
         sandwich_norm = False,
         softclamp_output = False,
         softclamp_output_value = 30.,
-        resi_dual = False,
-        resi_dual_scale = 1.,
         zero_init_branch_output = False,
         layer_dropout = 0.,
         cross_attn_tokens_dropout = 0.,
@@ -1775,12 +1882,9 @@ class AttentionLayers(Module):
 
         dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
         data_dependent_alibi = attn_kwargs.get('data_dependent_alibi', False)
-        neutreno_value_residual = attn_kwargs.get('neutreno_value_residual', False)
 
         assert len(kwargs) == 0, f'unrecognized kwargs passed in {kwargs.keys()}'
 
-        add_value_residual |= neutreno_value_residual
-
         self.dim = dim
         self.causal = causal
         self.layers = ModuleList([])
@@ -1831,19 +1935,11 @@ class AttentionLayers(Module):
             assert alibi_num_heads <= heads, 'number of ALiBi heads must be less than the total number of heads'
             self.rel_pos = AlibiPositionalBias(heads = alibi_num_heads, total_heads = heads, **rel_pos_kwargs)
 
-        assert at_most_one_of(sandwich_norm, resi_dual), 'either sandwich norm or resiDual is selected, but not both'
         assert not (not pre_norm and sandwich_norm), 'sandwich norm cannot be used when not using prenorm'
 
-        if resi_dual:
-            pre_norm = False
-
         self.pre_norm = pre_norm
         self.sandwich_norm = sandwich_norm
 
-        self.resi_dual = resi_dual
-        assert 0 < resi_dual_scale <= 1., 'resiDual prenorm residual must be scaled by a factor greater than 0 and less than or equal to 1.'
-        self.resi_dual_scale = resi_dual_scale
-
         self.residual_attn = residual_attn
         self.cross_residual_attn = cross_residual_attn
         assert not (flash_attn and (residual_attn or cross_residual_attn)), 'flash attention is not compatible with residual attention'
@@ -2002,7 +2098,7 @@ class AttentionLayers(Module):
 
         # whether it has post norm
 
-        self.final_norm = norm_fn() if pre_norm or resi_dual else nn.Identity()
+        self.final_norm = norm_fn() if pre_norm else nn.Identity()
 
         # whether unet or not
 
@@ -2175,7 +2271,7 @@ class AttentionLayers(Module):
         # handle left padded sequences
 
         if exists(seq_start_pos):
-            seq_arange = torch.arange(x.shape[-2], device = x.device, dtype = torch.long)
+            seq_arange = arange(x.shape[-2], device = x.device, dtype = torch.long)
             left_pad_mask = seq_arange >= seq_start_pos[..., None]
 
             if exists(self_attn_kv_mask):
@@ -2193,7 +2289,7 @@ class AttentionLayers(Module):
                 mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
 
                 if not exists(pos):
-                    pos = torch.arange(x.shape[1] + mem_len, device = x.device) - mem_len
+                    pos = arange(x.shape[1] + mem_len, device = x.device) - mem_len
 
                 rotary_pos_emb = self.rotary_pos_emb(pos)
 
@@ -2213,7 +2309,7 @@ class AttentionLayers(Module):
         attn_cache = []
 
         if exists(cache):
-            assert not self.training and self.causal and not any([*map(exists, (mask, attn_mask))])
+            assert self.causal and not any([*map(exists, (mask, attn_mask))])
 
             if exists(context):
                 context = context[:, :0]
@@ -2231,13 +2327,7 @@ class AttentionLayers(Module):
         is_multistream = streams > 1
 
         if is_multistream:
-            x = repeat(x, 'b n d -> b n s d', s = streams)
-            x = x + self.stream_emb
-            x = rearrange(x, 'b n s d -> (b s) n d')
-
-        # outer residual - for resiDual paper
-
-        outer_residual = x * self.resi_dual_scale
+            x = einx.add('b n d, s d -> (b s) n d', x, self.stream_emb)
 
         # get layers to be executed
 
@@ -2359,9 +2449,6 @@ class AttentionLayers(Module):
             if not exists(first_cross_attn_inter) and layer_type == 'c':
                 first_cross_attn_inter = inter
 
-            if self.resi_dual:
-                outer_residual = outer_residual + out * self.resi_dual_scale
-
             if exists(post_branch_norm):
                 out = post_branch_norm(out)
 
@@ -2395,10 +2482,7 @@ class AttentionLayers(Module):
         if is_multistream:
             x = reduce(x, '(b s) n d -> b n d', 'sum', s = streams)
 
-        if self.resi_dual:
-            x = x + final_norm(outer_residual)
-        else:
-            x = final_norm(x)
+        x = final_norm(x)
 
         if not return_hiddens:
             return x
@@ -2444,7 +2528,7 @@ class PrefixDecoder(AttentionLayers):
             if isinstance(prefix_attn_len, int):
                 prefix_attn_len = torch.full((b,), prefix_attn_len, device = device)
 
-            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_attn_len, 'b -> b 1 1 1')
+            prefix_mask = arange(n, device = device) < rearrange(prefix_attn_len, 'b -> b 1 1 1')
             forwarded_mask = forwarded_mask | prefix_mask
 
         if exists(attn_mask):
@@ -2773,13 +2857,13 @@ class TransformerWrapper(Module):
             prepend_seq, prepend_dim = prepend_embeds.shape[1:]
             assert prepend_dim == x.shape[-1], 'prepended embeddings need to have same dimensions as text model dimensions'
 
-            x = torch.cat((prepend_embeds, x), dim = -2)
+            x = cat((prepend_embeds, x), dim = -2)
 
             if exists(prepend_mask) or exists(mask):
                 mask = default(mask, lambda: torch.ones((b, n), device = device, dtype = torch.bool))
                 prepend_mask = default(prepend_mask, lambda: torch.ones((b, prepend_seq), device = device, dtype = torch.bool))
 
-                mask = torch.cat((prepend_mask, mask), dim = -1)
+                mask = cat((prepend_mask, mask), dim = -1)
 
         # whether to reduce the gradient going to the embedding, from cogview paper, corroborated by GLM-130B model
 
@@ -2945,7 +3029,7 @@ class TransformerWrapper(Module):
 
         if return_mems:
             hiddens = intermediates.hiddens
-            new_mems = [torch.cat(pair, dim = -2) for pair in zip(mems, hiddens)] if exists(mems) else hiddens
+            new_mems = [cat(pair, dim = -2) for pair in zip(mems, hiddens)] if exists(mems) else hiddens
             new_mems = [t[..., -self.max_mem_len:, :].detach() for t in new_mems]
 
             if not return_intermediates: