KenT

KenT is a recurrent language model from the SSM family. Rather than stacking more layers, KenT concentrates its capacity in the depth of the transformation inside a single cell, and keeps temporal memory as simple as possible.

This repo contains trust_remote_code code; weights are stored in FP16.

Architectural Philosophy

  • One cell, no layer stacking. KenT is a single recurrent cell; the model is not replicated into multiple layers. Depth comes not from the number of layers but from a chain of transformation blocks applied in sequence inside the cell.
  • 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.
  • 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.

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-100M-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-100M-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.
  • 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.
Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support