Update V2.py

V2.py

@@ -128,6 +128,7 @@ 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, hyper_dim=128, dropout=0.0):
         super().__init__()
@@ -143,70 +144,66 @@ class HyperConv1D(layers.Layer):
         self.dynamic_proj = layers.Dense(d_model)
         self.kernel_generator = layers.Dense(k, dtype='float32')
 
-        # Hypernetwork: token-wise transform before pooling
+        # Hypernetwork for token-wise features
         self.hyper = tf.keras.Sequential([
             layers.Dense(hyper_dim, activation='gelu'),
             layers.Dense(d_model)
         ], name="hyper")
 
+        # Attention pooling for global context
         self.attn_pool = layers.Dense(1)
-        self.scale_dense = layers.Dense(d_model)
 
+        # LayerNorm + Dropout
         self.norm = layers.LayerNormalization()
         self.dropout = layers.Dropout(dropout)
+        self.dense = layers.Dense(d_model)
 
     def call(self, x, training=None):
         x_in = x
         x_dtype = x.dtype  # 입력 dtype 저장
 
-    # 1) input projection
+        # 1) Input projection
         x_proj = self.input_proj(x)  # (B, L, D)
-    
         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 생성 후 x_proj dtype으로 맞춤
+        # ------------------------------
+        # 2) DynamicConv local mixing
+        # ------------------------------
         kernels = self.kernel_generator(self.dynamic_dense(x_proj))
         kernels = tf.cast(kernels, x_proj.dtype)
         kernels = tf.nn.softmax(kernels, axis=-1)
 
-    # padding & patch 추출
         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)
+        x_pad_4d = tf.expand_dims(x_pad, axis=1)
         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'
-    )
+            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 shape 맞추기
         kernels_exp = tf.expand_dims(kernels, axis=-1)
-        out_local = tf.reduce_sum(patches * kernels_exp, axis=2)  # (B,L,D)
+        out_local = tf.reduce_sum(patches * kernels_exp, axis=2)
         out_local = self.dynamic_proj(out_local)
 
-    # ------------------------------
-    # 3) Hyper scaling
-    # ------------------------------
+        # ------------------------------
+        # 3) Global context via attention pooling (scale 제거)
+        # ------------------------------
         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(h * global_z, axis=1)
-
-        scale = tf.expand_dims(tf.nn.sigmoid(self.scale_dense(global_z)), 1)
-        scale = tf.cast(scale, x_proj.dtype)  # dtype 맞춤
-        out_local = out_local * scale
-
-    # ------------------------------
-    # 4) Residual + SiLU + LayerNorm
-    # ------------------------------
+        scores = tf.nn.softmax(self.attn_pool(h), axis=1)  # (B, L, 1)
+        global_context = tf.reduce_sum(h * scores, axis=1)   # (B, D)
+        # token-wise concat
+        global_context_exp = tf.expand_dims(global_context, 1) * tf.ones([B, L, 1], dtype=x_proj.dtype)
+        out_local = tf.concat([out_local, global_context_exp], axis=-1)
+        out_local = self.dense(out_local)  # dimension 맞춤
+
+        # ------------------------------
+        # 4) Residual + SiLU + LayerNorm
+        # ------------------------------
         out = x_proj + out_local
         out = tf.nn.silu(out)
         out = self.norm(out)