Update V2.py

V2.py

@@ -138,10 +138,12 @@ class HyperConv1D(layers.Layer):
         # Input projection
         self.input_proj = layers.Dense(d_model, name="input_proj")
 
-        # Dynamic kernel conv
-        self.dynamic_dense = layers.Dense(d_model, activation='silu')
+        # Dynamic kernel conv (중간 차원)
+        self.d_mid = max(64, d_model // 8)
+        self.dynamic_dense = layers.Dense(self.d_mid, 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([
@@ -150,13 +152,13 @@ class HyperConv1D(layers.Layer):
         ], name="hyper")
 
         self.attn_pool = layers.Dense(1)
-        self.scale_dense = layers.Dense(d_model)
+        self.scale_dense = layers.Dense(d_model, bias_initializer=tf.keras.initializers.Constant(0.0))
 
+        # 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
@@ -170,8 +172,9 @@ class HyperConv1D(layers.Layer):
         # ------------------------------
         # 2) DynamicConv local mixing
         # ------------------------------
-        kernels = self.kernel_generator(self.dynamic_dense(x_proj))  # (B, L, k)
-        kernels = tf.nn.softmax(kernels, axis=-1)
+        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)
 
         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)
@@ -183,7 +186,7 @@ class HyperConv1D(layers.Layer):
             padding='VALID'
         )
         patches = tf.reshape(patches, [B, L, self.k, D])
-        kernels_exp = tf.expand_dims(kernels, axis=-1)
+        kernels_exp = tf.cast(tf.expand_dims(kernels, axis=-1), x_proj.dtype)
         out_local = tf.reduce_sum(patches * kernels_exp, axis=2)  # (B,L,D)
         out_local = self.dynamic_proj(out_local)
 
@@ -195,11 +198,12 @@ class HyperConv1D(layers.Layer):
         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)
-        out_local = out_local * scale
+        # 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)
 
         # ------------------------------
-        # 4) Residual + SiLU + LayerNorm
+        # 4) Residual + SiLU + LayerNorm + Dropout
         # ------------------------------
         out = x_proj + out_local
         out = tf.nn.silu(out)