Upload 3 files

README.md

@@ -14,26 +14,24 @@ tags:
 - custom_code
 - fp16
 - base
-datasets:
-- wikimedia/wikipedia
-- HuggingFaceFW/fineweb-edu
 ---
 
 # KenT-30M-Base
 
-KenT is a recurrent language model. This ~30M-parameter version stacks 24 identical layers; each layer expands the
-representation stream into `M` resolution positions, mixes them with a fixed DCT-II, transforms the features with a SwiGLU,
-and keeps a learned-forgetting memory (EMA) with one decay coefficient per resolution position.
+KenT is a recurrent language model. This ~30M-parameter version reads **one token per step** and stacks 24 identical layers;
+each layer expands the token vector into `M` resolution positions, mixes them with a fixed DCT-II, transforms the features
+with a SwiGLU, and keeps a learned-forgetting memory (EMA) whose decay is a learned **matrix** (one coefficient per
+resolution position and feature). The memory is read back through a DCT-II and projected to the token dimension.
 
 > This repo contains `trust_remote_code` code; weights are stored in **FP16**.
 
 ## Layer
 
 ```
-Z   = ResProj_in(Norm(X))              # n -> M x d
+Z   = ResProj_in(Norm(X))              # 1 -> M: token vector n -> M x d
 Z   = Z + SwiGLU(DCT_M(RMSNorm(Z)))    # WNN block: fixed DCT-II over the M axis, SwiGLU over the d axis
-H_t = (1 - a) * H_{t-1} + a * Z_t      # EMA, a = learned vector of M coefficients per layer
-X'  = X + ResProj_out(Norm(H_t))       # M x d -> n
+H_t = (1 - A) * H_{t-1} + A * Z_t      # EMA Learner Matrix: A = sigmoid(logits), shape M x d, per layer
+X'  = X + ResProj_out(DCT_M(Norm(H_t)))  # DCT-II Proj, then M -> 1; residual with the layer input
 ```
 
 | | value |
@@ -42,7 +40,9 @@ X'  = X + ResProj_out(Norm(H_t))       # M x d -> n
 | layers | 24 |
 | resolution `M` x features `d` | 12 x 96 (M·d = 1152 = 4.5n) |
 | SwiGLU hidden | 384 (4d) |
-| parameters | ~29.7M (embedding 12.9M + layers 16.9M) |
+| EMA decay | learned matrix `(12, 96)` per layer (`ema_per_feature=false` -> vector `(12,)`) |
+| recurrent state | `H`: `(12, 96)` per layer, constant size |
+| parameters | ~29.75M (embedding 12.9M + layers 16.9M) |
 
 ## Quickstart
 
@@ -66,4 +66,4 @@ Long documents: `python infer.py --model ./KenT-30M-Base --prompt-file doc.txt -
 ## Notes
 
 - Research model, trained from scratch; evaluate before real-world use.
-- The tokenizer and weights must match `vocab_size` in `config.json`.
+- The tokenizer and weights must match `vocab_size` in `config.json`.

config.json

@@ -26,6 +26,7 @@
   "dtype": "float16",
   "ema_alpha_first": 0.5,
   "ema_alpha_last": 0.01,
+  "ema_per_feature": true,
   "ema_block_size": 128,
   "use_gradient_checkpointing": false
 }

model.py

@@ -1,24 +1,25 @@
 # coding=utf-8
 """KenT: kiến trúc hồi quy hai luồng (representation X_t / temporal-memory H_t).
 
-Mỗi KenT layer (hộp "WNN" trong sơ đồ) thực hiện một bước:
+Mỗi bước nhận ĐÚNG MỘT token (cửa sổ trượt từng token một); lịch sử nằm trong trạng thái H của từng layer.
+Mỗi KenT layer (hộp "WNN" trong sơ đồ) thực hiện:
 
-    Z    = ResProj_in(Norm_X(X))                # (…, M, d)   n -> M resolution, mỗi resolution d feature
+    Z    = ResProj_in(Norm_X(X))                # 1 -> M: một vector token thành M resolution, mỗi resolution d feature
     Z    = Z + SwiGLU( DCT_M( RMSNorm(Z) ) )    # WNN Block: DCT-II cố định trên trục M, SwiGLU trên trục d, skip quanh cả cụm
-    H_t  = (1 - a) ⊙ H_{t-1} + a ⊙ Z_t          # EMA, a = sigmoid(ema_logit) là VECTOR M phần tử học được / layer
-    X'   = X + ResProj_out(Norm_H(H_t))         # M x d -> n  (Add với X vào layer, rồi sang layer kế / lm_head)
+    H_t  = (1 - A) ⊙ H_{t-1} + A ⊙ Z_t          # "EMA Learner Matrix": A = sigmoid(ema_logit), MA TRẬN M x d học được / layer
+    X'   = X + ResProj_out( DCT_M( Norm_H(H_t) ) )   # DCT-II Proj rồi M -> 1; Add với X vào layer (residual)
 
 Mỗi layer xếp K = `num_wnn_blocks` WNN Block nối tiếp (mặc định K = 1, đúng sơ đồ); độ sâu đến từ việc xếp chồng
-nhiều layer ("xN" trong sơ đồ).
+nhiều layer ("xN" trong sơ đồ). Layer cuối cho X' -> norm cuối -> lm_head (tied với embedding).
 
-Bộ nhớ H có cơ chế quên: a là một vector M phần tử (mỗi vị trí trên trục resolution/phổ có hệ số quên riêng,
-chung cho mọi feature d), học được và riêng từng layer. Khởi tạo: nội suy theo cấp số nhân từ `ema_alpha_first`
-(layer đầu, vị trí đầu) tới `ema_alpha_last` (layer cuối, vị trí cuối) theo trung bình của chỉ số layer và chỉ số vị trí.
+Bộ nhớ H có cơ chế quên: A là ma trận M x d (mỗi vị trí resolution VÀ mỗi feature có hệ số quên riêng), học được và
+riêng từng layer. Khởi tạo: nội suy theo cấp số nhân từ `ema_alpha_first` (layer đầu, vị trí đầu, feature đầu) tới
+`ema_alpha_last` (layer cuối, vị trí cuối, feature cuối) theo trung bình của chỉ số layer, vị trí và feature.
+`ema_per_feature=False` quay về vector M phần tử (chung cho mọi feature), dùng scan theo khối.
 
-EMA là một phép hồi quy tuyến tính nên vẫn song song hóa được theo thời gian: trong từng khối
-độ dài `ema_block_size` dùng ma trận phân rã tam giác dưới (riêng cho từng vị trí M), các khối được ghép bằng
-affine-scan log-depth. Kết quả bằng đúng hồi quy tuần tự. Trạng thái giải mã chỉ là H cho mỗi
-layer: O(1) theo độ dài.
+EMA là một phép hồi quy tuyến tính (đường chéo theo từng kênh) nên vẫn song song hóa được theo thời gian bằng
+affine-scan log-depth (Hillis-Steele): kết quả bằng đúng hồi quy tuần tự. Trạng thái giải mã chỉ là H cho mỗi
+layer, shape (M, d): O(1) theo độ dài.
 
 Trọng số được lưu ở FP16; EMA và normalization chạy ở FP32.
 """
@@ -55,9 +56,11 @@ class KentConfig(PretrainedConfig):
         use_bias: dùng bias trong Resolution Proj (SwiGLU không bias).
         initializer_range: độ lệch chuẩn khởi tạo trọng số.
         ema_alpha_first: hệ số a khởi tạo của layer đầu, vị trí đầu (quên nhanh).
-        ema_alpha_last: hệ số a khởi tạo của layer cuối, vị trí cuối (nhớ lâu); phần còn lại nội suy theo cấp số
-            nhân theo trung bình chỉ số layer và chỉ số vị trí. Vector a vẫn là tham số học được.
-        ema_block_size: độ dài khối của EMA song song (chỉ ảnh hưởng tốc độ/bộ nhớ).
+        ema_alpha_last: hệ số a khởi tạo của layer cuối, vị trí cuối, feature cuối (nhớ lâu); phần còn lại nội suy theo
+            cấp số nhân theo trung bình các chỉ số (layer, vị trí[, feature]). a vẫn là tham số học được.
+        ema_per_feature: True = EMA Learner Matrix (hệ số quên riêng cho mỗi (vị trí M, feature d), ema_logit shape
+            (M, d)); False = vector M phần tử, chung cho mọi feature (ema_logit shape (M,)).
+        ema_block_size: độ dài khối của EMA song song (chỉ dùng khi ema_per_feature=False; ảnh hưởng tốc độ/bộ nhớ).
         use_cache: trả về KentCache khi suy luận.
         use_gradient_checkpointing: bật gradient checkpointing theo từng KenT layer khi khởi tạo model
             (chỉ có tác dụng ở chế độ train; tương đương gọi `model.gradient_checkpointing_enable()`).
@@ -79,6 +82,7 @@ class KentConfig(PretrainedConfig):
         initializer_range: float = 0.02,
         ema_alpha_first: float = 0.5,
         ema_alpha_last: float = 0.01,
+        ema_per_feature: bool = True,
         ema_block_size: int = 128,
         use_cache: bool = True,
         use_gradient_checkpointing: bool = False,
@@ -100,6 +104,7 @@ class KentConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.ema_alpha_first = ema_alpha_first
         self.ema_alpha_last = ema_alpha_last
+        self.ema_per_feature = ema_per_feature
         self.ema_block_size = ema_block_size
         self.use_cache = use_cache
         self.use_gradient_checkpointing = use_gradient_checkpointing
@@ -119,10 +124,22 @@ class KentConfig(PretrainedConfig):
         return [float(self.ema_alpha_first * ratio ** (0.5 * (r_l + (p / (n_p - 1) if n_p > 1 else 0.0))))
                 for p in range(n_p)]
 
+    def alpha_init_matrix(self, layer_idx: int):
+        """Ma trận a (M x d) khởi tạo của layer: giảm theo cấp số nhân theo độ sâu, vị trí và feature."""
+        n_l, n_p, n_d = self.num_hidden_layers, self.resolution_size, self.feature_dim
+        r_l = layer_idx / (n_l - 1) if n_l > 1 else 0.0
+        ratio = self.ema_alpha_last / self.ema_alpha_first
+        return [[float(self.ema_alpha_first * ratio ** ((r_l + (p / (n_p - 1) if n_p > 1 else 0.0)
+                                                         + (j / (n_d - 1) if n_d > 1 else 0.0)) / 3.0))
+                 for j in range(n_d)] for p in range(n_p)]
+
     def alpha_init(self, layer_idx: int) -> float:
-        """Trung bình của vector a khởi tạo của layer (để log)."""
-        v = self.alpha_init_vector(layer_idx)
-        return float(sum(v) / len(v))
+        """Trung bình của a khởi tạo của layer (để log)."""
+        if self.ema_per_feature:
+            flat = [x for row in self.alpha_init_matrix(layer_idx) for x in row]
+        else:
+            flat = self.alpha_init_vector(layer_idx)
+        return float(sum(flat) / len(flat))
 
 
 # --------------------------------------------------------------------------- #
@@ -164,11 +181,37 @@ def _affine_scan(a: torch.Tensor, b: torch.Tensor):
 
 
 def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
-    """EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16)."""
+    """EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16).
+
+    a: (M, d) -> hệ số riêng từng kênh (EMA Learner Matrix, scan Hillis-Steele); a: (M,) -> theo khối (chung cho d).
+    """
     with torch.autocast(device_type=u.device.type, enabled=False):
+        if a.dim() == 2:
+            return _ema_scan_diag(u.float(), mask, h0.float(), a.float())
         return _ema_scan_fp32(u.float(), mask, h0.float(), a.float(), block)
 
 
+def _ema_scan_diag(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor):
+    """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t với a: (M, d) khác nhau theo từng kênh.
+
+    u: (B,T,M,d) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,M,d) fp32.
+    Mỗi bước là ánh xạ affine h -> A_t h + B_t (A_t = 1 - a m_t, B_t = a m_t u_t); inclusive scan log-depth
+    (Hillis-Steele) cho trạng thái sau mọi t. Mọi A_t nằm trong (0, 1] nên ổn định.
+    Trả về H (B,T,M,d) và h_cuối (B,M,d).
+    """
+    bsz, t, n, d = u.shape
+    am = a.view(1, 1, n, d) * mask.to(u.dtype).view(bsz, t, 1, 1)
+    dec = 1.0 - am
+    acc = am * u
+    off = 1
+    while off < t:
+        acc = torch.cat([acc[:, :off], dec[:, off:] * acc[:, :-off] + acc[:, off:]], dim=1)
+        dec = torch.cat([dec[:, :off], dec[:, off:] * dec[:, :-off]], dim=1)
+        off *= 2
+    h = dec * h0.unsqueeze(1) + acc
+    return h, h[:, -1]
+
+
 def _ema_scan_fp32(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
     """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t, tính song song cho mọi t.
 
@@ -308,12 +351,14 @@ class KentLayer(nn.Module):
         self.wnn = nn.ModuleList([KentWNNBlock(config) for _ in range(config.num_wnn_blocks)])
         self.norm_h = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
         self.res_out = nn.Linear(self.n * self.d, config.hidden_size, bias=config.use_bias)   # Resolution Proj M -> 1
-        a0 = torch.tensor(config.alpha_init_vector(layer_idx), dtype=torch.float32)
+        a0 = torch.tensor(
+            config.alpha_init_matrix(layer_idx) if config.ema_per_feature else config.alpha_init_vector(layer_idx),
+            dtype=torch.float32)
         self.ema_logit_init = torch.log(a0 / (1.0 - a0))
-        self.ema_logit = nn.Parameter(self.ema_logit_init.clone())             # (M,)
+        self.ema_logit = nn.Parameter(self.ema_logit_init.clone())             # (M, d) hoặc (M,)
 
     def alpha(self) -> torch.Tensor:
-        """Vector a (M,) trong (0,1), tính ở fp32."""
+        """EMA Learner Matrix a (M, d) (hoặc vector (M,)) trong (0,1), tính ở fp32."""
         return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999)
 
     def forward(self, x, mask, h0, chunked: bool = False):
@@ -331,7 +376,8 @@ class KentLayer(nn.Module):
             h0 = u.new_zeros((u.shape[0], self.n, self.d))
         h, h_last = ema_scan(u, m, h0, self.alpha(), self.block_size)
         h = h.reshape(b, t, self.n, self.d)
-        dx = self.res_out(self.norm_h(h).to(x.dtype).reshape(b, t, self.n * self.d))
+        spec = dct2(self.norm_h(h).to(x.dtype))       # DCT-II Proj trên trục M của H
+        dx = self.res_out(spec.reshape(b, t, self.n * self.d))
         return x + dx, h_last