go for an improvised solution for what may be an issue with autoregressive and laser

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.6.3"
+version = "0.6.4"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

train_latent_with_text.py

@@ -95,7 +95,7 @@ def encode_tokens(str: str) -> Tensor:
         dim = 128,
         depth = 8,
         dim_head = 64,
-        heads = 8
+        heads = 8,
     )
 ).cuda()
 

train_mnist.py

@@ -65,7 +65,7 @@ def forward(self, x):
         dim = 64,
         depth = 4,
         dim_head = 32,
-        heads = 8
+        heads = 8,
     )
 ).cuda()
 

train_mnist_vae.py

@@ -127,7 +127,7 @@ def forward(self, x):
         dim = 64,
         depth = 4,
         dim_head = 32,
-        heads = 8
+        heads = 8,
     )
 ).cuda()
 

train_mnist_with_unet.py

@@ -61,7 +61,7 @@ def forward(self, x):
         dim = 64,
         depth = 4,
         dim_head = 32,
-        heads = 8
+        heads = 8,
     )
 ).to(device)
 

transfusion_pytorch/transfusion.py

@@ -757,6 +757,7 @@ def __init__(
         use_flex_attn = False,
         gate_values = True,
         laser = False,
+        laser_softclamp_value = 15.,
         learned_value_residual_mix = False
     ):
         super().__init__()
@@ -785,6 +786,7 @@ def __init__(
         self.softcap_value = softcap_value
 
         self.laser = laser
+        self.laser_softclamp_value = laser_softclamp_value
 
         self.dropout = nn.Dropout(dropout)
 
@@ -850,8 +852,8 @@ def forward(
         # laser attention
 
         if self.laser:
-            v_max = v.amax(dim = -2, keepdim = True).detach()
-            v = (v - v_max).exp()
+            v = softclamp(v, self.laser_softclamp_value)
+            v = v.exp()
 
         # whether to use flex attention or not
 
@@ -890,7 +892,7 @@ def forward(
         # laser attention
 
         if self.laser:
-            out = log(out) + v_max
+            out = log(out)
 
         # maybe gate values