Update Mo.py

Mo.py

@@ -130,13 +130,22 @@ class LoU(layers.Layer):
         self.d_model = d_model
         self.clip_value = float(clip_value)
         self.eps = float(eps)
+
+        # Q/K/V 변환
         self.Q = layers.Dense(d_model, dtype='float32')
         self.K = layers.Dense(d_model, dtype='float32')
         self.V = layers.Dense(d_model, dtype='float32')
+
+        # 정규화
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
-        
+
+        # 비선형 표현력
         self.glu = SwiGLU(d_model, 3500)
+
+        # 학습 가능한 과거 토큰 가중치
+        self.alpha = self.add_weight(shape=(d_model,), initializer='ones', trainable=True)
+
     def call(self, x):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
@@ -145,34 +154,24 @@ class LoU(layers.Layer):
         q = self.Q(x_f32)
         k = self.K(x_f32)
         V = self.V(x_f32)
+
         g_q = (tf.nn.tanh(q) + 1.0) / 2.0
         g_k = (tf.nn.tanh(k) + 1.0) / 2.0
-        score = g_q * g_k
 
-        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
-        
-        out = self.norm(x_comb + residual)
+        # 과거 토큰 가중치 반영 점수
+        score = g_q * g_k * self.alpha  # element-wise scaling
+        # 누적합 대신 가중 평균
+        # score_t = sum_{i=0}^{t} alpha_i * V_i / sum_{i=0}^{t} alpha_i
+        score_cum = tf.math.cumsum(score * V, axis=1)
+        alpha_cum = tf.math.cumsum(score, axis=1)
+        score_weighted = score_cum / tf.maximum(alpha_cum, self.eps)
+
+        # 정규화 + 클리핑
+        score_norm = tf.clip_by_value(score_weighted, -self.clip_value, self.clip_value)
+        out = self.norm(score_norm + residual)
         out = self.glu(out)
         return tf.cast(out, x.dtype)
 
-
 class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()