KenT: GeLU MLP, 12 layers, N=24 D=208 (~301M)

README.md

@@ -13,7 +13,8 @@ tags:
   - linear-time
   - frequency-domain
   - dct
-  - swiglu
+  - mlp
+  - gelu
   - text-generation
   - custom_code
   - fp16
@@ -22,19 +23,19 @@ tags:
 
 # KenT
 
-KenT is a recurrent language model from the SSM family. KenT concentrates its capacity in a very wide spectral representation: few layers, each with a single transformation block, and a light learned-forgetting memory.
+KenT is a recurrent language model from the SSM family. KenT concentrates its capacity in a wide spectral representation stacked over a moderate depth: each layer has a single transformation block, plus a light learned-forgetting memory.
 
 > This repo contains `trust_remote_code` code; weights are stored in **FP16**.
 
 ## Architectural Philosophy
 
-- **Shallow stack, wide spectrum.** KenT stacks a few identical recurrent layers (8 in this ~300M model), each containing exactly
-  one transformation block. The layers are connected in sequence through the representation stream, and the last layer produces
+- **Wide spectrum, moderate depth.** KenT stacks identical recurrent layers (12 in this ~300M model), each containing exactly
+  one transformation block (a GeLU MLP with hidden size 4 x D). The layers are connected in sequence through the representation stream, and the last layer produces
   the final output. The representation stream feeding `lm_head` is wide (1024), and each layer's Bilinear projection expands it
-  into a large resolution axis (N = 32 DCT positions) with enough features per position (D = 240) before the DCT.
+  into a large resolution axis (N = 24 DCT positions) with enough features per position (D = 208) before the DCT.
 - **Invest in the transformation core.** Capacity is concentrated in the "hidden" part, where features are actually transformed,
-  rather than in the vocabulary table. Parameters were scaled from ~100M to ~300M by widening the representation stream and the
-  spectral axes (N and D together), not by stacking more layers.
+  rather than in the vocabulary table. Parameters were scaled from ~100M by widening the representation stream and the spectral
+  axes (N and D together) and then by adding layers; the MLP inside each block is kept small.
 - **Two streams, two roles.** One stream is a representation carrying information about the current token; the other is a memory
   stream that aggregates history. The two streams meet only through learned projections and are never mixed directly.
 - **Frequency as a fixed foundation, learning where it's needed.** Features are processed in the cosine-coefficient space using a
@@ -43,8 +44,8 @@ KenT is a recurrent language model from the SSM family. KenT concentrates its ca
 - **Memory with learned forgetting.** History is summarized with an exponential moving average
   `H_t = (1 - a) H_{t-1} + a u_t`, where `a` is a single learned scalar per layer. Its initial value decreases with depth
   (shallow layers forget fast, deep layers remember longer). The decoding state has **constant** size regardless of context length.
-- **Width is the main axis.** There is no inner chain of blocks inside a layer: each layer applies one DCT/SwiGLU block over a
-  large N x D grid, and memory timescale still varies across the 8 layers.
+- **Width and depth together.** There is no inner chain of blocks inside a layer: each layer applies one DCT/MLP block over a
+  large N x D grid, and memory timescale varies across the 12 layers.
 - **Chunk-parallelizable.** The moving average is a linear recurrence, so a long document can be split into blocks/chunks,
   processed concurrently, and merged with an affine scan, yielding the same result as sequential processing.
 

config.json

@@ -10,10 +10,10 @@
   },
   "vocab_size": 50257,
   "hidden_size": 1024,
-  "num_hidden_layers": 8,
-  "resolution_size": 32,
-  "feature_dim": 240,
-  "intermediate_size": 960,
+  "num_hidden_layers": 12,
+  "resolution_size": 24,
+  "feature_dim": 208,
+  "intermediate_size": 832,
   "rms_norm_eps": 1e-05,
   "use_bias": true,
   "initializer_range": 0.02,

model.py

@@ -5,7 +5,7 @@ Mỗi KenT layer thực hiện một bước, với ĐÚNG MỘT WNN Block:
 
     V   = B_in(Norm_X(X))                      # (…, N, d_h)
     Z   = A_N V                                # DCT-II cố định trên trục N
-    u   = A_N [ Z + SwiGLU(Norm(Z)) ]          # một WNN Block
+    u   = A_N [ Z + MLP(Norm(Z)) ]             # một WNN Block, MLP = Linear -> GeLU -> Linear
     H_t = (1 - a) H_{t-1} + a u_t              # EMA, a = a_layer học được, riêng từng layer
     X'  = X + B_out(Norm_H(H_t))
 
@@ -47,9 +47,9 @@ class KentConfig(PretrainedConfig):
         num_hidden_layers: số KenT layer xếp chồng (mỗi layer có đúng một WNN Block).
         resolution_size: N, độ dài trục mà DCT-II tác động lên.
         feature_dim: d_h, số chiều đặc trưng trên mỗi vị trí resolution.
-        intermediate_size: kích thước hidden của SwiGLU (trên trục d_h).
+        intermediate_size: kích thước hidden của MLP GeLU (trên trục d_h).
         rms_norm_eps: epsilon của RMSNorm.
-        use_bias: dùng bias trong Bilinear / SwiGLU.
+        use_bias: dùng bias trong Bilinear / MLP.
         initializer_range: độ lệch chuẩn khởi tạo trọng số.
         ema_alpha_first: a_layer khởi tạo của layer đầu (quên nhanh).
         ema_alpha_last: a_layer khởi tạo của layer cuối (nhớ lâu); các layer ở giữa nội suy
@@ -265,29 +265,28 @@ class KentBilinear(nn.Module):
         return a * b
 
 
-class KentSwiGLU(nn.Module):
-    """SwiGLU(x) = (SiLU(x W_g + b_g) ⊙ (x W_v + b_v)) W_o + b_o."""
+class KentMLP(nn.Module):
+    """MLP(x) = GeLU(x W_1 + b_1) W_2 + b_2."""
 
     def __init__(self, dim: int, hidden: int, bias: bool = True):
         super().__init__()
-        self.gate_value = nn.Linear(dim, 2 * hidden, bias=bias)
-        self.out = nn.Linear(hidden, dim, bias=bias)
+        self.fc_in = nn.Linear(dim, hidden, bias=bias)
+        self.fc_out = nn.Linear(hidden, dim, bias=bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        g, v = self.gate_value(x).chunk(2, dim=-1)
-        return self.out(F.silu(g) * v)
+        return self.fc_out(F.gelu(self.fc_in(x)))
 
 
 class KentWNNBlock(nn.Module):
-    """Z' = A_N [ Z + SwiGLU(Norm(Z)) ]  (residual trước DCT-II)."""
+    """Z' = A_N [ Z + MLP(Norm(Z)) ]  (residual trước DCT-II)."""
 
     def __init__(self, config: KentConfig):
         super().__init__()
         self.norm = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
-        self.swiglu = KentSwiGLU(config.feature_dim, config.intermediate_size, config.use_bias)
+        self.mlp = KentMLP(config.feature_dim, config.intermediate_size, config.use_bias)
 
     def forward(self, z: torch.Tensor) -> torch.Tensor:
-        return dct2(z + self.swiglu(self.norm(z)))
+        return dct2(z + self.mlp(self.norm(z)))
 
 
 class KentLayer(nn.Module):