| ---
|
| title: Titans Miras Demo
|
| emoji: 🔬
|
| colorFrom: blue
|
| colorTo: purple
|
| sdk: gradio
|
| sdk_version: 4.36.1
|
| app_file: app.py
|
| pinned: false
|
| ---
|
|
|
| # Titans + MIRAS: A Brain That Changes Itself While Thinking
|
|
|
| A minimal but faithful reimplementation of **Titans** (test-time learning) and **MIRAS** (associative memory framework) using open-source models on Hugging Face.
|
|
|
| ## What is this?
|
|
|
| This demo showcases a neural architecture that can **learn and update its memory while generating responses** - a brain that literally changes itself while thinking!
|
|
|
| ### Key Features
|
|
|
| - 🔄 **Test-time learning**: Memory updates during inference (not just training)
|
| - 🎯 **Retention gate**: Surprising/novel inputs are more memorable (inspired by human memory)
|
| - 💾 **Persistent memory**: State is saved across sessions
|
| - 🤖 **Fully OSS**: Uses distilgpt2 and runs entirely on Hugging Face
|
|
|
| ## Architecture
|
|
|
| ```
|
| User Input
|
| ↓
|
| [Base LM: distilgpt2] → Hidden States (768-dim)
|
| ↓
|
| [Key/Value Projections] → Memory Space (256-dim)
|
| ↓
|
| [MIRAS Memory Module] ← Test-time Gradient Updates
|
| ↓
|
| [Text Generation] → Response + Memory Stats
|
| ```
|
|
|
| ### Components
|
|
|
| 1. **Base Language Model**: distilgpt2 (frozen, no training)
|
| 2. **Projection Layers**: Map hidden states to memory space
|
| 3. **MIRAS Memory**: Associative memory with learnable key→value mapping
|
| 4. **Retention Gate**: Adjusts learning rate based on surprise (loss magnitude)
|
| 5. **Memory Store**: Persists memory state to disk
|
|
|
| ## How It Works
|
|
|
| 1. Input text is processed through distilgpt2
|
| 2. Last hidden state is projected to key/value pairs
|
| 3. Memory predicts value from key
|
| 4. Loss (prediction error) indicates surprise
|
| 5. Higher surprise → higher retention → faster learning
|
| 6. Memory updated via gradient descent (1e-3 base LR)
|
| 7. Response generated and memory saved
|
|
|
| ## References
|
|
|
| - **Titans**: [Learning to Memorize at Test Time](https://arxiv.org/abs/2501.00663)
|
| - **MIRAS**: [Framework for Associative Memory with Attentional Bias](https://arxiv.org/abs/2504.13173)
|
|
|
| ## Running Locally
|
|
|
| ```bash
|
| pip install -r requirements.txt
|
| python app.py
|
| ```
|
|
|
| Built with ❤️ exploring the future of adaptive AI systems.
|
|
|
