MAGS: Memory Augmented Graph Scaling

Enabling LLMs to learn

Major Note: Since we pivoted our idea during Week 4 / 5, we have ~2 weeks less progress, and do not have the code/slide/docs from Week 3 / 4 in our archive. However, we tried our best to still have our finished prototype by the end of the semester!

Team

• Pulkith Paruchuri
• Stanley Yu
• Pragya Singh
• Shailesh Kumar

Summary

A full report can be found at docs/report.pdf and docs/Final_Presentation.pdf.

Current LLMs struggle with learning and long-term retention, relying on inefficient and ephemeral context windows that limit adaptability and scalability. MAGS (Memory Augmented Generative System) proposes a neuroscience-inspired architecture using Liquid Knowledge Graphs (LKGs) and dual memory blocks—episodic (eMB) and semantic (sMB)—to emulate human memory and enable continuous, test-time learning. The system dynamically updates memory through a reinforcement-learning-trained model and Hebbian plasticity-inspired methods, managing graph growth to ensure efficiency. We tested the architecture on a combination of text-based and logic games to measure accuracy, adaptability, and explainability. We found promising results that match or exceed state-of-the-art approaches in many tasks.

Structure

• _archive: Previous work and documentation
• src/: Python source code and modules
• notebooks/: Final demo notebooks (inference the src/ modules)
• docs/: Final version of Documentation
• docs/report.pdf: Final report
• docs/Final_Presentation.pdf: Final presentation AI Model Structure (Conceptual)

├── _archive/                    # Previous work and documentation
├── docs/                        # Previous work and documentation
│   ├── report.pdf               # Final Paper
│   └── Final_Presentation.pdf   # Final slides
├── src/                         # Python source code and modules
│   ├── RL                       # GRPO and PPO Pipelines
│   ├── v2                       # Main Test Pipelines
│   └── demo_site                # Main interactive site
└── notebooks/                   # Final demo notebooks (inferences the `src/` modules)

Setup

• Create a virtual environment with python 3.12 (required)
• Install dependencies: pip install -r requirements.txt
• All code was run with a CPU and Metal GPU. It should work with any GPU and should be automatically detected
• Run the GRPO Pipeline in src/RL/GRPO.ipynb. Once the model finishes, download the GGUF to v2/mags/inference/models/ with the name gemma-3-finetuned.gguf. The current setup should finish in about 20 mins (although the accuracy for these hyperparamaters are not great). Alternatively, install the files from here link and follow the same instructions. Sorry! It was to big to push to github.
• Run src/mags_test_inference.py to make sure your setup succeeded (all test cases should pass like below)

• Install VLLM with your correct drivers here link
• Run the interactive terminal version at src\v2\gen_pipeline.py. If your VLLM or model is not configured, it falls back to Generic (Cloud) LLM (Not Fine-tuned instead prompted, in which case you need to pass your API key)

Sorry for not using Colab. It was very frustrating, with the sessions constantly terminating (and me losing all my checkpoints or training runs), or running out of memory (and even disk storage?) and crashing randomly. It also wasn't must faster than an intel mac with metal GPU optimization. I also worked all semester doing regular python library-type code, so it wasn't easy to port it over to a notebook, and I felt like the current terminal experience is pretty good too!!

Also please check out the code, I think it's pretty modular and scalable :)

Demo

• Insall and run the 'Go Live' feature in VSCode to run the demo on src/demo_site/index3.html. This uses the lightweight prompted version for the router so you'll need your API key.