README.md

---
language: en
license: apache-2.0
library_name: transformers
pipeline_tag: text-generation
tags:
  - fwkv
  - rosa
  - rwkv
  - linear-transformer
  - recurrent-neural-network
  - chat
  - ultrachat
  - custom-architecture
---

# FWKV-ROSA --- Read more [here](https://flamef0x.github.io/papers/)

A **56M‑parameter**, from‑scratch recurrent language model that combines a
**per‑channel leaky integrator (FWKV)** with the **RWKV‑8 ROSA copy‑signal**
mechanism. It is a research experiment designed to explore how far purely
recurrent architectures can go on small‑scale, curated conversational data.

The model was chat‑tuned on `HuggingFaceH4/ultrachat_200k` and uses a simple
two‑role template:

```
<|user|> Your message
<|assistant|> Model reply<|endoftext|>
```

## Model Description

- **Architecture:** 14 stacked FWKV blocks. Each block replaces the standard attention with a fixed, data‑independent decayed accumulator:
  `stateₜ = W·stateₜ₋₁ + kₜ·vₜ`, where `W = clamp(sigmoid(w), min=0.1)`.
  The recurrence is computed exactly via a vectorised parallel scan (no
  approximations).
- **ROSA (Rapid Online Suffix Automaton):** A parameter‑free, causal
  predictor that injects the token that historically followed the longest
  matching suffix of the current context. The ROSA signal is embedded and
  added to the input representation of each token.
- **Factorised embedding/head:** 128‑dimensional embedding space, projected
  to a 512‑dimensional model space, with tied weights.
- **Context length:** 1024 tokens. No positional embeddings are used, making
  the context window “free” in terms of parameters.
- **Tokenizer:** GPT‑2 tokenizer extended with the special tokens `<|user|>`
  and `<|assistant|>`.

## Uses

### Direct Use

FWKV-ROSA is intended for **research on efficient language models** and for
**educational demonstrations** of recurrent architectures. You can chat with
it in a multi‑turn setting using the template above.

### Out‑of‑Scope Use

- This model is **not** suitable for any production or safety‑critical
  application.
- It has not been aligned with RLHF or other safety methods and may generate
  inappropriate or harmful content.
- The limited size and training data mean it cannot be relied upon for factual
  knowledge or reasoning.

### Bias, Risks, and Limitations

- Trained on a relatively small synthetic dataset, the model can produce
  repetitive or nonsensical output.
- The ROSA copy mechanism may occasionally copy large chunks of the user’s
  prompt verbatim.
- Biases present in the original UltraChat data are likely reflected in the
  model’s responses.

## How to Get Started

The model relies on a custom architecture. To load it, you must provide the
`modeling_fwkv.py` file (found in the repository) and trust the remote code:

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "FlameF0X/FWKV-ROSA",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained("FWKV/FWKV-ROSA")
```

Then format your prompts exactly with the chat tokens:

```python
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device).eval()

prompt = "<|user|> What is the capital of France?\n<|assistant|>"
input_ids = tokenizer.encode(prompt)
# ROSA IDs must be computed – you can import `rosa` from modeling_fwkv
from modeling_fwkv import rosa
rosa_ids = torch.tensor([rosa(input_ids)], device=device)
out = model(input_ids=torch.tensor([input_ids], device=device),
            rosa_ids=rosa_ids, use_cache=True)
# Continue autoregressive sampling…
```

For a fully working chat demo, see the Gradio app provided in the repository.

## Training Details

### Dataset

- **Name:** [UltraChat 200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k)
- **Splits:** `train_sft` (20,000 examples), `test_sft` (1,000 examples)
- **Format:** multi‑turn conversations; only assistant tokens contribute to the loss.

### Training Procedure

| Hyperparameter | Value |
|----------------|-------|
| Architecture   | 14 FWKV blocks, d_model=512, d_emb=128 |
| FFN multiplier | 4 |
| WKV decay floor| 0.1 |
| Batch size (per GPU) | 8 |
| Gradient accumulation | 4 |
| Effective batch size  | 32 |
| Learning rate  | 3×10⁻⁴ (cosine schedule) |
| Weight decay   | 0.1 |
| Gradient clipping | 1.0 |
| Optimizer      | AdamW (fused) |
| Precision      | bfloat16 mixed |
| Epochs         | 2 |
| Hardware       | 1× NVIDIA T4 (15 GB) |
| Training time  | ~2 hours |
| Speed          | ~80,000 tokens/second |

Gradient checkpointing was enabled to fit the 1024‑token sequences in memory.

## Evaluation

| Metric                | Value   |
|-----------------------|---------|
| Validation loss       | 4.216   |
| Validation perplexity | 67.78   |

The perplexity is relatively high due to the small model size and limited
training data. It is comparable to other similarly‑sized recurrent LMs on
UltraChat.

## Environmental Impact

The training ran for about **2 hours** on a single **NVIDIA T4** GPU
(maximum power draw ~70 W), resulting in an estimated **0.14 kWh** of
electricity consumption and approximately **0.06 kg CO₂eq** (assuming a
grid carbon intensity of 0.4 kg/kWh). This is a negligible footprint.

## Technical Specifications

- **Model type:** Recurrent neural network (linear RNN)
- **Parameters:** 56.2 million
  - Backbone (FWKV blocks + embeddings): ~50M
  - ROSA embedding: ~6.2M
- **Checkpoint format:** PyTorch `safetensors`
- **Required files in the repo:**
  - `config.json`
  - `model.safetensors` (or `pytorch_model.bin`)
  - `modeling_fwkv.py`
  - `tokenizer.json` / `vocab.json` / `merges.txt`
- **Auto‑mapping:** The `config.json` includes
  `"auto_map": { "AutoModelForCausalLM": "modeling_fwkv.FWKVLanguageModel" }`,
  so loading with `trust_remote_code=True` will automatically locate the
  correct class.

## Citation

If you use FWKV-ROSA in your research, please cite it as:

```bibtex
@misc{fwkv-rosa,
  author = {FlameF0X},
  title = {FWKV-ROSA: A 56M Recurrent Chat LM with RWKV-style Decay and ROSA Copy Signal},
  year = {2025},
  howpublished = {\url{https://huggingface.co/FWKV/FWKV-ROSA}},
}
```

## Additional Information

This model was built as an experiment to test the combination of a simple
leaky integrator with the ROSA copy signal on a small, clean conversational
dataset. It demonstrates that a pure linear RNN can learn to produce
coherent multi‑turn dialogue without any attention mechanisms. Feedback and
contributions are welcome!