KenT-30M-Base / README.md
KenteaKe11's picture
Upload 3 files
9cd4795 verified
|
Raw History Blame Contribute Delete
2.57 kB
metadata
license: apache-2.0
language:
  - en
library_name: transformers
pipeline_tag: text-generation
tags:
  - kent
  - recurrent
  - linear-time
  - frequency-domain
  - dct
  - text-generation
  - custom_code
  - fp16
  - base

KenT-30M-Base

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))              # 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 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
hidden size n 256 (tied embedding / lm_head, vocab 50257)
layers 24
resolution M x features d 12 x 96 (M·d = 1152 = 4.5n)
SwiGLU hidden 384 (4d)
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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

repo = "Vi-LAB/KenT-30M-Base"  # or a local directory
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, dtype=torch.float16).to("cuda").eval()
gen_cfg = GenerationConfig.from_pretrained(repo)

inputs = tok("Hello, today", return_tensors="pt").to("cuda")
out = model.generate(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, generation_config=gen_cfg)
print(tok.decode(out[0], skip_special_tokens=True))

Long documents: python infer.py --model ./KenT-30M-Base --prompt-file doc.txt --chunk-size 1024 (chunks are processed in parallel and merged with an affine scan; the result equals sequential processing).

Notes

  • Research model, trained from scratch; evaluate before real-world use.
  • The tokenizer and weights must match vocab_size in config.json.