Update AlphaS2S.py

AlphaS2S.py

@@ -247,9 +247,20 @@ class CrossBlock(layers.Layer):
         g_k = (tf.nn.tanh(z) + 1.0) / 2.0
         score = (g_q * g_k)
         score = tf.cumsum(score, axis=1)
-        mean_last = tf.reduce_mean(score, axis=-1, keepdims=True)
-        denom = tf.maximum(mean_last, self.eps)
+
+        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
         return y
@@ -269,7 +280,6 @@ class LoU(layers.Layer):
         self.glu = SwiGLU(d_model, 320)
         self.cross = CrossBlock()
 
-    # LoU는 원래 Uni-directional Attention/Recurrent Block 역할
     def call(self, x, z):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
@@ -282,18 +292,30 @@ class LoU(layers.Layer):
         g_k = (tf.nn.tanh(k) + 1.0) / 2.0
         score = g_q * g_k
 
-        score = tf.cumsum(score, axis=1)
-        mean_last = tf.reduce_mean(score, axis=-1, keepdims=True)
-        denom = tf.maximum(mean_last, self.eps)
+        score = tf.cumsum(score, axis=1) # (B, L, D)
+        
+        # 💡 수정된 부분: 현재 토큰까지의 누적합 평균으로 정규화
+        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)
         x_comb = score_clipped * V
         
-        # LoU 블록에서는 x_comb + residual 후 CrossBlock을 통과
         out = self.norm(x_comb + residual)
         out = self.cross(out, z)
         out = self.glu(out)
         return tf.cast(out, x.dtype)
+
         
 # =======================
 # 4) AlphaS2S 모델 (기존 코드 유지)