README.md

---
license: apache-2.0
language:
- en
pipeline_tag: text-generation
library_name: transformers
tags:
- text-generation
- small-language-model
- from-scratch
- causal-lm
- decoder-only
- gqa
- rope
- swiglu
- exclusive-gated-attention
- engram
- kv-cache
---

<div align="center">

# Pragya · exp2

### A 28.5M-parameter decoder-only language model, trained from scratch. For testing, not for production.

[![Params](https://img.shields.io/badge/params-28.5M-blue?style=flat-square)](.)
[![Backbone](https://img.shields.io/badge/backbone-21.25M-blue?style=flat-square)](.)
[![Layers](https://img.shields.io/badge/layers-16-blue?style=flat-square)](.)
[![Context](https://img.shields.io/badge/context-1024_tokens-blue?style=flat-square)](.)
[![Vocab](https://img.shields.io/badge/vocab-12.3k-blue?style=flat-square)](.)
[![Val PPL](https://img.shields.io/badge/val_ppl-13.6-brightgreen?style=flat-square)](.)
[![Val Acc](https://img.shields.io/badge/val_acc-47.1%25-brightgreen?style=flat-square)](.)
[![License](https://img.shields.io/badge/license-apache--2.0-lightgrey?style=flat-square)](.)

</div>

> ### ⚠️ Not a product
> This is a scratch checkpoint. It is not aligned, factually reliable, or safe. It exists to test a from-scratch training and architecture stack and to measure controlled changes between runs.
> A massive issue is found that eos token is not working, prob is 0 so model just fills full context 
---

## Summary

| | |
|:---|:---|
| **Architecture** | `PragyaX1DecoderOnly` |
| Total parameters | 28.5M |
| Backbone (excluding embeddings and Engram tables) | 21.25M |
| Embeddings (tied input and LM head) | 4.72M |
| Engram lookup tables | 2.53M |
| Context | 1024 tokens |
| Vocabulary | 12,288 byte-level BPE |
| Validation perplexity | 13.6 |
| Validation accuracy | 47.1% |
| Intelligence Index | 12.41 |

Validation accuracy is lower than test3's 49.7% because this run's evaluation data is harder, so the two checkpoints' numbers are not directly comparable.

---

## Model card

| | |
|:---|:---|
| **Architecture** | `PragyaX1DecoderOnly` (decoder-only transformer) |
| **Layers** | 16 |
| **Hidden size** | 384 |
| **Attention heads** | 12 query, 6 key/value (GQA) |
| **Head dimension** | 32 |
| **Attention** | Exclusive Gated Attention (per-head sigmoid gate + XSA) |
| **FFN** | SwiGLU, hidden size 768 (`ffn_mult = 2`) |
| **Positional encoding** | RoPE, base 10,000 |
| **Engram memory** | Layers 1 and 8, 3-gram lookup over 9,807 classes |
| **Mixture of Experts** | Off |
| **Multi-token prediction** | Off |
| **Embeddings** | Tied (input embedding = LM head) |
| **Context length** | 1024 tokens |
| **Decoding** | KV cache (see below) |
| **Vocabulary** | 12,288 byte-level BPE (`pragya.tokenizer.json`) |

---

## Architecture

### Exclusive Gated Attention

Each attention block applies two per-head operations to the attention output, before the output projection. Both start as identity, so the model begins as plain causal attention.

**Gated attention.** A per-head gate is computed from the first 16 hidden dimensions of the block input:

```
gate    = 2 · sigmoid(W_gate · x[:16])
context = context · gate
```

With `W_gate = 0`, the gate is exactly 1, so the output passes through unchanged at step 0.

**Exclusive self-attention (XSA).** The component of the context that points along the token's own value vector is removed:

```
v_unit     = v / ‖v‖
correction = gate · (context · v_unit) · v_unit
context    = context − correction
```

The correction is zeroed at position 0, where causal attention can only see the token itself.

### Engram memory

Engram adds a hashed n-gram lookup at layers 1 and 8. Token sequences are normalised (NFKC, lowercasing, accent stripping) and mapped to 9,807 classes. Each class indexes a learned embedding table. The lookup table is stored in the checkpoint, and the tokenizer is checked against it at load time, so the model must be used with the tokenizer it was trained with.

### KV cache

Generation reuses the keys and values of earlier tokens. The prompt is processed once, and each new token costs one step rather than a full pass over the sequence. This includes Engram's n-gram history and the XSA position mask.

Verified against uncached decoding:
- Per-step logits agree to about 1e-6.
- Greedy `generate()` output is identical with and without the cache.

Within the 1,024-token window, cached decoding is exact. Past the window, the model re-encodes the last 1,024 tokens at each step, the same sliding-window behaviour as before. The cache is enabled by default in `generate()`; pass `use_cache=False` to turn it off.

---

## Training

| | |
|:---|:---|
| Training tokens | ~5B |
| Optimizer | MuonwithAdamw |
| Schedule | WSD |
| Steps | 19,000 |
| Batch | 320 |
| Precision | BF16 |

---

## Results

| | exp2 |
|:---|:---:|
| Validation perplexity | 13.6 |
| Validation accuracy | 47.1% |

### Benchmarks

Evaluated with lm-evaluation-harness, 0-shot, `acc_norm`.

| Benchmark | Score | Chance |
|:---|:---:|:---:|
| HellaSwag | 29.76% | 25% |
| ARC-Easy | 31.57% | 25% |
| ARC-Challenge | 23.98% | 25% |
| PIQA | 61.86% | 50% |
| ArithMark-3 | 38.30% | 25% |

**Intelligence Index: 12.41**

```
N(s, c) = 100 × (s − c) ÷ (100 − c)
Index = (N(HellaSwag, 25) + N(ARC, 25) + N(PIQA, 50) + 0.65 × N(ArithMark-3, 25)) ÷ 3.65
```

ARC is the mean of ARC-Easy and ARC-Challenge.

ArithMark-3 is scored with `acc_norm`. Its authors score by mean per-token log-likelihood, so this value may differ slightly from their leaderboard. Its standard error is about ±1.5 points.

Training used about 5B tokens. The AxiomicLabs reference models were trained on 15 to 75B tokens, so keep that gap in mind when comparing.

---

## Quick start

```bash
pip install torch transformers safetensors
```

The model and tokenizer load through `transformers` with remote code. Batched generation with padding is not supported, so generate one prompt at a time.

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

repo = "RuVI-AI/Pragya-EXP2"   # or a local folder
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).to("cpu").eval()

inputs = tok("The Internet is", return_tensors="pt")
out = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs["attention_mask"],
    max_new_tokens=60,
    do_sample=True,
    temperature=0.7,
    top_k=12,
    repetition_penalty=1.3,
    no_repeat_ngram_size=3,
    pad_token_id=tok.pad_token_id,
    top_p=0.95
)
print(tok.decode(out[0], skip_special_tokens=True))
```

The tokenizer prepends `<|bos|>` by default. The string `<|bos|>` inside a prompt is treated as ordinary text.

### Load output

Loading prints a `LOAD REPORT` that lists `model.lm_head.weight` as `MISSING`. This is expected: the checkpoint stores the embedding once, and the head is tied to it when the model loads. To check:

```python
print("tied:", model.model.transformer.wte.weight is model.model.lm_head.weight)  # True
```

On first load, Transformers also prints "A new version of the following files was downloaded". This is the remote code being fetched, not an error. Pin a `revision=` to stop the message once you have a version you trust.

---

## Sampling

These settings gave fluent, on-topic output in testing:

```text
temperature        0.7
top_k              12
repetition_penalty 1.3
no_repeat_ngram    3
top_p              0.95
```

Greedy decoding (`do_sample=False`) repeats phrases after a few sentences, so use sampling.

The model writes in a simple, children's-book register. This reflects the TinyStories-heavy data mix.

---

## Capabilities

**It can:** continue English prompts for a sentence or two, and produce plausible text in narrative and encyclopedic registers.

**It cannot:** answer questions reliably, follow instructions, stay coherent past the first paragraph, handle non-English text, or use context longer than 1024 tokens.

---

## Limitations

- **Small scale.** Output is grammatical for a sentence or two, then drifts into invented facts.
- **English only.** The tokenizer encodes other languages, but the model has no training signal for them.
- **No instruction tuning.** Any instruction-following is incidental.
- **No padded batches.** Attention masks are not supported, so padded batches are rejected.
- **No alignment or safety filtering.** The model can produce biased, false or toxic text.

---

## License

Apache 2.0. Each training dataset carries its own license. Check upstream terms before redistributing derivatives.

---

## Citation

```bibtex
@misc{pragya-exp2,
  title  = {Pragya exp2: a 28.5M from-scratch decoder-only language model with Exclusive Gated Attention and Engram memory},
  author = {Arush},
  year   = {2026},
  note   = {Experimental checkpoint, not for production use},
  howpublished = {\url{https://huggingface.co/RuVI-AI/Pragya-EXP2}}
}
```