Update V2.py

V2.py

@@ -129,53 +129,48 @@ ds = ds.map(lambda v1, v2: ((v1, v2), tf.zeros([BATCH_SIZE], dtype=tf.float32)),
 ds = ds.prefetch(tf.data.AUTOTUNE)
 
 
-class DynamicConv(layers.Layer):
-    def __init__(self, d_model, k=7):
+
+class MixerBlock(layers.Layer):
+    def __init__(self, seq_len, dim, token_mlp_dim, channel_mlp_dim, dropout=0.0):
         super().__init__()
-        assert k % 2 == 1
-        self.k = k
-        self.dense = layers.Dense(d_model, activation='silu')
-        self.proj = layers.Dense(d_model)
-        self.generator = layers.Dense(k, dtype='float32')
-
-        self.ln1 = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
-        self.ln2 = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
-        
-        
-    def call(self, x):
-        x_in = x
-        x = tf.cast(x, tf.float32)
-        x = self.ln1(x)
+        self.seq_len = seq_len
+        self.dim = dim
+        self.token_mlp_dim = token_mlp_dim
+        self.channel_mlp_dim = channel_mlp_dim
 
-        B = tf.shape(x)[0]
-        L = tf.shape(x)[1]
-        D = tf.shape(x)[2]
+        self.ln1 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
+        # token-mixing MLP: operate over tokens => apply Dense on transposed axis
+        self.token_fc1 = layers.Dense(token_mlp_dim, activation='gelu', dtype=tf.float32)
+        self.token_fc2 = layers.Dense(seq_len, dtype=tf.float32)
 
-        kernels = self.generator(self.dense(x))
-        kernels = tf.nn.softmax(kernels, axis=-1)
+        self.ln2 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
+        # channel-mixing MLP: operate per-token over channels
+        self.channel_fc1 = layers.Dense(channel_mlp_dim, activation='gelu', dtype=tf.float32)
+        self.channel_fc2 = layers.Dense(dim, dtype=tf.float32)
 
-        pad = (self.k - 1) // 2
-        x_pad = tf.pad(x, [[0,0],[pad,pad],[0,0]])
+        self.dropout = layers.Dropout(dropout)
 
-        x_pad_4d = tf.expand_dims(x_pad, axis=1)
-        patches = tf.image.extract_patches(
-            images=x_pad_4d,
-            sizes=[1,1,self.k,1],
-            strides=[1,1,1,1],
-            rates=[1,1,1,1],
-            padding='VALID'
-        )
-        patches = tf.reshape(patches, [B, L, self.k, D])
+    def call(self, x, training=None):
+        # x: (B, L, D)
+        B = tf.shape(x)[0]
+        L = tf.shape(x)[1]
+        D = tf.shape(x)[2]
 
-        kernels_exp = tf.expand_dims(kernels, axis=-1)
-        out = tf.reduce_sum(patches * kernels_exp, axis=2)
-        out = self.proj(out)
-        out = tf.nn.gelu(out)
-        out = x + self.ln2(out)
+        # Token-mixing
+        y = self.ln1(x)                          # (B, L, D)
+        y_t = tf.transpose(y, perm=[0,2,1])     # (B, D, L)
+        y_t = self.token_fc1(y_t)               # (B, D, token_mlp_dim)
+        y_t = self.token_fc2(y_t)               # (B, D, L)
+        y = tf.transpose(y_t, perm=[0,2,1])     # (B, L, D)
+        x = x + self.dropout(y, training=training)
 
-        # 🔥 원래 dtype으로 돌려줌
-        return tf.cast(out, x_in.dtype)
+        # Channel-mixing
+        z = self.ln2(x)
+        z = self.channel_fc1(z)
+        z = self.channel_fc2(z)
+        x = x + self.dropout(z, training=training)
 
+        return x
 
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):
@@ -192,7 +187,7 @@ class SentenceEncoder(Model):
         self.embed = layers.Embedding(vocab_size, embed_dim)
         self.pos_embed = layers.Embedding(input_dim=max_len, output_dim=embed_dim)
         self.dropout = layers.Dropout(dropout_rate)
-        self.blocks = [DynamicConv(d_model=embed_dim, k=7) for _ in range(4)]
+        self.blocks = [MixerBlock(seq_len=MAX_LEN, dim=embed_dim, token_mlp_dim=256, channel_mlp_dim=embed_dim, dropout=0.1) for _ in range(3)]
         self.attn_pool = layers.Dense(1)
         
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)