---
license: apache-2.0
language:
  - vi
  - en
library_name: transformers
pipeline_tag: text-generation
tags:
  - kent
  - ssm
  - state-space-model
  - recurrent
  - linear-time
  - frequency-domain
  - dct
  - swiglu
  - text-generation
  - custom_code
  - fp16
  - 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

```bash
pip install torch transformers
```

### Standard text generation (`generate`)

```python
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`):

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

```python
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:

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