Update AlphaS2S.py

AlphaS2S.py

@@ -226,6 +226,18 @@ class gMLPBlock(layers.Layer):
         
         return residual + out
 
+class CrossBlock(layers.Layer):
+    def __init__(self, d_model): # 💡 d_model 인자 추가
+        super().__init__()
+        # 💡 수정: 출력 차원을 1에서 d_model로 변경 (채널별 게이팅 허용)
+        self.alpha = layers.Dense(d_model, activation='sigmoid', dtype='float32') 
+    def call(self, x, z):
+        # a의 shape: (Batch, Seq_len, D_model)
+        a = self.alpha(x)
+        # y: 각 채널이 독립적인 가중치 (a)로 X와 Z를 융합
+        y = a * x + (1.0 - a) * z
+        return y
+
 
 class LoU(layers.Layer):
     def __init__(self, d_model, clip_value=5.0, eps=1e-6):
@@ -241,6 +253,7 @@ class LoU(layers.Layer):
         
         self.alpha_linear = layers.Dense(1, activation='sigmoid', dtype='float32')
         self.glu = SwiGLU(d_model, d_model)
+        self.cross = CrossBlock(d_model)
 
     def _ema_over_time(self, score, alpha_dynamic):
         seq = tf.transpose(score, perm=[1, 0, 2])
@@ -285,7 +298,7 @@ class LoU(layers.Layer):
         
         # LoU 블록에서는 x_comb + residual 후 CrossBlock을 통과
         out = self.norm(x_comb + residual)
-        out = out + z
+        out = self.cross(out, z)
         out = self.glu(out)
         return tf.cast(out, x.dtype)