Update Model.py

Model.py

@@ -145,7 +145,7 @@ class Lo(layers.Layer):
         super().__init__()
         # 내부 계산은 float32로 유지
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
-        self.p = layers.Dense(96, use_bias=True, dtype='float32')
+        self.p = layers.Dense(128, use_bias=True, dtype='float32')
         self._out_dtype = 'float32'
 
     def call(self, x):
@@ -174,21 +174,15 @@ class LoSoU(layers.Layer):
         self.eps = float(eps)
 
         # projection / gating layers in float32
-        self.Q = layers.Dense(96, dtype='float32')
-        self.K = layers.Dense(96, dtype='float32')
-        self.V = Lo(d_model)  # Lo already handles casting to model dtype; we'll cast back to float32
+        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.Qr = Lo(d_model)
+        self.Kr = Lo(d_model)
+        self.Vr = Lo(d_model)
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
         self.O = layers.Dense(d_model, dtype='float32')
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
-
-        # 동적 alpha 계산을 위한 레이어
-        # alpha는 [0, 1] 범위여야 하므로 sigmoid 사용
-        # 입력 x의 d_model 차원을 사용하여 각 샘플에 대해 alpha 계산
-        # 예: (B, L, d_model) -> (B, L, 1) -> (B, L, 1) with sigmoid
-        # 또는 (B, L, d_model) -> (B, L, d_model) -> global reduce -> (B, L, 1)
-        # 간단히 각 위치에 대해 동일한 alpha 사용 (입력의 평균 기반)
-        # 또는 위치별로 다르게 사용 (각 위치에 대해 계산)
-        # 여기서는 위치별로 다르게 계산 (B, L, 1)
         self.alpha_linear = layers.Dense(1, activation='sigmoid', dtype='float32')
 
     def _ema_over_time(self, score, alpha_dynamic):
@@ -231,9 +225,13 @@ class LoSoU(layers.Layer):
         residual = x_f32
 
         # Q, K, V
-        q = self.Q(x_f32)   # (B, L, 96)
-        k = self.K(x_f32)   # (B, L, 96)
-        V = tf.cast(self.V(x), tf.float32)  # ensure V's output is float32
+        q = self.Q(x_f32)
+        k = self.K(x_f32)
+        V = self.V(x_f32)
+
+        q = self.Qr(q)
+        k = self.Kr(k)
+        v = self.Vr(v)
 
         # gating signals in (0,1)
         g_q = tf.nn.sigmoid(q)