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