Update AlphaS2S.py

AlphaS2S.py

@@ -240,29 +240,10 @@ class CrossBlock(layers.Layer):
         super().__init__()
         self.clip_value = clip_value
         self.eps = eps
+        self.attn = layers.MultiHeadAttention(8, 20)
         # 💡 수정: 출력 차원을 1에서 d_model로 변경
     def call(self, x, z):
-        # a의 shape: (Batch, Seq_len, D_model)
-        g_q = (tf.nn.tanh(x) + 1.0) / 2.0
-        g_k = (tf.nn.tanh(z) + 1.0) / 2.0
-        score = (g_q * g_k)
-        score = tf.cumsum(score, axis=1)
-
-        seq_len = tf.shape(score)[1]
-        # [1, 2, 3, ..., L]을 D_model 차원으로 확장
-        count_for_mean = tf.cast(tf.range(seq_len) + 1, score.dtype)
-        count_for_mean = tf.reshape(count_for_mean, (1, seq_len, 1))
-        
-        # 누적합을 현재까지의 토큰 개수로 나누어 평균 누적합 계산 (B, L, D)
-        score_mean = score / count_for_mean
-        
-        # 정규화 분모 설정
-        denom = tf.maximum(score_mean, self.eps)
-        score_norm = score / denom
-        # -----------------------------------------------
-
-        score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
-        y = score_clipped * z
+        y = self.attn(x, z, z)
         return y
 
 class LoU(layers.Layer):
@@ -274,7 +255,7 @@ class LoU(layers.Layer):
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         
-        self.glu = SwiGLU(d_model, 320)
+        self.glu = SwiGLU(d_model, 350)
         self.cross = CrossBlock()
 
     def call(self, x, z):
@@ -327,7 +308,6 @@ class AlphaS2S(tf.keras.Model):
         
         # 디코더 실행
         y = self.dec_embedding(dec_inputs) + self.dec_pos_embedding(dec_pos)
-        # Note: LoU는 내부적으로 EMA를 사용하며, 일반적인 Cross-Attention 블록의 역할을 수행
         for layer in self.dec_layers: y = layer(y, enc_out, training=training) 
         
         return self.final_layer(y)