Update mcdpgdn.py

mcdpgdn.py

@@ -1,11 +1,5 @@
 import torch
-import numpy as np
 from torch import nn
-from torch.nn import functional as F
-import einops
-from einops.layers.torch import Rearrange
-import math
-
 
 def l2norm(x, dim=-1, eps=1e-6):
     return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
@@ -24,24 +18,16 @@ class GatedDeltaNet(nn.Module):
         self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
         self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
         self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
-        ####################################################
-        ### NEW: Gates for delta rule and output gating
+     
         self.W_gate = nn.Linear(d_in, d_out, bias=False)
         self.W_beta = nn.Linear(d_in, d_out, bias=False)
 
-        # Note: The decay gate alpha corresponds to
-        # A_log + W_alpha(x) + dt_bias
         self.W_alpha = nn.Linear(d_in, num_heads, bias=False)
         self.dt_bias = nn.Parameter(torch.ones(num_heads))
         A_init = torch.empty(num_heads).uniform_(0, 16)
         self.A_log = nn.Parameter(torch.log(A_init))
-        # We could implement this as
-        # W_alpha = nn.Linear(d_in, num_heads, bias=True)
-        # but the bias is separate for interpretability and
-        # to mimic the official implementation
-
+       
         self.norm = nn.RMSNorm(self.head_dim, eps=1e-6)
-        ####################################################
 
         self.out_proj = nn.Linear(d_out, d_out)
         self.dropout = nn.Dropout(dropout)
@@ -51,38 +37,32 @@ class GatedDeltaNet(nn.Module):
         queries = self.W_query(x)
         keys = self.W_key(x)
         values = self.W_value(x)
-        ####################################################
-        ### NEW: Compute delta rule gates
+       
         beta = torch.sigmoid(self.W_beta(x))
         alpha_log = -self.A_log.exp().view(1, 1, -1) * F.softplus(
             self.W_alpha(x) + self.dt_bias
         )
         alpha = alpha_log.exp()
         gate = self.W_gate(x)
-        ####################################################
-
+       
         keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
         values = values.view(b, num_tokens, self.num_heads, self.head_dim)
         queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
         beta = beta.view(b, num_tokens, self.num_heads, self.head_dim)
-        gate = gate.view(b, num_tokens, self.num_heads, self.head_dim)  # NEW
+        gate = gate.view(b, num_tokens, self.num_heads, self.head_dim)  
 
         keys = keys.transpose(1, 2)
         queries = queries.transpose(1, 2)
         values = values.transpose(1, 2)
         beta = beta.transpose(1, 2)
 
-        ####################################################
-        ### NEW: QKNorm-like normalization for delta rule
         queries = l2norm(queries, dim=-1) / (self.head_dim ** 0.5)
         keys = l2norm(keys, dim=-1)
-        ####################################################
-
+       
         S = x.new_zeros(b, self.num_heads, self.head_dim, self.head_dim)
 
         outs = []
-        ####################################################
-        ### NEW: Gated delta rule update
+       
         for t in range(num_tokens):
             k_t = keys[:, :, t]
             q_t = queries[:, :, t]
@@ -95,18 +75,15 @@ class GatedDeltaNet(nn.Module):
             delta = (v_t - kv_mem) * b_t
             S = S + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
             y_t = (S * q_t.unsqueeze(-1)).sum(dim=-2)
-            ####################################################
+          
             outs.append(y_t)
 
         context = torch.stack(outs, dim=2).transpose(1, 2).contiguous()
         context = context.view(b, num_tokens, self.num_heads, self.head_dim)
 
-        ####################################################
-        ### NEW: Apply RMSNorm and SiLU gate
         context = self.norm(context)
         context = context * F.silu(gate)
-        ####################################################
-
+      
         context = context.view(b, num_tokens, self.d_out)
         context = self.dropout(context)
         out = self.out_proj(context)