Update README.md

3fb305f about 6 hours ago

5.71 kB

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

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