take the free lunch

Author: lucidrains

README.md

@@ -249,3 +249,12 @@ loss = model(seq, memory_replay_backprop = True) # memory efficient training fro
     url     = {https://api.semanticscholar.org/CorpusID:272987528}
 }
 ```
+
+```bibtex
+@inproceedings{Zhou2024ValueRL,
+    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},
+    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:273532030}
+}
+```

recurrent_memory_transformer_pytorch/recurrent_memory_transformer.py

@@ -11,6 +11,7 @@
 from torch import nn, einsum, Tensor
 
 from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
 
 from recurrent_memory_transformer_pytorch.attend import Attend
 
@@ -120,6 +121,7 @@ def __init__(
         dim_head = 64,
         heads = 8,
         dropout = 0.,
+        accept_value_residual = False,
         use_flash_attn = False,
         use_custom_causal_attn_mask = False
     ):
@@ -141,22 +143,45 @@ def __init__(
         self.to_kv = Linear(dim, dim_inner * 2)
         self.to_out = Linear(dim_inner, dim)
 
+        # learned value residual mixing
+
+        self.learned_value_residual_mix = None
+
+        if accept_value_residual:
+            self.learned_value_residual_mix = nn.Sequential(
+                Linear(dim, heads),
+                Rearrange('b n h -> b h n 1'),
+                nn.Sigmoid()
+            )
+
     def forward(
         self,
         x,
         rotary_emb: tuple[Tensor, Tensor] | None = None,
         mask = None,
-        xl_memories = None
+        xl_memories = None,
+        value_residual = None
     ):
+        assert not (exists(value_residual) ^ exists(self.learned_value_residual_mix))
+
         h = self.heads
         x = self.norm(x)
 
-
         q = self.to_q(x)
         k, v = self.to_kv(x).chunk(2, dim = -1)
 
+        # split heads
+
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
 
+        # handle value residual
+
+        orig_v = v
+
+        if exists(self.learned_value_residual_mix):
+            mix = self.learned_value_residual_mix(x)
+            v = v.lerp(value_residual, mix)
+
         # add a null key / value
         # to protect against an entirely masked out sequence
         # as well as giving attention ability to attend to nothing
@@ -191,7 +216,7 @@ def forward(
 
         out = rearrange(out, 'b h n d -> b n (h d)')
 
-        return self.to_out(out), next_xl_memories
+        return self.to_out(out), next_xl_memories, orig_v
 
 # transformer
 
@@ -269,14 +294,17 @@ def __init__(
 
         self.layers = ModuleList([])
 
-        for _ in range(depth):
+        for layer_index in range(depth):
+            is_first = layer_index == 0
+
             self.layers.append(ModuleList([
                 init_hyper_conn(dim = dim, branch = Attention(
                     dim = dim,
                     dim_head = dim_head,
                     causal = causal,
                     heads = heads,
                     use_flash_attn = use_flash_attn,
+                    accept_value_residual = not is_first,
                     use_custom_causal_attn_mask = memory_not_causal,
                     dropout = attn_dropout
                 )),
@@ -435,15 +463,20 @@ def forward(
         if has_xl_memories and self.enhanced_xl_recurrence and len(xl_memories) > 1:  # simply shift all the xl memories down by one, so lower layer gets access to representations from layer above
             xl_memories = [*xl_memories[1:], xl_memories[0]]
 
+        # value residual
+
+        value_residual = None
+
         # expand streams for hyper connections
 
         x = self.expand_streams(x)
 
         # attention and feedforward
 
         for attn, ff in self.layers:
-            x, xl_memories = attn(x, mask = mask, xl_memories = next(xl_memories_iter, None), rotary_emb = rotary_emb)
+            x, xl_memories, attn_values = attn(x, mask = mask, xl_memories = next(xl_memories_iter, None), rotary_emb = rotary_emb, value_residual = value_residual)
 
+            value_residual = default(value_residual, attn_values)
             new_xl_memories.append(xl_memories)
 
             x = ff(x)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'recurrent-memory-transformer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.6.0',
+  version = '0.7.0',
   license='MIT',
   description = 'Recurrent Memory Transformer - Pytorch',
   author = 'Phil Wang',