Update Mo.py

Mo.py

@@ -124,51 +124,50 @@ class SwiGLU(layers.Layer):
         x_val, x_gate = tf.split(x_proj, 2, axis=-1)
         return self.out(x_val * tf.nn.silu(x_gate))
         
-class LoU(layers.Layer):
+class LoUScan(layers.Layer):
     def __init__(self, d_model, clip_value=5.0, eps=1e-6):
         super().__init__()
         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, 320)
-
-        # 학습 가능한 과거 토큰 가중치
-        self.alpha = self.add_weight(shape=(d_model,), initializer='ones', trainable=True)
+        self.glu = SwiGLU(d_model, 3500)
 
     def call(self, x):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
-        x_f32 = self.norm1(x)
+        x_f32 = self.norm1(x_f32)
 
         q = self.Q(x_f32)
         k = self.K(x_f32)
-        V = self.V(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 * 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)
+        score = g_q * g_k  # gating
+
+        # tf.scan으로 순차 누적합 (인과적)
+        def step(carry, inputs):
+            prev_sum = carry
+            s, v_t = inputs
+            new_sum = prev_sum + s * v_t
+            # 정규화
+            out = new_sum / tf.maximum(tf.reduce_sum(score[:tf.shape(prev_sum)[0]], axis=0, keepdims=True), self.eps)
+            return new_sum, out
+
+        # 초기값
+        init = tf.zeros_like(v[0])
+        _, outputs = tf.scan(step, (score, v), initializer=init, axis=0)
+
+        # 안정화
+        outputs = tf.clip_by_value(outputs, -self.clip_value, self.clip_value)
+        out = self.norm(outputs + residual)
         out = self.glu(out)
         return tf.cast(out, x.dtype)
 
@@ -187,7 +186,7 @@ class Lo(layers.Layer):
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.lou = LoU(d_model)
+        self.lou = LoUScan(d_model)
         self.lo = Lo(d_model)
 
     def call(self, x):