🧬 CellularFlow: Memory-Augmented Continual Learning LLM

CellularFlow is a memory-augmented neural architecture designed as a continual-learning alternative to standard Transformers. By replacing dense Feed-Forward Networks (FFN/MLP) with Multi-Head Associative DNA Memory Banks and an Episodic Memory Slot Buffer, CellularFlow decouples factual knowledge storage from sequence reasoning.

It achieves state-of-the-art catastrophic forgetting mitigation (83.9% retention across sequential domains) and enables zero-backprop streaming learning during inference.

---

🌟 Key Capabilities at a Glance

---

📐 Architecture Overview

CellularFlow fuses two computational pathways into a unified Hybrid CMC Layer:

                       Input Sequence: X (B, T, d)
                                   │
                    ┌──────────────┴──────────────┐
                    ▼                             ▼
       ┌─────────────────────────┐   ┌─────────────────────────┐
       │   Multi-Head DNA Memory │   │   Episodic Memory Slot  │
       │   Associative Banks     │   │   Buffer (Fast-Write)   │
       └────────────┬────────────┘   └────────────┬────────────┘
                    │                             │
                    └──────────────┬──────────────┘
                                   │ (Gated Memory Enrichment)
                                   ▼
       ┌───────────────────────────────────────────────────────┐
       │  Causal Multi-Head Self-Attention with RoPE (FlashAttn)│
       └───────────────────────────┬───────────────────────────┘
                                   │
                                   ▼
                       Output Sequence: Y (B, T, d)

1. Multi-Head DNA Memory (CMCLayer)

Each head ($H$) maintains learned key-value associative banks: $$\text{scores}_h = \frac{\text{Norm}(x_h) \cdot \text{Norm}(K_h)^T}{\tau_h} + \mathcal{N}(0, 0.1)$$ $$\text{weights}_h = \text{Softmax}(\text{Top-K}(\text{scores}_h))$$ $$\text{out}_h = \text{weights}_h \cdot V_h$$

• Specialized Subspaces: Heads specialize independently across syntax, semantics, and domain knowledge.
• Exploration Noise: Gaussian perturbation prevents dead memory slots during Top-K sparse routing.

2. The Three Continual-Learning Modes

• Mode 1 — Live Learning (inference_write=True): Updates DNA memory values on the fly during inference via Exponential Moving Average (EMA) with zero backward pass. Protected by Spherical Anisotropy Regularization to prevent vector collapse.
• Mode 2 — Selective Fine-Tuning (set_mode("selective")): Freezes ~85% of the backbone (projections, embeddings, LayerNorms) and trains only the DNA banks. Retains foundational knowledge while rapidly absorbing new domains.
• Mode 3 — Episodic Fact Injection (inject_fact): Writes facts into slot-based episodic memory with temporal age decay (exp(-0.005 * age)) and consolidates top facts into DNA banks post-epoch.

---

🚀 Quickstart

1. Installation

Requires Python $\ge$ 3.11 and PyTorch $\ge$ 2.4.0.

# Clone the repository
git clone https://github.com/celcilin/cellularflow.git
cd cellularflow

# Install dependencies using UV (recommended) or pip
pip install -e .

# For GPU acceleration (CUDA 12.4+):
pip install torch --index-url https://download.pytorch.org/whl/cu124

2. Python API Usage

import torch
from cellularflow import CellularFlowLM, CellularFlowTrainer, BPEDataset

# 1. Initialize tokenizer & dataset
dataset = BPEDataset("Alice was beginning to get very tired of sitting by her sister...", context_len=256)

# 2. Instantiate CellularFlow LM
model = CellularFlowLM(
    vocab_size   = dataset.vocab,
    dim          = 512,
    n_layers     = 6,
    n_heads      = 8,
    n_entries    = 128,
    context_len  = 256,
    use_episodic = True
)

# 3. Pretraining
trainer = CellularFlowTrainer(model, dataset, device="cuda" if torch.cuda.is_available() else "cpu")
trainer.pretrain(epochs=100, seed_dna=True)

# 4. Fast Generation (with KV-Cache)
prompt = "The journey into"
print(trainer.generate(prompt, max_new=100, temperature=0.8))

# 5. Continual Learning: Mode 3 Fact Injection
trainer.inject_fact("The hidden archives are kept inside Vault 42.")

# 6. Continual Learning: Mode 2 Selective Fine-Tuning (Backbone Frozen)
trainer.selective_finetune("Technical medical notes on neurology...", epochs=10)

# 7. Continual Learning: Mode 1 Live Streaming Learning (0 Backprop)
trainer.live_learn("Streaming log telemetry received in real time...")

---

🖥️ Interactive Web Dashboard & Tools

CellularFlow includes an interactive glassmorphic web dashboard for real-time inference, fact injection, and memory inspection:

# Start the FastAPI server
uvicorn server.app:app --host 0.0.0.0 --port 8000

Open http://localhost:8000 in your browser to interactively generate text, inspect layer-wise episodic slot utilization, and test live fact injections.

# Interactive CLI Playground
python analysis.py --checkpoint checkpoint/CMC_BaseModel.pt --interactive

---

📊 Benchmarks & Validated Results

1. Parameter Efficiency vs. Baseline Transformer

Evaluated on a standardized 62KB multi-domain corpus:

2. Catastrophic Forgetting (5-Domain Sequential Learning)

Trained sequentially across Literature, Science, History, Technical, and Poetry:

---

📂 Project Structure

cellularflow/
├── cellularflow/
│   ├── core.py               # CMCLayer, HybridCMCLayer, EpisodicMemory, CellularFlowLM
│   ├── trainer.py            # Pretraining, selective fine-tuning, live learning, mixed precision
│   ├── extensions.py         # Blockwise Attention, Compressed KV (CKV), MTP, Beaconing
│   ├── corpus.py             # Multi-domain benchmark corpora
│   └── swarm.py              # DNASwarm evolutionary optimizer
├── benchmarks/
│   └── evaluate_checkpoint.py# Evaluation harness for perplexity, accuracy, and memory norms
├── server/
│   └── app.py                # FastAPI server + WebSocket endpoint
├── dashboard/
│   ├── index.html            # Web dashboard UI
│   ├── app.js                # Frontend WebSocket and API client
│   └── styles.css            # Dark glassmorphic design system
├── sft/
│   ├── sft_dataset.py        # ChatML templates and target loss masking
│   └── sft_trainer.py        # Supervised fine-tuning curriculum engine
├── scripts/
│   ├── train_tokenizer.py    # ByteLevelBPE tokenizer builder
│   └── test_extensions.py    # Architecture extension verification
├── analysis.py               # CLI exploration & interactive REPL
├── pyproject.toml            # Project build & dependency definitions
└── CONTRIBUTING.md           # Contribution guidelines & developer standards

---

🤝 Contributing

We welcome contributions from researchers, engineers, and developers worldwide! Please review CONTRIBUTING.md for instructions on setting up your environment, adhering to XLA/TPU graph rules, and submitting pull requests.

---

👤 Author

Celcilin C S

• GitHub: @celcilin
• Email: celcilin204@gmail.com

---

📜 License

This project is licensed under the MIT License — see the LICENSE file for details.