AAIE-Distilled Dense
A ~354M-parameter dense transformer trained on general web text via knowledge distillation
against a Qwen/Qwen2.5-1.5B teacher. This is the raw base checkpoint — it has not seen any
instruction-tuning or task-specific fine-tuning, so it behaves like a classic base language
model: it continues text plausibly, but will not reliably follow instructions phrased as
questions or requests (e.g. "Provide feedback on the following...").
If you want a model that follows general instructions, see the companion
AAIE-Distilled Dense Instruct model (same base weights, further trained on SmolTalk). If you
want the assignment-feedback-specialized model, that's a separate export built on top of both
of these stages — see this project's TECHNICAL_REPORT.md.
Architecture
GQA attention (8 query heads / 2 KV heads) + RoPE + SwiGLU FFN + RMSNorm + tied embeddings,
20 layers, d_model=512, 354M parameters. Tokenizer: same as Qwen/Qwen2.5-0.5B (151,936 vocab).
Data
| Dataset | Size | Role |
|---|---|---|
HuggingFaceFW/fineweb-edu (sample-10BT subset) | ~70B tokens (100,000 steps x 131,072 tokens/step) | General web-text language modeling. |
Qwen/Qwen2.5-1.5B (frozen teacher) | — | Distillation signal: loss = 0.5 * CE(labels) + 0.5 * KD(teacher logits, T=2.0). |
Training settings
| Value | |
|---|---|
| Base | random initialization |
| Optimizer | AdamW, weight decay 0.1 |
| LR schedule | warmup 1,000 steps -> 3e-4 peak, cosine decay to 3e-5 |
| Batch | micro-batch 4 x grad-accum 32 x seq-len 1024 = 131,072 tokens/step |
| Steps | 100,000 |
| Grad clip | 1.0 |
| Sequence length | 1024 |
| Hardware | single A100 (40GB), SLURM cluster |
Usage
This is a base model — prompt it as a text continuation, not an instruction:
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "your-username/aaie-ddense-pretrain" # after pushing, see push_to_hub.py
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, device_map="auto")
inputs = tokenizer("The purpose of a database index is", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=60, use_cache=True, do_sample=True, temperature=0.8)
print(tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
trust_remote_code=True is required (custom architecture, see modeling_aaieddense.py).
use_cache=True (the default) enables real KV-caching for .generate() — see
GQAAttentionCached in modeling_aaieddense.py.
Known limitations
- Not instruction-tuned: asking it a question or giving it an instruction will often just produce a plausible-sounding continuation of the prompt text rather than an answer/response — this is expected base-model behavior, not a bug.
- Small model, moderate pretraining budget: 354M params on ~70B tokens is well below compute-optimal scale for models trained by major labs; expect weaker world knowledge and reasoning than similarly-sized production models trained on much larger corpora.
- Fluency and factual reliability were not evaluated for this checkpoint specifically — see AAIE-Distilled Dense Instruct and the task-specific model for the evaluations that were run (LLM-judge scoring on IT/CS assignment feedback).