Skip to content

Repository files navigation

continual-learning-poc

A small, from-scratch decoder-only transformer in Rust (candle), built to study one question: if a trained model is warm-started into a new training run instead of being retrained from scratch, how much of what it already knew survives — and which mitigation strategy preserves the most of it?

This is a research sandbox, not a product. It trains a tiny transformer on a synthetic names-facts corpus, saves it as real model artifacts (.safetensors + config.json + tokenizer.json), then continues training it on a second, disjoint corpus while measuring catastrophic forgetting against several mitigation techniques.


Table of Contents


The Research Question

Stage 1 trains a model on names-facts and saves a real checkpoint. Stage 2 loads that checkpoint and continues training on a different corpus. How much of Stage 1's knowledge is still there afterward — and does that change depending on whether Stage 2 uses plain fine-tuning, replay, partial freezing, a LoRA adapter, or an EWC-style penalty?

See ARCHITECTURE.md for the full technical design and ROADMAP.md for the phased build plan.


Quickstart

git clone https://github.com/AarambhDevHub/continual-learning-poc.git
cd continual-learning-poc
cargo build --release
# Step 0 — generate the synthetic Stage 1 dataset (names-facts corpus + probes)
cargo run --release -- generate-data

# Stage 1 — train from scratch on the names-facts corpus, save a real checkpoint
cargo run --release -- train-stage1 --config nano --steps 2000

# Baseline eval — record Stage 1 probe accuracy before any continual training
cargo run --release -- eval --checkpoint checkpoints/stage1 \
    --probes data/stage1_forgetting_probes.jsonl

Stage 2 commands (train-stage2, report) are planned for Phase 6+ and are not yet implemented. See ROADMAP.md for the build plan.

No GPU required — the whole loop is sized to run comfortably on an 8 GB CPU-only laptop in well under 30 minutes per experiment.


Project Structure

continual-learning-poc/
├── src/
│   ├── model/              # decoder-only transformer (Phase 3 ✅)
│   │   ├── embedding.rs    #   TokenPositionEmbedding
│   │   ├── attention.rs    #   CausalSelfAttention (MHA + causal mask)
│   │   ├── ffn.rs          #   SwiGluFfn
│   │   ├── block.rs        #   TransformerBlock (pre-norm + residual)
│   │   └── transformer.rs  #   PocModel (full pipeline)
│   ├── dataset/            # synthetic names-facts corpus + probe generator (Phase 2 ✅)
│   ├── train.rs            # AdamW training loop + Trainer struct (Phase 4 ✅)
│   ├── eval.rs             # probe accuracy + perplexity harness (Phase 5 ✅)
│   └── checkpoint.rs       # save/load .safetensors + config + tokenizer (Phase 4 ✅)
├── data/                   # generated .jsonl corpora and probes (gitignored)
└── checkpoints/            # saved model.safetensors + config.json + tokenizer.json (gitignored)

The mitigation/ directory (replay, freeze, LoRA, EWC-lite) is planned for Phase 8–11 and does not exist yet.

Full breakdown in ARCHITECTURE.md.


Mitigation Strategies Compared

Strategy Idea
Naive (baseline) Warm-start, then fine-tune on Stage 2 with no protection
Replay Mix Stage 1 data back into Stage 2 batches
Partial freeze Freeze early layers, train only late layers + head
LoRA adapter Freeze the base model entirely, train a small adapter
EWC-lite Penalize movement of parameters that mattered most in Stage 1

Details and expected tradeoffs in ARCHITECTURE.md §10.


Documentation

  • ARCHITECTURE.md — model design, checkpoint format, dataset design, data flow
  • ROADMAP.md — phased build plan with goals, tasks, tests, and milestones
  • CHANGELOG.md — release history

Status

Active research project. Phase 5 (Evaluation Harness) is completeeval.rs loads any checkpoint, scores cloze probes with exact first-token match, and computes perplexity. The Stage 1 baseline is recorded: probe_accuracy ≈ 37–44%, perplexity ≈ 7.5–14.5 (varies with the seeded BPE/dataset split). 38/38 tests pass.

Phase Description Status
0 Project setup + core types ✅ Done
1 BPE tokenizer ✅ Done
2 Dataset generator (names-facts + probes) ✅ Done
3 Model architecture ✅ Done
4 Stage 1 training + checkpoint save ✅ Done
5 Evaluation harness Done
6 Stage 2 corpus + naive continual training 🔜 Next
7 Forgetting measurement ⬜ Planned
8 Mitigation: replay buffer ⬜ Planned
9 Mitigation: partial freeze ⬜ Planned
10 Mitigation: LoRA/DoRA adapter ⬜ Planned
11 Mitigation: EWC-lite (stretch) ⬜ Planned
12 Comparison report ⬜ Planned

Follow ROADMAP.md for per-phase task checklists and milestones.

Contributing

Contributions, issue reports, and forgetting-mitigation ideas are welcome. See CONTRIBUTING.md for setup instructions and coding conventions.

License

Licensed under the MIT License.


Support This Project

This project is developed and maintained independently, in the open, alongside the rest of the Aarambh Dev Hub projects. If it's useful to you and you'd like to support continued development:

Thanks for checking this project out — happy coding!

About

No description, website, or topics provided.

Resources

Code of conduct

Contributing

Security policy

Stars

1 star

Watchers

0 watching

Forks

Releases

Sponsor this project

Packages

Contributors

Languages