Update Model.py

Model.py

@@ -69,7 +69,7 @@ vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
 max_len = 512
-batch_size = 256
+batch_size = 32
 
 def text_to_ids(text):
     return sp.encode(text, out_type=int)
@@ -124,71 +124,58 @@ class SwiGLU(layers.Layer):
         out = self.W1(tf.nn.silu(a) * b)     
         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.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
-        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):
-        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
-            new = alpha_t * x_t + (1.0 - alpha_t) * prev_ema
-            return new
-
-        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)
-        ema = tf.transpose(ema_seq, perm=[1, 0, 2])
-        return ema
+class SparseCausalAttention(Layer):
+    def __init__(self, num_heads, head_dim, window_size=16, **kwargs):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.window_size = window_size  # 로컬 윈도우 크기
+
+    def build(self, input_shape):
+        self.q_dense = Dense(self.num_heads * self.head_dim)
+        self.k_dense = Dense(self.num_heads * self.head_dim)
+        self.v_dense = Dense(self.num_heads * self.head_dim)
+        self.out_dense = Dense(input_shape[-1])
 
     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.sigmoid(q)
-        # g_k = tf.nn.sigmoid(k)
-
-        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
-
-        alpha_dynamic = self.alpha_linear(x_f32)
-        score_ema = self._ema_over_time(score, alpha_dynamic)
-        mean_last = tf.reduce_mean(score_ema, axis=-1, keepdims=True)
-        denom = tf.maximum(mean_last, self.eps)
-        score_norm = score_ema / denom
-        score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
-        x_comb = score_clipped * V
-        out = self.proj(x_comb)
-        out = self.norm(out + residual)
-        return tf.cast(out, x.dtype)
+        batch_size, seq_len, dim = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2]
+        
+        # Q, K, V
+        q = tf.reshape(self.q_dense(x), (batch_size, seq_len, self.num_heads, self.head_dim))
+        k = tf.reshape(self.k_dense(x), (batch_size, seq_len, self.num_heads, self.head_dim))
+        v = tf.reshape(self.v_dense(x), (batch_size, seq_len, self.num_heads, self.head_dim))
+
+        # Transpose for matmul: (batch, heads, seq, head_dim)
+        q = tf.transpose(q, perm=[0, 2, 1, 3])
+        k = tf.transpose(k, perm=[0, 2, 1, 3])
+        v = tf.transpose(v, perm=[0, 2, 1, 3])
+
+        # 스케일
+        scale = tf.math.sqrt(tf.cast(self.head_dim, tf.float32))
+        q = q / scale
+
+        # 희소 마스크 계산: 로컬 윈도우
+        # 각 토큰 i는 max(i-window_size,0) ~ i까지 attention
+        attn_scores = tf.matmul(q, k, transpose_b=True)  # (batch, heads, seq, seq)
+        mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)  # causal mask
+        # 윈도우 크기 제한
+        band_mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), self.window_size, 0)
+        mask = mask * band_mask
+        mask = tf.reshape(mask, (1, 1, seq_len, seq_len))  # 브로드캐스트 가능
+        attn_scores = tf.where(mask > 0, attn_scores, tf.fill(tf.shape(attn_scores), -1e9))
+
+        attn_probs = tf.nn.softmax(attn_scores, axis=-1)
+        attn_output = tf.matmul(attn_probs, v)  # (batch, heads, seq, head_dim)
+
+        # 합치기
+        attn_output = tf.transpose(attn_output, perm=[0, 2, 1, 3])
+        attn_output = tf.reshape(attn_output, (batch_size, seq_len, self.num_heads*self.head_dim))
+        return self.out_dense(attn_output)
 
 class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.d = layers.Dense(256, activation='silu')
+        self.d = layers.Dense(64, activation='silu')
         self.w = layers.Dense(d_model)
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
 
@@ -200,7 +187,7 @@ class Lo(layers.Layer):
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.lou = LoU(d_model)
+        self.lou = SparseCausalAttention(num_heads=2, head_dim=64)
         self.glu = SwiGLU(d_model)
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.lo = Lo(d_model)
@@ -256,8 +243,8 @@ 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=700,
-    n_layers=16
+    d_model=128,
+    n_layers=2
 )
 
 # 옵티마이저 설정