Update Model.py

Model.py

@@ -116,37 +116,35 @@ with strategy.scope():
 class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        # 내부 계산은 float32로 유지
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
         self.p = layers.Dense(128, use_bias=True, dtype='float32')
         self._out_dtype = 'float32'
 
     def call(self, x):
-        # x may be bfloat16; cast to float32 for stable intermediate computation
         x_f32 = tf.cast(x, tf.float32)
         x = self.proj(x_f32)
         x = tf.nn.gelu(x)
         x = self.p(x)
-        # cast back to model dtype for consistency
         return tf.cast(x, self._out_dtype)
 
+class SwiGLU(layers.Layer):
+    def __init__(self, d_model):
+        super().__init__()
+        self.W = layers.Dense(3500, dtype='float32')
+        self.W1 = layers.Dense(d_model, dtype='float32')
+    def call(self, x):
+        x = tf.cast(x, tf.float32)
+        x = self.W(x)
+        a, b = tf.split(x, 2, axis=-1)
+        out = self.W1(tf.nn.silu(a) * b)     
+        return tf.cast(out, x.dtype)
+
 class LoU(layers.Layer):
-    """
-    안정화된 LoSoU 레이어 (동적 alpha 사용)
-    - alpha 값을 입력에 따라 동적으로 계산: alpha = sigmoid(Linear(x))
-    - 누적합 대신 지수이동평균(EMA) 사용 (alpha: smoothing factor)
-    - 내부 계산은 float32로 수행 (TPU bfloat16 안정성 향상)
-    - EMA 결과 클리핑 및 작은 epsilon 적용
-    - 안전한 split 처리 (짝수 차원 가정; 아니라면 마지막 차원 pad 필요)
-    """
     def __init__(self, d_model, clip_value=5.0, eps=1e-6):
         super().__init__()
-        # 대부분 연산을 float32로 수행
         self.d_model = d_model
         self.clip_value = float(clip_value)
         self.eps = float(eps)
-
-        # projection / gating layers in 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')
@@ -155,100 +153,55 @@ class LoU(layers.Layer):
         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.glu = SwiGLU(d_model)
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
-        
         self.alpha_linear = layers.Dense(1, activation='sigmoid', dtype='float32')
 
     def _ema_over_time(self, score, alpha_dynamic):
-        # score: (B, L, D) float32 in [0,1] roughly
-        # alpha_dynamic: (B, L, 1) float32 in [0,1]
-
-        # transpose to (L, B, D) to scan over time steps
-        seq = tf.transpose(score, perm=[1, 0, 2])  # (L, B, D)
-        alpha_seq = tf.transpose(alpha_dynamic, perm=[1, 0, 2])  # (L, B, 1)
+        seq = tf.transpose(score, perm=[1, 0, 2])
+        alpha_seq = tf.transpose(alpha_dynamic, perm=[1, 0, 2])
 
         def step(prev_ema, inputs):
             x_t, alpha_t = inputs
-            # prev_ema: (B, D), x_t: (B, D), alpha_t: (B, 1)
             new = alpha_t * x_t + (1.0 - alpha_t) * prev_ema
             return new
 
-        # 초기값을 첫 step 값으로 설정
-        init = seq[0]  # (B, D)
-        first_alpha = alpha_seq[0]  # (B, 1)
-
-        # scan의 elems는 (L-1, B, D) 및 (L-1, B, 1) 이어야 함
-        remaining_seq = seq[1:]  # (L-1, B, D)
-        remaining_alpha = alpha_seq[1:]  # (L-1, B, 1)
-
-        # elems는 두 텐서의 튜플로 구성: (x_t, alpha_t)
+        init = seq[0]
+        first_alpha = alpha_seq[0]
+        remaining_seq = seq[1:]
+        remaining_alpha = alpha_seq[1:]
         elems = (remaining_seq, remaining_alpha)
-
         ema_seq = tf.scan(fn=step, elems=elems, initializer=init)
-        # 초기값 포함
-        ema_seq = tf.concat([tf.expand_dims(init, 0), ema_seq], axis=0)  # (L, B, D)
-
-        # transpose back to (B, L, D)
+        ema_seq = tf.concat([tf.expand_dims(init, 0), ema_seq], axis=0)
         ema = tf.transpose(ema_seq, perm=[1, 0, 2])
         return ema
 
     def call(self, x):
-        # x: (B, L, d_model) maybe bfloat16 or float32
-        # cast to float32 for all internal computations
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
         x_f32 = self.norm1(x)
 
-        # Q, K, V
         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)
         g_k = tf.nn.sigmoid(k)
-
-        # elementwise product -> bounded roughly [0,1]
         score = g_q * g_k
-
-        # 동적 alpha 계산: (B, L, d_model) -> (B, L, 1)
-        alpha_dynamic = self.alpha_linear(x_f32) # (B, L, 1)
-        # 필요시 alpha_dynamic에 대한 후처리 (예: min/max 등) 가능
-        # ex: alpha_dynamic = tf.clip_by_value(alpha_dynamic, 0.01, 0.99)
-
-        # EMA across time (stable alternative to cumsum)
+        alpha_dynamic = self.alpha_linear(x_f32)
         score_ema = self._ema_over_time(score, alpha_dynamic)
-
-        # optionally normalize by (mean + eps) across last dim to reduce scale variations
-        mean_last = tf.reduce_mean(score_ema, axis=-1, keepdims=True)  # (B, L, 1)
+        mean_last = tf.reduce_mean(score_ema, axis=-1, keepdims=True)
         denom = tf.maximum(mean_last, self.eps)
         score_norm = score_ema / denom
-
-        # clip to avoid extremes
         score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
-
-        # combine with V
-        x_comb = score_clipped * V  # (B, L, d_model)
-
-        out = self.proj(x_comb)  # (B, L, d_model)
-
-        # ensure out dim even for split
-        d = out.shape[-1]  # this is an int (static shape)
-        if d is not None and d % 2 == 1:
-            out = tf.pad(out, [[0,0],[0,0],[0,1]])
-
-        a, b = tf.split(out, 2, axis=-1)
-        gated = tf.nn.silu(a) * b
-        out = self.O(gated)
-
+        x_comb = score_clipped * V
+        out = self.proj(x_comb)
+        out = self.glu(out)
         out = self.norm(out + residual)
-
-        # cast back to original dtype for downstream layers
         return tf.cast(out, x.dtype)
 
 class ReLM(tf.keras.Model):
@@ -257,20 +210,15 @@ class ReLM(tf.keras.Model):
         self.token_embedding = layers.Embedding(vocab_size, d_model)
         self.pos_embedding = layers.Embedding(max_seq_len, d_model)
         self.blocks = [LoU(d_model) for _ in range(n_layers)]
-
-        # LayerNormalization은 float32로 해서 정밀도 문제 방지
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
 
     def call(self, x, training=False):
         batch_size, seq_len = tf.shape(x)[0], tf.shape(x)[1]
         positions = tf.range(seq_len)[tf.newaxis, :]
-
         x = self.token_embedding(x) + self.pos_embedding(positions)
         for block in self.blocks:
             x = block(x)
-
         x = self.ln_f(x)
-
         embedding_matrix = tf.cast(self.token_embedding.embeddings, x.dtype)
         logits = tf.matmul(x, embedding_matrix, transpose_b=True)
         return tf.cast(logits, tf.float32)
@@ -301,7 +249,7 @@ def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
 model = ReLM(
     vocab_size=vocab_size,
     max_seq_len=max_len,
-    d_model=128,
+    d_model=700,
     n_layers=2
 )