KenT_400M_Base

Vi-LAB/KenT_400M_Base — mô hình ngôn ngữ ~399M tham số theo kiến trúc Oscillatory SSM + Context Queue. (English version below / Bản tiếng Anh ở bên dưới)

---

🇻🇳 Tiếng Việt (chính)

Kiến trúc

{token,pos} ─┬─ token    : Embedding(vocab→D+1)                    = X [B,D+1,1]
             └─ position : Sinusoidal(T)(1→E) → [1,t,t²,…,t^D]     = P [B,E,D+1]
PX = P·X  (D+1→1)  →  [B,E]       (vector token = hệ số đa thức, tính tại E điểm sin/cos)
 ├─ Direction : N × [RMSNorm → FFN(GeLU, m=K·E) → Add] → d
 └─ Magnitude : N × [RMSNorm → FFN(GeLU, m=K·E) → Add] → m
Δx[i] = Δx[i-1] + m[i-1] ⊙ L2Norm(d[i-1])          ← trạng thái SSM (cộng dồn vận tốc)
Q = Linear(E→D+1)(RMSNorm(Δx)) + Embedding(token)   ← skip-connection
Context Queue (FIFO, tối đa C slot lưu K,V) ── Cross-attn (attend toàn bộ C vector)
h = Q + CrossAttn(RMSNorm(Q), K, V)  ──Cache──▶ push (K,V) vào queue (slot cũ nhất bị đẩy ra)
Generate: Linear(D+1→vocab)(RMSNorm(h)) → Tokenizer (argmax / sampling)

Context Queue là hàng đợi FIFO: mỗi bước sinh, vector h được đẩy vào cuối hàng; khi đầy C slot thì slot cũ nhất bị loại. Trạng thái suy luận có kích thước cố định (Δx + queue + pos), không tăng theo độ dài chuỗi.

Tham số	Giá trị	Ý nghĩa
hidden_dim (D)	767	số bậc đa thức → vector D+1 = 768
embed_dim (E)	1024	số điểm sinusoidal / chiều luồng vận tốc
num_layers (N)	19	số khối FFN mỗi nhánh
ffn_mult (K)	4	m = K·E
context_slots (C)	256	số slot tối đa của Context Queue
num_heads	12	số head của cross-attn

Tổng ≈ 399M tham số (vocab 50 304, không tie embedding).

Cấu trúc thư mục

File	Vai trò
model.safetensors	Trọng số
config.json	Cấu hình model
model.py	Kiến trúc (KenTConfig, KenTForCausalLM, ContextQueue) — tự tải khi trust_remote_code=True
generate_config.json	Cài đặt sinh mặc định
infer.py	Công cụ suy luận một file (CLI + load / generate_text)

Quick start

Không cần clone repo hay có file local nào.

pip install torch transformers safetensors
python -c "import urllib.request as u; u.urlretrieve('https://huggingface.co/Vi-LAB/KenT_400M_Base/resolve/main/infer.py','infer.py')"
python infer.py "Hello"

Lần chạy đầu infer.py tự tải trọng số, code kiến trúc và generate_config.json từ Hub về cache; các lần sau dùng lại. Tự chọn GPU/CPU và stream văn bản ra màn hình.

python infer.py                                        # nhập prompt liên tục
python infer.py "Hello" --greedy --max_new_tokens 128
python infer.py "Hello" --temperature 0.7 --seed 1
python infer.py "Hello" --random_init               # thử pipeline, không cần trọng số (đầu ra là nhiễu)

Dùng trong code

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("Vi-LAB/KenT_400M_Base", trust_remote_code=True).eval()

ids = torch.tensor([[50256] + tok.encode("Hello", add_special_tokens=False)])   # 50256 = <|endoftext|> (BOS)
out = model.generate(ids, max_new_tokens=64)           # mặc định lấy từ generate_config.json
print(tok.decode(out[0], skip_special_tokens=True))

Hoặc dùng helper của infer.py:

from infer import load, generate_text
model, tok = load()
print(generate_text(model, tok, "Hello", stream=False))

Suy luận từng bước với trạng thái cố định:

o = model(ids, use_cache=True, last_only=True)                       # nạp prompt
o = model(next_token[:, None], state=o.state, use_cache=True, last_only=True)

trust_remote_code=True là bắt buộc vì kiến trúc KenT nằm trong model.py của repo, không có sẵn trong transformers.

Chạy offline / máy chủ không có mạng

huggingface-cli download Vi-LAB/KenT_400M_Base --local-dir ./KenT_400M_Base     # trên máy có mạng
python infer.py "Hello" --model ./KenT_400M_Base                            # trên máy đích

Yêu cầu phần cứng

Trọng số ~1,6 GB (fp32). Chạy được trên CPU; GPU tự dùng autocast bf16. Trạng thái suy luận kích thước cố định nên bộ nhớ không tăng theo độ dài chuỗi.

Lưu ý

• Tokenizer: GPT-2 (BPE, 50 257 token). vocab_size = 50304 là 50 257 làm tròn lên bội số của 64 (các hàng thừa không dùng tới). GPT-2 không có token pad/bos riêng nên pad = bos = eos = 50256 (<|endoftext|>). infer.py tự thêm BOS vào đầu prompt; nếu tự gọi generate, hãy đặt BOS ở đầu chuỗi.
• GPT-2 tokenizer thiên về tiếng Anh nên tiếng Việt bị tách thành nhiều token (chuỗi dài hơn).
• Đừng ép trọng số sang bf16/fp16 bằng torch_dtype/.half(): nhánh Velocity Predictor chạy fp32 và sẽ lỗi lệch dtype. Hãy giữ fp32 và dùng torch.autocast (đã làm sẵn trong infer.py).
• Batch suy luận nên có cùng độ dài prompt (không left-padding), vì queue được dùng chung theo thời gian.
• Một vài chi tiết sơ đồ không nêu rõ nên được chọn như sau: m là vector E chiều nhân từng phần tử; K, V được tính một lần khi push vào queue; có thêm một slot "null" học được để softmax xác định khi queue rỗng; mọi Linear không bias.

---

🇬🇧 English

Architecture

KenT is a ~399M-parameter recurrent language model: the current token (as polynomial coefficients) is evaluated at sinusoidal position points, passed through two residual FFN stacks (Direction d and Magnitude m), and integrated into an SSM state Δx[i] = Δx[i-1] + m[i-1] ⊙ L2Norm(d[i-1]). The state becomes a query Q that cross-attends over a Context Queue: a FIFO of at most C cached (K, V) slots. Each step pushes the new vector into the queue (the oldest slot is evicted when full). Inference state is fixed-size (Δx + queue + position).

Files

model.safetensors (weights) · config.json · model.py (architecture, auto-downloaded with trust_remote_code=True) · generate_config.json (default sampling settings) · infer.py (single-file inference CLI + load / generate_text).

Quick start

No clone or local files needed.

pip install torch transformers safetensors
python -c "import urllib.request as u; u.urlretrieve('https://huggingface.co/Vi-LAB/KenT_400M_Base/resolve/main/infer.py','infer.py')"
python infer.py "Hello"

The first run downloads weights, architecture code and generate_config.json from the Hub into the cache. Use --model ./dir for an offline copy and --random_init to smoke-test the pipeline without weights (output is noise).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("Vi-LAB/KenT_400M_Base", trust_remote_code=True).eval()
ids = torch.tensor([[50256] + tok.encode("Hello", add_special_tokens=False)])   # 50256 = BOS
print(tok.decode(model.generate(ids, max_new_tokens=64)[0], skip_special_tokens=True))

Notes

• Tokenizer: GPT-2 BPE (50,257 tokens); vocab_size=50304 is that rounded up to a multiple of 64. pad = bos = eos = 50256 (<|endoftext|>).
• Keep weights in fp32 and use torch.autocast on GPU; casting all weights to bf16/fp16 breaks the fp32 velocity-predictor path.
• Ambiguities in the diagram were resolved as: elementwise vector magnitude, K/V computed at push time, a learned null slot for the empty queue, bias-free linears.