stack-attention-pytorch 0.0.2__tar.gz → 0.0.4__tar.gz

{stack_attention_pytorch-0.0.2 → stack_attention_pytorch-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: stack-attention-pytorch
-Version: 0.0.2
+Version: 0.0.4
 Summary: Implementation of Differentiable Stacks for augmenting Transformers / Attention
 Project-URL: Homepage, https://pypi.org/project/stack-attention-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/stack-attention-pytorch
@@ -47,6 +47,32 @@ Description-Content-Type: text/markdown
 
 For following a line of research that augments attention with a differentiable stack, beginning with DuSell et al. at ETH Zurich
 
+## Install
+
+```bash
+$ pip install stack-attention
+```
+
+## Usage
+
+```python
+import torch
+from stack_attention.stack_trans_layer import StackTransLayer
+
+tokens = torch.randn(2, 512, 256)
+
+layer = StackTransLayer(256)
+
+out1, state = layer(tokens)
+
+out2, state = layer(
+    tokens,
+    stack_states = state
+)
+
+assert out1.shape == out2.shape == tokens.shape
+```
+
 ## Citations
 
 ```bibtex

{stack_attention_pytorch-0.0.2 → stack_attention_pytorch-0.0.4}/README.md

@@ -2,6 +2,32 @@
 
 For following a line of research that augments attention with a differentiable stack, beginning with DuSell et al. at ETH Zurich
 
+## Install
+
+```bash
+$ pip install stack-attention
+```
+
+## Usage
+
+```python
+import torch
+from stack_attention.stack_trans_layer import StackTransLayer
+
+tokens = torch.randn(2, 512, 256)
+
+layer = StackTransLayer(256)
+
+out1, state = layer(tokens)
+
+out2, state = layer(
+    tokens,
+    stack_states = state
+)
+
+assert out1.shape == out2.shape == tokens.shape
+```
+
 ## Citations
 
 ```bibtex

{stack_attention_pytorch-0.0.2 → stack_attention_pytorch-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stack-attention-pytorch"
-version = "0.0.2"
+version = "0.0.4"
 description = "Implementation of Differentiable Stacks for augmenting Transformers / Attention"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{stack_attention_pytorch-0.0.2 → stack_attention_pytorch-0.0.4}/stack_attention/stack_trans_layer.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from functools import partial
 
 import torch
-from torch import cat, nn, Tensor
+from torch import cat, nn, Tensor, tensor
 from torch.nn import Module, Linear, RMSNorm, Parameter
 import torch.nn.functional as F
 
@@ -58,7 +58,9 @@ class StackTransLayer(Module):
         num_stacks = 4,     # heads
         dim_stack = 16,     # dim head
         stack_size = 24,
-        prenorm = False
+        prenorm = False,
+        add_residual = True,
+        learned_residual_gate = True
     ):
         super().__init__()
 
@@ -87,10 +89,17 @@ class StackTransLayer(Module):
 
         self.null_stack = Parameter(torch.randn(dim_stack) * 1e-2)
 
-        # combine
+        # combining stack reads across number of stacks
 
         self.combine = LinearNoBias(dim_inner, dim)
 
+        # maybe combining with residual
+
+        self.add_residual = add_residual
+        self.learned_residual_gate = learned_residual_gate and add_residual
+
+        self.residual_scale = Parameter(tensor(0.)) if self.learned_residual_gate else None
+
     def forward(
         self,
         tokens,
@@ -101,7 +110,7 @@ class StackTransLayer(Module):
     ):
         assert action_temperature > 0.
 
-        orig, device = tokens, tokens.device
+        residual, device = tokens, tokens.device
 
         # maybe pre norm
 
@@ -182,18 +191,24 @@ class StackTransLayer(Module):
         read_stack_attn_logits = self.to_read_stack_attn(next_stack_with_null)
         read_stack_attn_logits = rearrange(read_stack_attn_logits, '... 1 -> ...')
 
-        read_stack_attn_logits = read_stack_attn_logits.masked_fill(next_mask_with_null, mask_value(next_stack_with_null))
+        read_stack_attn_logits = read_stack_attn_logits.masked_fill(~next_mask_with_null, mask_value(next_stack_with_null))
 
         read_stack_attn = read_stack_attn_logits.softmax(dim = -1)
 
         read_stack_out = einsum(next_stack_with_null, read_stack_attn, 'b h s d, b h s -> b h d')
 
-        # combine
+        # combine heads
 
-        out = rearrange(stack_inputs, 'b h d -> b (h d)')
+        out = rearrange(read_stack_out, 'b h d -> b (h d)')
 
         out = self.combine(out)
 
         out = inverse_pack(out)
 
+        # maybe add residual
+
+        if self.add_residual:
+            residual_scale = self.residual_scale.exp() if exists(self.residual_scale) else 1.
+            out = out + residual * residual_scale
+
         return out, next_stack_states