ASAC-pytorch 0.0.2__tar.gz → 0.0.4__tar.gz

asac_pytorch-0.0.4/ASAC/ASAC.py

@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+from collections import namedtuple
+import torch
+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.layers.torch import Rearrange
+
+from x_transformers import Decoder
+
+from x_mlps_pytorch import MLP
+
+from vector_quantize_pytorch import VectorQuantize
+
+from ema_pytorch import EMA
+
+from torch_einops_utils import pack_with_inverse, maybe
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# return types
+
+AttentionReturn = namedtuple('AttentionReturn', ['attended', 'indices', 'aux_loss', 'aux_loss_breakdown'])
+ASACReturn = namedtuple('ASACReturn', ['logits', 'aux_loss', 'aux_loss_breakdown'])
+
+# feedforward
+
+def FeedForward(dim, expansion_factor = 4.):
+    dim_inner = int(dim * expansion_factor)
+    return nn.Sequential(
+        nn.RMSNorm(dim),
+        nn.Linear(dim, dim_inner),
+        nn.GELU(),
+        nn.Linear(dim_inner, dim)
+    )
+
+# embedding
+
+def PatchEmbedding(dim, patch_size, channels = 3):
+    patch_dim = channels * (patch_size ** 2)
+
+    return nn.Sequential(
+        Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+        nn.RMSNorm(patch_dim),
+        Linear(patch_dim, dim),
+        nn.RMSNorm(dim),
+    )
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        dim_head = 64,
+        heads = 8,
+        k_rmsnorm = True,
+        attn_schema: Module | None = None,
+        attn_add_residual = True # they had to add a residual for stability
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        dim_inner = dim_head * heads
+
+        self.norm = nn.RMSNorm(dim)
+
+        self.to_qkv = Linear(dim, dim_inner * 3, bias = False)
+        self.combine_heads = Linear(dim_inner, dim, bias = False)
+
+        self.k_rmsnorm = nn.RMSNorm(dim_head) if k_rmsnorm else None
+
+        self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.merge_heads = Rearrange('b h n d -> b n (h d)')
+
+        self.attn_schema = attn_schema
+        self.attn_add_residual = attn_add_residual and attn_schema
+
+        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
+    ):
+        tokens = self.norm(tokens)
+
+        tokens, inverse_pack = pack_with_inverse(tokens, 'b * d')
+
+        q, k, v = self.to_qkv(tokens).chunk(3, dim = -1)
+        q, k, v = (self.split_heads(t) for t in (q, k, v))
+
+        k = maybe(self.k_rmsnorm)(k)
+
+        q = q * self.scale
+
+        sim = einsum(q, k, 'b h i d, b h j d -> b h i j')
+
+        orig_sim = sim
+
+        # the proposal
+
+        aux_loss = self.zero
+        aux_loss_breakdown = (self.zero, self.zero)
+        indices = None
+
+        if exists(self.attn_schema):
+            sim, indices, aux_loss, aux_loss_breakdown = self.attn_schema(orig_sim)
+
+        if self.attn_add_residual:
+            sim = sim + orig_sim
+
+        # modulate
+
+        if exists(pre_softmax_attn_gates):
+            sim = sim + pre_softmax_attn_gates
+
+        # attend
+
+        attn = sim.softmax(dim = -1)
+
+        # modulate
+
+        if exists(post_softmax_attn_gates):
+            attn = attn * post_softmax_attn_gates
+
+        # aggregate and combine out
+
+        out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
+
+        out = self.merge_heads(out)
+        attended = self.combine_heads(out)
+
+        # bring back the packed dimensions
+
+        attended = inverse_pack(attended)
+
+        return AttentionReturn(attended, indices, aux_loss, aux_loss_breakdown)
+
+# attention autoencoder
+
+class AttentionSchema(Module):
+    def __init__(
+        self,
+        dim,
+        dim_bottleneck,
+        kl_div_loss = True,
+        detach_target = True,
+        encoder: Module | None = None,
+        decoder: Module | None = None,
+        recon_loss_weight = 1.,
+        commit_loss_weight = 1.,
+        **vq_kwargs
+    ):
+        super().__init__()
+
+        if not exists(encoder):
+            encoder = MLP(dim, dim_bottleneck, activation = nn.LeakyReLU())
+
+        self.encoder = encoder
+
+        self.vq = VectorQuantize(dim_bottleneck, **vq_kwargs)
+
+        if not exists(decoder):
+            decoder = MLP(dim_bottleneck, dim, activation = nn.LeakyReLU())
+
+        self.decoder = decoder
+
+        self.kl_div_loss = kl_div_loss
+        self.detach_target = detach_target
+
+        self.recon_loss_weight = recon_loss_weight
+        self.commit_loss_weight = commit_loss_weight
+
+    def forward(
+        self,
+        attn_sim,
+        return_loss = None
+    ):
+        return_loss = default(return_loss, self.training)
+
+        attn_features, inverse_pack = pack_with_inverse(attn_sim, 'b *')
+
+        encoded = self.encoder(attn_features)
+
+        quantized, indices, commit_loss = self.vq(encoded)
+
+        decoded = self.decoder(quantized)
+
+        recon = inverse_pack(decoded)
+
+        # loss, mse as in paper or reverse kl
+
+        if return_loss:
+            if self.detach_target:
+                attn_sim = attn_sim.detach()
+
+            if self.kl_div_loss:
+                recon_loss = F.kl_div(
+                    attn_sim.log_softmax(dim = -1),
+                    recon.softmax(dim = -1),
+                    reduction = 'batchmean'
+                )
+            else:
+                recon_loss = F.mse_loss(recon, attn_sim)
+
+        # total
+
+        total_loss = recon_loss * self.recon_loss_weight + commit_loss * self.commit_loss_weight
+
+        return recon, indices, total_loss, (recon_loss, commit_loss)
+
+# class
+
+class ASAC(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        heads,
+        to_embedding,
+        seq_len = None,
+        dim_head = 64,
+        num_classes = 10,
+        use_asac = False,
+        dim_bottleneck = 256,
+        vq_codebook_size = 256,
+        recon_loss_weight = 1.,
+        commit_loss_weight = 1.
+    ):
+        super().__init__()
+
+        self.depth = depth
+
+        self.to_embedding = to_embedding
+        self.pos_embedding = nn.Parameter(torch.randn(seq_len, dim)) if exists(seq_len) else None
+
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            attn_schema = AttentionSchema(
+                dim = heads * (seq_len ** 2),
+                dim_bottleneck = dim_bottleneck,
+                codebook_size = vq_codebook_size,
+                recon_loss_weight = recon_loss_weight,
+                commit_loss_weight = commit_loss_weight
+            ) 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),
+                FeedForward(dim)
+            ]))
+
+        self.to_logits = nn.Sequential(
+            nn.RMSNorm(dim),
+            Linear(dim, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.to_embedding(x)
+
+        if exists(self.pos_embedding):
+            x = x + self.pos_embedding
+
+        total_aux_loss = 0.
+        total_recon_loss = 0.
+        total_commit_loss = 0.
+
+        for attn, ff in self.layers:
+            attn_out, indices, aux_loss, (recon_loss, commit_loss) = attn(x)
+
+            x = attn_out + x
+            x = ff(x) + x
+
+            total_aux_loss = total_aux_loss + aux_loss
+            total_recon_loss = total_recon_loss + recon_loss
+            total_commit_loss = total_commit_loss + commit_loss
+
+        x = reduce(x, 'b n d -> b d', 'mean')
+
+        logits = self.to_logits(x)
+
+        return ASACReturn(logits, total_aux_loss, (total_recon_loss / self.depth, total_commit_loss / self.depth))

asac_pytorch-0.0.4/ASAC/__init__.py

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ASAC-pytorch
-Version: 0.0.2
+Version: 0.0.4
 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
@@ -55,12 +55,12 @@ Implementation of [Attention Schema-based Attention Control (ASAC)](https://arxi
 
 ```bibtex
 @misc{saxena2025attentionschemabasedattentioncontrol,
-    title   = {Attention Schema-based Attention Control (ASAC): A Cognitive-Inspired Approach for Attention Management in Transformers}, 
+    title   = {Attention Schema-based Attention Control (ASAC): A Cognitive-Inspired Approach for Attention Management in Transformers},
     author   = {Krati Saxena and Federico Jurado Ruiz and Guido Manzi and Dianbo Liu and Alex Lamb},
     year     = {2025},
     eprint   = {2509.16058},
     archivePrefix = {arXiv},
     primaryClass = {cs.AI},
-    url      = {https://arxiv.org/abs/2509.16058}, 
+    url      = {https://arxiv.org/abs/2509.16058},
 }
 ```

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

@@ -6,12 +6,12 @@ Implementation of [Attention Schema-based Attention Control (ASAC)](https://arxi
 
 ```bibtex
 @misc{saxena2025attentionschemabasedattentioncontrol,
-    title   = {Attention Schema-based Attention Control (ASAC): A Cognitive-Inspired Approach for Attention Management in Transformers}, 
+    title   = {Attention Schema-based Attention Control (ASAC): A Cognitive-Inspired Approach for Attention Management in Transformers},
     author   = {Krati Saxena and Federico Jurado Ruiz and Guido Manzi and Dianbo Liu and Alex Lamb},
     year     = {2025},
     eprint   = {2509.16058},
     archivePrefix = {arXiv},
     primaryClass = {cs.AI},
-    url      = {https://arxiv.org/abs/2509.16058}, 
+    url      = {https://arxiv.org/abs/2509.16058},
 }
 ```

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

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

asac_pytorch-0.0.2/ASAC/ASAC.py

@@ -1,165 +0,0 @@
-from __future__ import annotations
-
-import torch
-from torch import nn, tensor
-from torch.nn import Module, Linear
-import torch.nn.functional as F
-
-from einops import einsum
-from einops.layers.torch import Rearrange
-
-from x_transformers import Decoder
-
-from x_mlps_pytorch import MLP
-
-from vector_quantize_pytorch import VectorQuantize
-
-from ema_pytorch import EMA
-
-from torch_einops_utils import pack_with_inverse
-
-# helpers
-
-def exists(v):
-    return v is not None
-
-def default(v, d):
-    return v if exists(v) else d
-
-# attention
-
-class Attention(Module):
-    def __init__(
-        self,
-        dim,
-        dim_head = 64,
-        heads = 8,
-        attn_schema: Module | None = None,
-        attn_add_residual = True # they had to add a residual for stability
-    ):
-        super().__init__()
-        self.scale = dim_head ** -0.5
-        dim_inner = dim_head * heads
-
-        self.to_qkv = Linear(dim, dim_inner * 3, bias = False)
-        self.combine_heads = Linear(dim_inner, dim, bias = False)
-
-        self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
-        self.merge_heads = Rearrange('b h n d -> b n (h d)')
-
-        self.attn_schema = attn_schema
-        self.attn_add_residual = attn_add_residual and attn_schema
-
-        self.register_buffer('zero', tensor(0.), persistent = False)
-
-    def forward(
-        self,
-        tokens, # (b h w d)
-    ):
-        tokens, inverse_pack = pack_with_inverse(tokens, 'b * d')
-
-        q, k, v = self.to_qkv(tokens).chunk(3, dim = -1)
-        q, k, v = (self.split_heads(t) for t in (q, k, v))
-
-        q = q * self.scale
-
-        sim = einsum(q, k, 'b h i d, b h j d -> b h i j')
-
-        orig_sim = sim
-
-        # the proposal
-
-        aux_loss = self.zero
-
-        if exists(self.attn_schema):
-            sim, indices, aux_loss = self.attn_schema(orig_sim)
-
-        if self.attn_add_residual:
-            sim = sim + orig_sim
-
-        # attend
-
-        attn = sim.softmax(dim = -1)
-
-        # aggregate and combine out
-
-        out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
-
-        out = self.merge_heads(out)
-        attended = self.combine_heads(out)
-
-        # bring back the packed dimensions
-
-        attended = inverse_pack(attended)
-
-        return attended, indices, aux_loss
-
-# attention autoencoder
-
-class AttentionSchema(Module):
-    def __init__(
-        self,
-        dim,
-        dim_bottleneck,
-        kl_div_loss = True,
-        detach_target = True,
-        **vq_kwargs
-    ):
-        super().__init__()
-        self.encoder = MLP(dim, dim_bottleneck, activation = nn.LeakyReLU())
-
-        self.vq = VectorQuantize(dim_bottleneck, **vq_kwargs)
-
-        self.decoder = MLP(dim_bottleneck, dim, activation = nn.LeakyReLU())
-
-        self.kl_div_loss = kl_div_loss
-        self.detach_target = detach_target
-
-    def forward(
-        self,
-        attn_sim,
-        return_loss = None
-    ):
-        return_loss = default(return_loss, self.training)
-
-        attn_features, inverse_pack = pack_with_inverse(attn_sim, 'b *')
-
-        encoded = self.encoder(attn_features)
-
-        quantized, indices, commit_loss = self.vq(encoded)
-
-        decoded = self.decoder(quantized)
-
-        recon = inverse_pack(decoded)
-
-        # loss, mse as in paper or reverse kl
-
-        if return_loss:
-            if self.detach_target:
-                attn_sim = attn_sim.detach()
-
-            if self.kl_div_loss:
-                recon_loss = F.kl_div(
-                    attn_sim.log_softmax(dim = -1),
-                    recon,
-                    reduction = 'batchmean'
-                )
-            else:
-                recon_loss = F.mse_loss(recon, attn_sim)
-
-        # total
-
-        total_loss = recon_loss + commit_loss
-
-        loss_breakdown = (recon_loss, commit_loss)
-
-        return recon, indices, total_loss
-
-# class
-
-class ASAC(Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        return x

asac_pytorch-0.0.2/ASAC/__init__.py

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