fix architecture

README.md

@@ -22,14 +22,14 @@ tags:
 
 # KenT
 
-KenT is a recurrent language model from the SSM family. Rather than stacking more layers, KenT concentrates its capacity in the **depth of the transformation inside a single cell**, and keeps temporal memory as simple as possible.
+KenT is a recurrent language model from the SSM family. KenT concentrates its capacity in compact cells built from a short chain of transformation blocks, and keeps temporal memory as simple as possible.
 
 > This repo contains `trust_remote_code` code; weights are stored in **FP16**.
 
 ## Architectural Philosophy
 
-- **One cell, no layer stacking.** KenT is a single recurrent cell; the model is not replicated into multiple layers.
-  Depth comes not from the number of layers but from a chain of transformation blocks applied in sequence inside the cell.
+- **A short chain of compact cells.** KenT stacks a few identical recurrent cells, each containing only two transformation blocks.
+  The cells are connected in sequence through the representation stream, and the last cell produces the final output.
 - **Invest in the transformation core.** Capacity is concentrated in the "hidden" part, where features are actually transformed,
   rather than in the vocabulary table or in repeating the architecture. The width and depth of this core are scaled together to
   stay balanced, rather than pushing a single dimension.

config.json

@@ -10,10 +10,10 @@
   },
   "vocab_size": 50257,
   "hidden_size": 512,
-  "num_hidden_layers": 1,
-  "resolution_size": 32,
+  "num_hidden_layers": 5,
+  "resolution_size": 8,
   "feature_dim": 512,
-  "num_wnn_blocks": 16,
+  "num_wnn_blocks": 2,
   "intermediate_size": 2048,
   "rms_norm_eps": 1e-05,
   "use_bias": true,

model.py

@@ -39,10 +39,10 @@ class KentConfig(PretrainedConfig):
     Args:
         vocab_size: kích thước từ vựng.
         hidden_size: số chiều của representation stream X_t (d_x).
-        num_hidden_layers: số KenT layer xếp chồng.
+        num_hidden_layers: số KenT cell nối liền nhau (cell cuối sinh đầu ra cuối).
         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.
-        num_wnn_blocks: L, số WNN Block nối tiếp trong một timestep.
+        num_wnn_blocks: L, số WNN Block nối tiếp trong MỖI cell.
         intermediate_size: kích thước hidden của SwiGLU (trên trục d_h).
         rms_norm_eps: epsilon của RMSNorm.
         use_bias: dùng bias trong Bilinear / SwiGLU.
@@ -57,10 +57,10 @@ class KentConfig(PretrainedConfig):
         self,
         vocab_size: int = 50257,
         hidden_size: int = 512,
-        num_hidden_layers: int = 1,
-        resolution_size: int = 32,
+        num_hidden_layers: int = 5,
+        resolution_size: int = 8,
         feature_dim: int = 512,
-        num_wnn_blocks: int = 16,
+        num_wnn_blocks: int = 2,
         intermediate_size: int = 2048,
         rms_norm_eps: float = 1e-5,
         use_bias: bool = True,