Update V2.py

V2.py

@@ -128,57 +128,87 @@ ds = ds.batch(BATCH_SIZE, drop_remainder=True)
 ds = ds.map(lambda v1, v2: ((v1, v2), tf.zeros([BATCH_SIZE], dtype=tf.float32)), num_parallel_calls=tf.data.AUTOTUNE)
 ds = ds.prefetch(tf.data.AUTOTUNE)
 
-
-class HyperConv1D(layers.Layer):
-    def __init__(self, d_model, k=7, mem_size=64, hyper_dim=128, dropout=0.0):
+class HyperDynamicConv1D(layers.Layer):
+    def __init__(self, d_model, k=7, hyper_dim=128, dropout=0.0):
         super().__init__()
         assert k % 2 == 1
         self.k = k
         self.d_model = d_model
-        self.mem_size = mem_size
 
         # Input projection
         self.input_proj = layers.Dense(d_model, name="input_proj")
 
-        # Local depthwise conv
-        self.local_conv = layers.DepthwiseConv1D(kernel_size=k, padding='same', activation='silu')
-        self.local_proj = layers.Dense(d_model, name="local_proj")
+        # Dynamic kernel conv
+        self.dynamic_dense = layers.Dense(d_model, activation='silu')
+        self.dynamic_proj = layers.Dense(d_model)
+        self.kernel_generator = layers.Dense(k, dtype='float32')
 
-        # Hypernetwork: global -> scale vector
+        # Hypernetwork: token-wise transform before pooling
         self.hyper = tf.keras.Sequential([
             layers.Dense(hyper_dim, activation='gelu'),
             layers.Dense(d_model)
         ], name="hyper")
 
-
-        self.norm = layers.LayerNormalization()
         self.attn_pool = layers.Dense(1)
         self.scale_dense = layers.Dense(d_model)
 
-    def call(self, x):
+        self.norm = layers.LayerNormalization()
+        self.dropout = layers.Dropout(dropout)
+
+    def call(self, x, training=None):
         x_in = x
         x_dtype = x.dtype
 
         # 1) input projection
         x_proj = self.input_proj(x)
 
-
-        out_local = self.local_conv(x_proj)
-
+        B = tf.shape(x_proj)[0]
+        L = tf.shape(x_proj)[1]
+        D = self.d_model
+        pad = (self.k - 1) // 2
+
+        # ------------------------------
+        # 2) DynamicConv local mixing
+        # ------------------------------
+        kernels = self.kernel_generator(self.dynamic_dense(x_proj))  # (B, L, k)
+        kernels = tf.nn.softmax(kernels, axis=-1)
+
+        x_pad = tf.pad(x_proj, [[0,0],[pad,pad],[0,0]])
+        x_pad_4d = tf.expand_dims(x_pad, axis=1)  # (B,1,L+k-1,D)
+        patches = tf.image.extract_patches(
+            images=x_pad_4d,
+            sizes=[1,1,self.k,1],
+            strides=[1,1,1,1],
+            rates=[1,1,1,1],
+            padding='VALID'
+        )
+        patches = tf.reshape(patches, [B, L, self.k, D])
+        kernels_exp = tf.expand_dims(kernels, axis=-1)
+        out_local = tf.reduce_sum(patches * kernels_exp, axis=2)  # (B,L,D)
+        out_local = self.dynamic_proj(out_local)
+
+        # ------------------------------
+        # 3) Hyper scaling
+        # ------------------------------
         h = self.hyper(x_proj)
         global_z = self.attn_pool(h)
         global_z = tf.nn.softmax(global_z, axis=1)
-        global_z = tf.reduce_sum(x_proj * global_z, axis=1)
+        global_z = tf.reduce_sum(h * global_z, axis=1)
 
         scale = tf.expand_dims(tf.nn.sigmoid(self.scale_dense(global_z)), 1)
         out_local = out_local * scale
-        out_local = self.local_proj(out_local)
 
+        # ------------------------------
+        # 4) Residual + SiLU + LayerNorm
+        # ------------------------------
         out = x_proj + out_local
         out = tf.nn.silu(out)
         out = self.norm(out)
+        out = self.dropout(out, training=training)
+
         return tf.cast(out, x_dtype)
 
+
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):
         super().__init__(**kwargs)