ASAC-pytorch 0.0.7__tar.gz → 0.0.10__tar.gz

{asac_pytorch-0.0.7 → asac_pytorch-0.0.10}/ASAC/ASAC.py

@@ -6,10 +6,10 @@ from torch import nn, tensor
 from torch.nn import Module, Linear, ModuleList
 import torch.nn.functional as F
 
-from einops import einsum, reduce
+from einops import einsum, reduce, rearrange
 from einops.layers.torch import Rearrange
 
-from x_transformers import Decoder
+from x_transformers import Decoder, AutoregressiveWrapper, TransformerWrapper
 
 from x_mlps_pytorch import MLP
 
@@ -29,8 +29,8 @@ def default(v, d):
 
 # return types
 
-AttentionReturn = namedtuple('AttentionReturn', ['attended', 'indices', 'aux_loss', 'aux_loss_breakdown'])
-ASACReturn = namedtuple('ASACReturn', ['logits', 'aux_loss', 'aux_loss_breakdown'])
+AttentionReturn = namedtuple('AttentionReturn', ['attended', 'indices', 'aux_loss', 'aux_loss_breakdown', 'attn_sim'])
+ASACReturn = namedtuple('ASACReturn', ['logits', 'aux_loss', 'aux_loss_breakdown', 'attn_sims', 'attn_schema_indices', 'attn_schema_autoregressive_loss'])
 
 # feedforward
 
@@ -65,7 +65,8 @@ class Attention(Module):
         heads = 8,
         k_rmsnorm = True,
         attn_schema: Module | None = None,
-        attn_add_residual = True # they had to add a residual for stability
+        attn_add_residual = True, # they had to add a residual for stability
+        stochastic_sample_attn = False
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -84,13 +85,16 @@ class Attention(Module):
         self.attn_schema = attn_schema
         self.attn_add_residual = attn_add_residual and attn_schema
 
+        self.stochastic_sample_attn = stochastic_sample_attn
+
         self.register_buffer('zero', tensor(0.), persistent = False)
 
     def forward(
         self,
         tokens, # (b h w d)
         pre_softmax_attn_gates = None,
-        post_softmax_attn_gates = None
+        post_softmax_attn_gates = None,
+        attn_schema_target = None
     ):
         tokens = self.norm(tokens)
 
@@ -114,7 +118,7 @@ class Attention(Module):
         indices = None
 
         if exists(self.attn_schema):
-            sim, indices, aux_loss, aux_loss_breakdown = self.attn_schema(orig_sim)
+            sim, indices, aux_loss, aux_loss_breakdown = self.attn_schema(orig_sim, target_sim = attn_schema_target)
 
         if self.attn_add_residual:
             sim = (sim + orig_sim) * 0.5
@@ -126,7 +130,10 @@ class Attention(Module):
 
         # attend
 
-        attn = sim.softmax(dim = -1)
+        if self.stochastic_sample_attn:
+            attn = F.gumbel_softmax(sim, tau = 1., hard = True, dim = -1)
+        else:
+            attn = sim.softmax(dim = -1)
 
         # modulate
 
@@ -144,7 +151,13 @@ class Attention(Module):
 
         attended = inverse_pack(attended)
 
-        return AttentionReturn(attended, indices, aux_loss, aux_loss_breakdown)
+        return AttentionReturn(
+            attended,
+            indices,
+            aux_loss,
+            aux_loss_breakdown,
+            orig_sim
+        )
 
 # attention autoencoder
 
@@ -186,7 +199,8 @@ class AttentionSchema(Module):
     def forward(
         self,
         attn_sim,
-        return_loss = None
+        return_loss = None,
+        target_sim = None
     ):
         return_loss = default(return_loss, self.training)
 
@@ -204,18 +218,20 @@ class AttentionSchema(Module):
 
         recon_loss = self.zero
 
+        target = default(target_sim, attn_sim)
+
         if return_loss:
             if self.detach_target:
-                attn_sim = attn_sim.detach()
+                target = target.detach()
 
             if self.kl_div_loss:
                 recon_loss = F.kl_div(
-                    attn_sim.log_softmax(dim = -1),
+                    target.log_softmax(dim = -1),
                     recon.softmax(dim = -1),
                     reduction = 'none'
                 ).sum(dim = -1).mean()
             else:
-                recon_loss = F.mse_loss(recon, attn_sim)
+                recon_loss = F.mse_loss(recon, target)
 
         # total
 
@@ -241,10 +257,15 @@ class ASAC(Module):
         vq_codebook_size = 256,
         recon_loss_weight = 1.,
         commit_loss_weight = 1.,
-        kl_div_loss = True
+        kl_div_loss = True,
+        stochastic_sample_attn = False,
+        awareness_model_depth = 2,
+        **awareness_model_kwargs
     ):
         super().__init__()
 
+        assert depth >= 2, 'depth must be at least 2'
+
         self.depth = depth
 
         self.to_embedding = to_embedding
@@ -263,16 +284,40 @@ class ASAC(Module):
             ) if use_asac and exists(seq_len) else None
 
             self.layers.append(ModuleList([
-                Attention(dim, dim_head = dim_head, heads = heads, attn_schema = attn_schema),
+                Attention(dim, dim_head = dim_head, heads = heads, attn_schema = attn_schema, stochastic_sample_attn = stochastic_sample_attn),
                 FeedForward(dim)
             ]))
 
+        # autoregressive awareness model (attention schema theory)
+
+        self.awareness_transformer = None
+        self.awareness_model = None
+
+        if use_asac and exists(seq_len):
+            self.awareness_transformer = TransformerWrapper(
+                num_tokens = vq_codebook_size,
+                max_seq_len = depth,
+                attn_layers = Decoder(
+                    dim = dim,
+                    depth = awareness_model_depth,
+                    heads = heads,
+                    rotary_pos_emb = True,
+                    **awareness_model_kwargs
+                )
+            )
+
+            self.awareness_model = AutoregressiveWrapper(self.awareness_transformer)
+
+        # zero buffer for auxiliary losses
+
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
         self.to_logits = nn.Sequential(
             nn.RMSNorm(dim),
             Linear(dim, num_classes)
         )
 
-    def forward(self, x):
+    def forward(self, x, attn_schema_targets = None):
         x = self.to_embedding(x)
 
         if exists(self.pos_embedding):
@@ -280,8 +325,16 @@ class ASAC(Module):
 
         total_aux_loss = total_recon_loss = total_commit_loss = 0.
 
-        for attn, ff in self.layers:
-            attn_out, indices, aux_loss, (recon_loss, commit_loss) = attn(x)
+        attn_schema_targets = default(attn_schema_targets, [None] * self.depth)
+        attn_sims = []
+        attn_schema_indices = []
+
+        for (attn, ff), target in zip(self.layers, attn_schema_targets):
+            attn_out, indices, aux_loss, (recon_loss, commit_loss), attn_sim = attn(x, attn_schema_target = target)
+
+            attn_sims.append(attn_sim)
+            if exists(indices):
+                attn_schema_indices.append(indices)
 
             x = attn_out + x
             x = ff(x) + x
@@ -294,4 +347,51 @@ class ASAC(Module):
 
         logits = self.to_logits(x)
 
-        return ASACReturn(logits, total_aux_loss, (total_recon_loss / self.depth, total_commit_loss / self.depth))
+        attn_schema_autoregressive_loss = self.zero
+
+        if attn_schema_indices:
+            attn_schema_indices = rearrange(attn_schema_indices, 'depth b ... -> b (depth ...)')
+
+            if exists(self.awareness_model):
+                attn_schema_autoregressive_loss = self.awareness_model(attn_schema_indices)
+        else:
+            attn_schema_indices = None
+
+        return ASACReturn(
+            logits,
+            total_aux_loss,
+            (total_recon_loss / self.depth, total_commit_loss / self.depth),
+            attn_sims,
+            attn_schema_indices,
+            attn_schema_autoregressive_loss
+        )
+
+class EMA_ASAC(Module):
+    def __init__(
+        self,
+        asac_model,
+        ema_decay = 0.999,
+        **ema_kwargs
+    ):
+        super().__init__()
+        self.asac = asac_model
+        self.ema_model = EMA(asac_model, beta = ema_decay, **ema_kwargs)
+
+    def update(self):
+        self.ema_model.update()
+
+    def forward(self, x, use_ema = False):
+        if use_ema:
+            return self.ema_model(x)
+
+        if not self.training:
+            return self.asac(x)
+
+        # get EMA targets
+        with torch.no_grad():
+            self.ema_model.eval()
+            ema_outputs = self.ema_model(x)
+            ema_targets = [sim.detach() for sim in ema_outputs.attn_sims]
+            self.ema_model.train()
+
+        return self.asac(x, attn_schema_targets = ema_targets)

asac_pytorch-0.0.10/ASAC/__init__.py

@@ -0,0 +1 @@
+from ASAC.ASAC import ASAC, PatchEmbedding, EMA_ASAC

{asac_pytorch-0.0.7 → asac_pytorch-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ASAC-pytorch
-Version: 0.0.7
+Version: 0.0.10
 Summary: Implementation of Attention Schema-based Attention Control (ASAC)
 Project-URL: Homepage, https://pypi.org/project/ASAC/
 Project-URL: Repository, https://codeberg.org/lucidrains/ASAC

{asac_pytorch-0.0.7 → asac_pytorch-0.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ASAC-pytorch"
-version = "0.0.7"
+version = "0.0.10"
 description = "Implementation of Attention Schema-based Attention Control (ASAC)"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

asac_pytorch-0.0.7/ASAC/__init__.py

@@ -1 +0,0 @@
-from ASAC.ASAC import ASAC, PatchEmbedding