Update Mo.py

Mo.py

@@ -124,57 +124,53 @@ 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 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)
-
-        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)  # 사용자 정의 GLU
-
-    def call(self, x):
-        x_f32 = tf.cast(x, tf.float32)
-        residual = x_f32
-        x_f32 = self.norm1(x_f32)
-
-        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  # element-wise gating
-
-        # 배치별 순차적 scan 적용 (인과적)
-        def process_sequence(inputs):
-            score_seq, v_seq = inputs
-            seq_len = tf.shape(v_seq)[0]
-            init = tf.zeros_like(v_seq[0])
-
-            def step(carry, elems):
-                s_t, v_t = elems
-                new_sum = carry + s_t * v_t  # 현재까지 누적
-                out = new_sum / tf.maximum(tf.reduce_sum(score_seq[:tf.shape(v_seq)[0]], axis=0, keepdims=True), self.eps)
-                return new_sum, out
-
-            _, outputs = tf.scan(step, (score_seq, v_seq), initializer=init)
-            return outputs
-
-        # 배치 차원 처리
-        outputs = tf.map_fn(lambda inp: process_sequence(inp), (score, v), dtype=tf.float32)
-
-        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)
-
+class LoU(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)
+        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)
+    def call(self, x):
+        x_f32 = tf.cast(x, tf.float32)
+        residual = x_f32
+        x_f32 = self.norm1(x)
+
+        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)
+        out = self.glu(out)
+        return tf.cast(out, x.dtype)
 
 class Lo(layers.Layer):
     def __init__(self, d_model):
@@ -191,7 +187,7 @@ class Lo(layers.Layer):
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.lou = LoUScan(d_model)
+        self.lou = LoU(d_model)
         self.lo = Lo(d_model)
 
     def call(self, x):