Vi-LAB/KenT-300M-Base

KenT

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

• 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. 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 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.

Repo Structure

File	Role
model.py	Model implementation (KentConfig, KentModel, KentForCausalLM, KentCache)
config.json	Standard model configuration
generation_config.json	Default generation configuration
tokenizer.json, tokenizer_config.json	Tokenizer
infer.py	Parallel chunked prefix loading and the generation loop
model.safetensors	FP16 weights

Quickstart

pip install torch transformers

Standard text generation (generate)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

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
).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))

The recurrent state is returned as cache_params (similar to other SSM models in transformers):

out = model(**inputs, use_cache=True)
cache = out.cache_params                       # recurrent state
out = model(input_ids=next_token, cache_params=cache, use_cache=True)

Long documents: parallel chunked prefix loading

When processing in batches, each row of the batch is one chunk of the same document; the chunks are processed in parallel with the chunked=True flag, and the final state is returned in cache_params.

from infer import prefill, stream

ids = tok(open("doc.txt", encoding="utf-8").read(), return_tensors="pt").input_ids[0]

# 1) load the prefix only, get the last token's logits and the state
logits, cache = prefill(model, ids, chunk_size=1024, max_parallel_chunks=16)

# 2) load the prefix, then generate (streaming)
for token_id in stream(model, ids, gen_cfg):
    print(tok.decode([token_id]), end="", flush=True)

Command line:

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 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.