Update V2.py

V2.py

@@ -138,12 +138,10 @@ class HyperConv1D(layers.Layer):
         # Input projection
         self.input_proj = layers.Dense(d_model, name="input_proj")
 
-        # Dynamic kernel conv (중간 차원)
-        self.d_mid = max(64, d_model // 8)
-        self.dynamic_dense = layers.Dense(self.d_mid, activation='silu')
+        # 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')
-        self.kernel_temp = self.add_weight("kernel_temp", shape=(), initializer=tf.constant_initializer(1.0), trainable=True)
 
         # Hypernetwork: token-wise transform before pooling
         self.hyper = tf.keras.Sequential([
@@ -152,13 +150,13 @@ class HyperConv1D(layers.Layer):
         ], name="hyper")
 
         self.attn_pool = layers.Dense(1)
-        self.scale_dense = layers.Dense(d_model, bias_initializer=tf.keras.initializers.Constant(0.0))
+        self.scale_dense = layers.Dense(d_model)
 
-        # Normalization + dropout
         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
@@ -172,9 +170,8 @@ class HyperConv1D(layers.Layer):
         # ------------------------------
         # 2) DynamicConv local mixing
         # ------------------------------
-        kernels = self.kernel_generator(self.dynamic_dense(x_proj))  # (B,L,k)
-        kernels = tf.cast(kernels, tf.float32)
-        kernels = tf.nn.softmax(kernels / tf.maximum(self.kernel_temp, 1e-6), axis=-1)
+        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)
@@ -186,7 +183,7 @@ class HyperConv1D(layers.Layer):
             padding='VALID'
         )
         patches = tf.reshape(patches, [B, L, self.k, D])
-        kernels_exp = tf.cast(tf.expand_dims(kernels, axis=-1), x_proj.dtype)
+        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)
 
@@ -198,12 +195,11 @@ class HyperConv1D(layers.Layer):
         global_z = tf.nn.softmax(global_z, axis=1)
         global_z = tf.reduce_sum(h * global_z, axis=1)
 
-        # residual-gate 스타일 scale: 1 + α*tanh(...)
-        scale = 1.0 + 0.1 * tf.tanh(self.scale_dense(global_z))
-        out_local = out_local * tf.expand_dims(scale, 1)
+        scale = tf.expand_dims(tf.nn.sigmoid(self.scale_dense(global_z)), 1)
+        out_local = out_local * scale
 
         # ------------------------------
-        # 4) Residual + SiLU + LayerNorm + Dropout
+        # 4) Residual + SiLU + LayerNorm
         # ------------------------------
         out = x_proj + out_local
         out = tf.nn.silu(out)
@@ -211,8 +207,7 @@ class HyperConv1D(layers.Layer):
         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)