KenT-300M: EMA memory with learned per-layer alpha, 1 WNN per layer, 24 layers

README.md

@@ -22,28 +22,30 @@ tags:
 
 # KenT
 
-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.
+KenT is a recurrent language model from the SSM family. KenT concentrates its capacity in depth: many compact layers, each with a single transformation block, and a light learned-forgetting memory.
 
 > This repo contains `trust_remote_code` code; weights are stored in **FP16**.
 
 ## Architectural Philosophy
 
-- **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.
+- **Deep stack of compact layers.** KenT stacks many identical recurrent layers (24 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
+  the final output. Depth is prioritized: the deeper, the better, while width is kept at a level that stays balanced with it.
 - **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.
+  rather than in the vocabulary table. Parameters were scaled from ~100M to ~300M mainly by adding depth, with width unchanged,
+  so that both axes stay balanced (width still matters, but is not where the extra capacity went).
 - **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
   **fixed, non-learned** operator, interleaved with nonlinear blocks. The fixed part carries built-in structure; the learned part
   focuses solely on transforming features.
-- **Memory via simple accumulation.** History is summarized with a cumulative average: no dedicated parameters and no implicit
-  forgetting mechanism. The decoding state has **constant** size regardless of context length.
-- **Two separate depth axes.** The depth of feature transformation within a single step and the length of history remembered are
-  two independent axes: increasing one does not change the other.
-- **Chunk-parallelizable.** Because memory is an associative accumulation, a long document can be split into multiple chunks,
-  processed concurrently, and then merged, yielding the same result as sequential processing.
+- **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.
+- **Depth is the main axis.** There is no inner chain of blocks inside a layer any more: feature depth and memory timescale both
+  grow with the number of 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.
 
 The points above are **design hypotheses** that need to be verified experimentally, not performance guarantees.
 Detailed configuration and full formulas are not described here.
@@ -71,7 +73,7 @@ pip install torch transformers
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 
-repo = "Vi-LAB/KenT-100M-Base"  # replace with a repo id or a local directory
+repo = "Vi-LAB/KenT-300M-Base"  # replace with a repo id or a local directory
 tok = AutoTokenizer.from_pretrained(repo)
 model = AutoModelForCausalLM.from_pretrained(
     repo, trust_remote_code=True, dtype=torch.float16
@@ -116,13 +118,13 @@ for token_id in stream(model, ids, gen_cfg):
 Command line:
 
 ```bash
-python infer.py --model Vi-LAB/KenT-100M-Base --prompt-file doc.txt --chunk-size 1024
+python infer.py --model Vi-LAB/KenT-300M-Base --prompt-file doc.txt --chunk-size 1024
 ```
 
 `cache.clone()` lets you reuse the same prefix to generate multiple different branches without reloading it.
 
 ## Notes
 
-- Memory is a uniform cumulative average over every timestep seen; the model does not learn to forget or to select on its own.
+- Memory is an exponential moving average with a learned per-layer coefficient; it forgets old timesteps but does not select content on its own.
 - The tokenizer and weights must match `vocab_size` in `config.json`.
 - This is a research model; please evaluate it before using it for real-world tasks.

config.json

@@ -10,19 +10,20 @@
   },
   "vocab_size": 50257,
   "hidden_size": 512,
-  "num_hidden_layers": 5,
+  "num_hidden_layers": 24,
   "resolution_size": 8,
   "feature_dim": 512,
-  "num_wnn_blocks": 2,
   "intermediate_size": 2048,
   "rms_norm_eps": 1e-05,
   "use_bias": true,
   "initializer_range": 0.02,
-  "scan_group_size": 8,
   "use_cache": true,
   "tie_word_embeddings": true,
   "pad_token_id": 50256,
   "bos_token_id": 50256,
   "eos_token_id": 50256,
-  "dtype": "float16"
-}
+  "dtype": "float16",
+  "ema_alpha_first": 0.5,
+  "ema_alpha_last": 0.01,
+  "ema_block_size": 128
+}

infer.py

@@ -4,13 +4,13 @@
 
 Quy ước:
   * Trục batch khi nạp prefix = các chunk liên tiếp của MỘT tài liệu.
-  * Mỗi layer tính cumsum cục bộ trong chunk (song song), sau đó ghép các chunk bằng
-    exclusive scan đệ quy; kết quả bằng đúng CMA tuần tự.
-  * Sau prefix, trạng thái chỉ còn (H, t) mỗi layer -> giải mã O(1) bộ nhớ theo độ dài.
+  * Mỗi layer tính EMA song song theo khối, các khối/chunk được ghép bằng affine-scan;
+    kết quả bằng đúng EMA tuần tự H_t = (1 - a) H_{t-1} + a u_t.
+  * Sau prefix, trạng thái chỉ còn H mỗi layer -> giải mã O(1) bộ nhớ theo độ dài.
 
 Ví dụ:
-    python infer.py --model Vi-LAB/KenT-100M-Base --prompt "Xin chào"
-    python infer.py --model ./KenT-100M-Base --prompt-file doc.txt --chunk-size 1024
+    python infer.py --model Vi-LAB/KenT-300M-Base --prompt "Xin chào"
+    python infer.py --model ./KenT-300M-Base --prompt-file doc.txt --chunk-size 1024
 """
 
 import argparse
@@ -66,7 +66,7 @@ def prefill(
     """Nạp prefix 1D `input_ids` (T,). Trả về (logits token cuối (V,), cache).
 
     Các chunk được xử lý song song theo nhóm tối đa `max_parallel_chunks`; giữa các nhóm,
-    trạng thái (H, t) được truyền tiếp qua cache. Truyền `cache` có sẵn để nạp tiếp
+    trạng thái H được truyền tiếp qua cache. Truyền `cache` có sẵn để nạp tiếp
     một đoạn tài liệu mới vào cùng trạng thái.
     """
     device = next(model.parameters()).device
@@ -131,7 +131,7 @@ def stream(
     max_parallel_chunks: Optional[int] = None,
     cache=None,
 ) -> Iterator[int]:
-    """Sinh từng token: prefill chunk-song-song, rồi giải mã từng bước với trạng thái (H, t)."""
+    """Sinh từng token: prefill chunk-song-song, rồi giải mã từng bước với trạng thái H."""
     chunk_size = chunk_size or getattr(gen_cfg, "kent_prefill_chunk_size", None) or 512
     max_parallel_chunks = max_parallel_chunks or getattr(gen_cfg, "kent_max_parallel_chunks", None) or 16
     device = next(model.parameters()).device

model.py

@@ -1,19 +1,24 @@
 # 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 thực hiện một bước:
+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
-    for l in 0..L-1:  Z = A_N [ Z + SwiGLU(Norm(Z)) ]
-    H_t = ((t-1) H_{t-1} + u_t) / t           # CMA, u_t = Z
+    u   = A_N [ Z + SwiGLU(Norm(Z)) ]          # một WNN Block
+    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))
 
-Vì CMA là tổng tích lũy nên toàn bộ chuỗi trong một layer có thể tính song song
-(cumsum), và prefix dài có thể được chia thành các chunk rồi ghép bằng exclusive
-scan đệ quy. Trạng thái giải mã chỉ là (H, t) cho mỗi layer: O(1) theo độ dài.
+Độ sâu đến từ việc xếp chồng nhiều layer (mỗi layer một WNN), không còn chuỗi L WNN trong
+một layer. Bộ nhớ H có cơ chế quên: a_layer = sigmoid(ema_logit) là một số vô hướng học được
+cho mỗi layer; khởi tạo giảm dần theo độ sâu (layer nông quên nhanh, layer sâu nhớ lâu).
 
-Trọng số được lưu ở FP16; mọi phép cộng dồn CMA và normalization chạy ở FP32.
+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, 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.
+
+Trọng số được lưu ở FP16; EMA và normalization chạy ở FP32.
 """
 
 import math
@@ -39,15 +44,17 @@ 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 cell nối liền nhau (cell cuối sinh đầu ra cuối).
+        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.
-        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.
         initializer_range: độ lệch chuẩn khởi tạo trọng số.
-        scan_group_size: kích thước nhóm của exclusive scan đệ quy giữa các chunk.
+        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
+            theo cấp số nhân. a_layer 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ớ).
         use_cache: trả về KentCache khi suy luận.
     """
 
@@ -57,15 +64,16 @@ class KentConfig(PretrainedConfig):
         self,
         vocab_size: int = 50257,
         hidden_size: int = 512,
-        num_hidden_layers: int = 5,
+        num_hidden_layers: int = 24,
         resolution_size: int = 8,
         feature_dim: int = 512,
-        num_wnn_blocks: int = 2,
         intermediate_size: int = 2048,
         rms_norm_eps: float = 1e-5,
         use_bias: bool = True,
         initializer_range: float = 0.02,
-        scan_group_size: int = 8,
+        ema_alpha_first: float = 0.5,
+        ema_alpha_last: float = 0.01,
+        ema_block_size: int = 128,
         use_cache: bool = True,
         pad_token_id: int = 50256,
         bos_token_id: int = 50256,
@@ -78,12 +86,13 @@ class KentConfig(PretrainedConfig):
         self.num_hidden_layers = num_hidden_layers
         self.resolution_size = resolution_size
         self.feature_dim = feature_dim
-        self.num_wnn_blocks = num_wnn_blocks
         self.intermediate_size = intermediate_size
         self.rms_norm_eps = rms_norm_eps
         self.use_bias = use_bias
         self.initializer_range = initializer_range
-        self.scan_group_size = scan_group_size
+        self.ema_alpha_first = ema_alpha_first
+        self.ema_alpha_last = ema_alpha_last
+        self.ema_block_size = ema_block_size
         self.use_cache = use_cache
         super().__init__(
             pad_token_id=pad_token_id,
@@ -93,6 +102,14 @@ class KentConfig(PretrainedConfig):
             **kwargs,
         )
 
+    def alpha_init(self, layer_idx: int) -> float:
+        """a_layer khởi tạo của layer `layer_idx` (giảm theo cấp số nhân theo độ sâu)."""
+        n = self.num_hidden_layers
+        if n <= 1:
+            return float(self.ema_alpha_first)
+        r = layer_idx / (n - 1)
+        return float(self.ema_alpha_first * (self.ema_alpha_last / self.ema_alpha_first) ** r)
+
 
 # --------------------------------------------------------------------------- #
 # Fixed DCT-II operator
@@ -115,66 +132,66 @@ def dct2(z: torch.Tensor) -> torch.Tensor:
 
 
 # --------------------------------------------------------------------------- #
-# CMA: song song theo thời gian + ghép chunk bằng exclusive scan đệ quy
+# EMA: song song theo thời gian (khối tam giác + affine-scan giữa các khối)
 # --------------------------------------------------------------------------- #
-def _exclusive_scan(x: torch.Tensor, group: int) -> torch.Tensor:
-    """Exclusive prefix-sum dọc dim 0, tính đệ quy theo nhóm.
+def _affine_scan(a: torch.Tensor, b: torch.Tensor):
+    """Inclusive scan dọc dim 1 của các ánh xạ affine h -> a*h + b (Hillis-Steele, log K bước).
 
-    Các nhóm được scan song song; tổng của từng nhóm lại được scan đệ quy ở tầng trên,
-    độ sâu đệ quy O(log_group K).
+    a: (B,K) | b: (B,K,F). Phần tử k được áp dụng SAU các phần tử < k.
+    Trả về (A, B) sao cho trạng thái sau khối k là A[k]*h_vào + B[k].
     """
-    k = x.shape[0]
-    if k <= group:
-        return torch.cumsum(x, dim=0) - x
-    pad = (-k) % group
-    if pad:
-        x = torch.cat([x, x.new_zeros((pad,) + tuple(x.shape[1:]))], dim=0)
-    g = x.shape[0] // group
-    xg = x.reshape((g, group) + tuple(x.shape[1:]))
-    inner = torch.cumsum(xg, dim=1) - xg              # exclusive trong nhóm
-    outer = _exclusive_scan(xg.sum(dim=1), group)     # exclusive giữa các nhóm (đệ quy)
-    out = inner + outer.unsqueeze(1)
-    return out.reshape((g * group,) + tuple(x.shape[1:]))[:k]
-
-
-def cma_scan(u, mask, h0, n0):
-    """CMA song song theo thời gian cho từng hàng batch độc lập.
-
-    u: (B,T,N,D) fp32 | mask: (B,T) long {0,1} | h0: (B,N,D) fp32 | n0: (B,) long
-    Trả về H (B,T,N,D), h_cuối (B,N,D), n_cuối (B,).
-    """
-    s0 = h0 * n0.to(u.dtype)[:, None, None]
-    s = torch.cumsum(u * mask.to(u.dtype)[:, :, None, None], dim=1) + s0[:, None]
-    n = torch.cumsum(mask, dim=1) + n0[:, None]
-    h = s / n.clamp_min(1).to(u.dtype)[:, :, None, None]
-    return h, h[:, -1].clone(), n[:, -1].clone()
+    k = a.shape[1]
+    off = 1
+    while off < k:
+        b = torch.cat([b[:, :off], a[:, off:, None] * b[:, :-off] + b[:, off:]], dim=1)
+        a = torch.cat([a[:, :off], a[:, off:] * a[:, :-off]], dim=1)
+        off *= 2
+    return a, b
 
 
-def cma_chunk_scan(u, mask, h0, n0, group: int):
-    """CMA cho K chunk liên tiếp của MỘT tài liệu (trục batch = trục chunk).
+def ema_scan(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.
 
-    1. cumsum cục bộ trong từng chunk (song song trên mọi chunk);
-    2. tổng từng chunk -> exclusive scan đệ quy -> carry cho mỗi chunk;
-    3. cộng trạng thái khởi tạo (h0, n0) rồi chia đếm để được H toàn cục.
+    u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32
+    a: số vô hướng fp32 trong (0,1).
+    Trả về H (B,T,N,D) và h_cuối (B,N,D).
 
-    u: (K,C,N,D) | mask: (K,C) | h0: (1,N,D) | n0: (1,)
-    Trả về H (K,C,N,D), h_cuối (1,N,D), n_cuối (1,).
+    Trong một khối: H_t = sum_{s<=t} a m_s [prod_{s<r<=t}(1 - a m_r)] u_s + [prod_{r<=t}(1 - a m_r)] h_vào,
+    với tích tính qua hiệu của cumsum log (luôn <= 0 nên ổn định). Các khối được nối bằng affine-scan.
     """
-    local_s = torch.cumsum(u * mask.to(u.dtype)[:, :, None, None], dim=1)
-    local_n = torch.cumsum(mask, dim=1)
-    carry_s = _exclusive_scan(local_s[:, -1], group) + h0 * n0.to(u.dtype)[:, None, None]
-    carry_n = _exclusive_scan(local_n[:, -1], group) + n0
-    s = local_s + carry_s[:, None]
-    n = local_n + carry_n[:, None]
-    h = s / n.clamp_min(1).to(u.dtype)[:, :, None, None]
-    return h, h[-1:, -1].clone(), n[-1:, -1].clone()
+    bsz, t, n, d = u.shape
+    f = n * d
+    blk = max(1, min(block, t))
+    pad = (-t) % blk
+    uf = u.reshape(bsz, t, f)
+    m = mask.to(u.dtype)
+    if pad:  # padding ở cuối: mask = 0 nên trạng thái không đổi
+        uf = F.pad(uf, (0, 0, 0, pad))
+        m = F.pad(m, (0, pad))
+    k = (t + pad) // blk
+    uk = uf.reshape(bsz, k, blk, f)
+    am = a * m.reshape(bsz, k, blk)                              # (B,K,blk)
+
+    cum = torch.cumsum(torch.log1p(-am), dim=2)                  # log prod_{r<=t}(1 - a m_r)
+    tril = torch.ones(blk, blk, dtype=torch.bool, device=u.device).tril()
+    decay = torch.exp((cum.unsqueeze(-1) - cum.unsqueeze(-2)).masked_fill(~tril, float("-inf")))
+    local = torch.matmul(decay * am.unsqueeze(-2), uk)           # (B,K,blk,F): khối bắt đầu từ H = 0
+
+    a_blk = torch.exp(cum[..., -1])                              # (B,K): hệ số phân rã cả khối
+    a_p, b_p = _affine_scan(a_blk, local[:, :, -1])
+    h0f = h0.reshape(bsz, 1, f)
+    s = a_p.unsqueeze(-1) * h0f + b_p                            # trạng thái sau mỗi khối
+    h_in = torch.cat([h0f, s[:, :-1]], dim=1)                    # trạng thái vào mỗi khối
+    h = local + torch.exp(cum).unsqueeze(-1) * h_in.unsqueeze(2)
+    h = h.reshape(bsz, k * blk, f)[:, :t].reshape(bsz, t, n, d)
+    return h, s[:, -1].reshape(bsz, n, d)
 
 
 # --------------------------------------------------------------------------- #
-# Cache (state tuple H, t cho từng layer)
+# Cache (trạng thái H cho từng layer)
 # --------------------------------------------------------------------------- #
 class KentCache:
-    """Trạng thái hồi quy của KenT: H (mean, fp32) theo layer và bộ đếm t."""
+    """Trạng thái hồi quy của KenT: H (EMA, fp32) theo layer."""
 
     def __init__(self, config: KentConfig, batch_size: int = 1, device=None):
         shape = (batch_size, config.resolution_size, config.feature_dim)
@@ -182,12 +199,11 @@ class KentCache:
             torch.zeros(shape, dtype=torch.float32, device=device)
             for _ in range(config.num_hidden_layers)
         ]
-        self.t = torch.zeros(batch_size, dtype=torch.long, device=device)
         self.steps = 0  # số vị trí đã nạp (đếm ở phía Python, không cần sync GPU)
 
     @property
     def batch_size(self) -> int:
-        return self.t.shape[0]
+        return self.h[0].shape[0]
 
     @property
     def seen_tokens(self) -> int:
@@ -199,24 +215,20 @@ class KentCache:
     def reset(self):
         for h in self.h:
             h.zero_()
-        self.t.zero_()
         self.steps = 0
 
     def clone(self) -> "KentCache":
         new = KentCache.__new__(KentCache)
         new.h = [h.clone() for h in self.h]
-        new.t = self.t.clone()
         new.steps = self.steps
         return new
 
     def to(self, device):
         self.h = [h.to(device) for h in self.h]
-        self.t = self.t.to(device)
         return self
 
     def reorder_cache(self, beam_idx: torch.LongTensor):
         self.h = [h.index_select(0, beam_idx.to(h.device)) for h in self.h]
-        self.t = self.t.index_select(0, beam_idx.to(self.t.device))
 
 
 # --------------------------------------------------------------------------- #
@@ -273,35 +285,41 @@ class KentWNNBlock(nn.Module):
 
 
 class KentLayer(nn.Module):
-    """Một bước KenT: (X_t, H_{t-1}, t) -> (X_{t+1}, H_t)."""
+    """Một bước KenT: (X_t, H_{t-1}) -> (X_{t+1}, H_t), với một WNN Block và EMA học được."""
 
-    def __init__(self, config: KentConfig):
+    def __init__(self, config: KentConfig, layer_idx: int = 0):
         super().__init__()
         self.n = config.resolution_size
         self.d = config.feature_dim
-        self.group = config.scan_group_size
+        self.block_size = config.ema_block_size
         self.norm_x = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
         self.b_in = KentBilinear(config.hidden_size, self.n * self.d, config.use_bias)
-        self.blocks = nn.ModuleList([KentWNNBlock(config) for _ in range(config.num_wnn_blocks)])
+        self.wnn = KentWNNBlock(config)
         self.norm_h = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
         self.b_out = KentBilinear(self.n * self.d, config.hidden_size, config.use_bias)
+        a0 = config.alpha_init(layer_idx)
+        self.ema_logit_init = math.log(a0 / (1.0 - a0))
+        self.ema_logit = nn.Parameter(torch.tensor(self.ema_logit_init))
 
-    def forward(self, x, mask, h0, n0, chunked: bool = False):
+    def alpha(self) -> torch.Tensor:
+        """a_layer in (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):
         b, t, _ = x.shape
         v = self.b_in(self.norm_x(x)).view(b, t, self.n, self.d)
-        z = dct2(v)                                   # Z^(0) = A_N V
-        for blk in self.blocks:
-            z = blk(z)
-        u = z.float()                                 # CMA cộng dồn ở fp32
-        if h0 is None:
-            h0 = u.new_zeros((1 if chunked else b, self.n, self.d))
-            n0 = torch.zeros(h0.shape[0], dtype=torch.long, device=x.device)
-        if chunked:
-            h, h_last, n_last = cma_chunk_scan(u, mask, h0, n0, self.group)
+        u = self.wnn(dct2(v)).float()                 # Z^(0) = A_N V -> một WNN Block; EMA ở fp32
+        if chunked:  # trục batch là các chunk liên tiếp: nối thành một chuỗi cho EMA
+            u = u.reshape(1, b * t, self.n, self.d)
+            m = mask.reshape(1, b * t)
         else:
-            h, h_last, n_last = cma_scan(u, mask, h0, n0)
+            m = mask
+        if h0 is None:
+            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.b_out(self.norm_h(h).to(x.dtype).reshape(b, t, self.n * self.d))
-        return x + dx, h_last, n_last
+        return x + dx, h_last
 
 
 # --------------------------------------------------------------------------- #
@@ -342,6 +360,9 @@ class KentPreTrainedModel(PreTrainedModel):
             nn.init.normal_(module.weight, mean=0.0, std=std)
         elif isinstance(module, KentRMSNorm):
             nn.init.ones_(module.weight)
+        elif isinstance(module, KentLayer):
+            with torch.no_grad():
+                module.ema_logit.fill_(module.ema_logit_init)
         if isinstance(module, KentBilinear) and getattr(module, "_residual_out", False):
             # tích của hai nhánh tuyến tính => thu nhỏ mỗi nhánh (2L)^(-1/4)
             with torch.no_grad():
@@ -352,7 +373,7 @@ class KentModel(KentPreTrainedModel):
     def __init__(self, config: KentConfig):
         super().__init__(config)
         self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList([KentLayer(config) for _ in range(config.num_hidden_layers)])
+        self.layers = nn.ModuleList([KentLayer(config, i) for i in range(config.num_hidden_layers)])
         for layer in self.layers:
             layer.b_out._residual_out = True
         self.norm_f = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
@@ -381,7 +402,7 @@ class KentModel(KentPreTrainedModel):
         chunked=False: trục batch là các chuỗi độc lập; (B, T).
         chunked=True : trục batch là K chunk liên tiếp của MỘT tài liệu; (K, C).
                        Cache vào/ra có batch size 1 (trạng thái sau chunk cuối).
-        attention_mask: 1 = token thật, 0 = padding (không đóng góp vào CMA).
+        attention_mask: 1 = token thật, 0 = padding (H giữ nguyên tại vị trí đó).
         """
         use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training)
         output_hidden_states = bool(output_hidden_states)
@@ -405,22 +426,17 @@ class KentModel(KentPreTrainedModel):
             cache_params = KentCache(self.config, 1 if chunked else b, x.device)
 
         all_hidden = () if output_hidden_states else None
-        n_last = None
         for i, layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden += (x,)
             h0 = cache_params.h[i] if cache_params is not None else None
-            n0 = cache_params.t if cache_params is not None else None
             if self.gradient_checkpointing and self.training:
-                x, h_last, n_last = self._gradient_checkpointing_func(
-                    layer.__call__, x, mask, h0, n0, chunked
-                )
+                x, h_last = self._gradient_checkpointing_func(layer.__call__, x, mask, h0, chunked)
             else:
-                x, h_last, n_last = layer(x, mask, h0, n0, chunked)
+                x, h_last = layer(x, mask, h0, chunked)
             if use_cache:
                 cache_params.h[i] = h_last
         if use_cache:
-            cache_params.t = n_last
             cache_params.steps += b * t if chunked else t
 
         x = self.norm_f(x)
@@ -555,4 +571,5 @@ __all__ = [
     "KentForCausalLM",
     "KentPreTrainedModel",
     "dct2",
+    "ema_scan",
 ]