Blame for README.md

---
license: apache-2.0
language:
- en
library_name: transformers
pipeline_tag: text-generation
tags:
- kent
- ssm
- state-space-model
- recurrent
- linear-time
- frequency-domain
- dct
- mlp
- swiglu
- text-generation
- custom_code
- fp16
- base
datasets:
- wikimedia/wikipedia
---

# KenT

KenT is a recurrent language model from the SSM family. KenT concentrates its capacity in a wide spectral representation stacked over a moderate depth: each layer has a single transformation block, plus a light learned-forgetting memory.

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

## Architectural Philosophy

- **Wide spectrum, moderate depth.** KenT stacks identical recurrent layers (24 in this ~1B model), each containing exactly
  one transformation block (a SwiGLU MLP with hidden size ~8/3 x D). The layers are connected in sequence through the representation stream, and the last layer produces
  the final output. The representation stream feeding `lm_head` is wide (1536), and each layer's Bilinear projection expands it
  into a large resolution axis (N = 24 DCT positions) with enough features per position (D = 256) before the DCT.
- **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 by widening the representation stream and the spectral
  axes (N and D together) and then by adding layers; the MLP inside each block is kept small.
- **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.
- **Width and depth together.** There is no inner chain of blocks inside a layer: each layer applies one DCT/MLP block over a
  large N x D grid, and memory timescale varies across the 24 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-1B-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-1B-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.