Learn Artificial Hippocampus Networks (AHN) with a Pure PyTorch Qwen2 Model

Learn Artificial Hippocampus Networks (AHN) with a Pure PyTorch Qwen2 Model

Key Features of this Educational Implementation

• Standalone & Self-Contained: A single Python script with no external dependencies beyond PyTorch, Matplotlib, and TQDM.
• Pure PyTorch: All components, including the attention mechanism and the AHN recurrence, are implemented from scratch in pure PyTorch, making the code easy to read and debug.
• Minimal Dependencies: No custom CUDA kernels or complex libraries are required. It runs on a standard CPU or a free-tier Google Colab T4 GPU.
• Educational Focus: The code is heavily commented and structured to prioritize clarity and understanding over production-level performance.
• End-to-End Example: Includes data generation, tokenization, training, evaluation, and inference, providing a complete, working example of an AHN-enhanced model.

Understanding the AHN Innovation through this Script

1. The Multi-Memory System: Short-Term + Long-Term

• Short-Term Memory (Lossless): For tokens within the sliding window, the standard self.self_attn (a Qwen2Attention module) is used. This is the equivalent of the high-fidelity, quadratically expensive short-term memory.
• Long-Term Memory (Compressed): For tokens beyond the sliding window, the output is augmented by the self.ahn module (GatedDeltaNet). This module processes information sequentially and compresses it into a fixed-size state, acting as the efficient long-term memory.

2. The Recurrent Memory Mechanism (GatedDeltaNet)

• Gate Computation: The get_beta_alpha method shows how the input-dependent alpha (decay/forget gate) and beta (input gate) are calculated directly from the model's hidden states.
  python

  1# In GatedDeltaNet.get_beta_alpha() 2def get_beta_alpha(self, hidden_states: torch.Tensor): 3 beta = self.b_proj(hidden_states).sigmoid() # Input gate 4 dt = F.softplus(self.a_proj(hidden_states) + self.dt_bias) 5 alpha = (-torch.exp(self.A_log) * dt).exp() # Decay/forget gate 6 return beta, alpha
• The Recurrence Formula: The forward method of GatedDeltaNet contains the core state-update loop. This is the mathematical heart of the AHN, where the long-term memory state is updated at each timestep. This script's vectorized loop makes the operation clear:
  python

  1# In GatedDeltaNet.forward() 2# state is the compressed long-term memory 3# kv_product is the new information from the current token (Key * Value) 4 5state = alpha_t * state + beta_t * kv_product

3. Combining the Memories

How to Run the Script

1. Copy the text contents of Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.p script into a Google Colab code cell:
   Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py -> colab.
2. Install dependencies:
   pip install torch matplotlib tqdm
3. Run the script:
   python Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py [meaning, click the triangle in colab code cell]
   The script will automatically generate a synthetic dataset and vocab.json, and train the tokenizer, train the model, save plots, and run a sample generation. It is optimized to run smoothly in a Google Colab notebook with a CPU or a T4 GPU.

What to Expect

• best_model.pt: The trained model weights.
• vocab.json: The vocabulary for the character-level tokenizer.
• training_plots.png: A plot of the training and validation loss over epochs.

Prompt: '12+34='
Generated: '12+34=46'

Limitations & Differences from Original Research

• Dataset: Uses a tiny, synthetic arithmetic dataset instead of a massive web-scale text corpus.
• Performance: Implements the AHN recurrence in a pure Python loop, which is much slower than the custom CUDA kernels used in the original paper.
• Architecture: Uses simple addition to combine memories instead of the paper's AHNRouter, and uses standard attention instead of FlashAttention.
• Scale: Trains a small model on a single GPU instead of a large model with distributed training (DeepSpeed).

Citation

@misc{anonymous2025artificial,
      title={Artificial Hippocampus Networks: A Brain-Inspired Architecture for Efficient Long-Context Modeling},
      author={Anonymous},
      year={2025},
      eprint={2510.07318},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

Analysis: Deviations of the Educational Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py Script from the Original Model

A. Project Structure and Dependencies

• Educational Script: A single, monolithic Python file with minimal dependencies (torch, matplotlib, tqdm). This is excellent for portability and learning but is not scalable.
• Original Code: A full-fledged research repository. It relies heavily on the Hugging Face transformers ecosystem, which it monkey-patches and extends. It uses a complex directory structure (src/, eval/, examples/) and has a deep dependency stack including accelerate, deepspeed for distributed training, wandb for logging, and specialized libraries like fla (Flash Linear Attention) and mamba_ssm for optimized CUDA kernels.

B. Data Handling and Tokenization

• Educational Script: Implements a CharTokenizer from scratch and trains it on a synthetic arithmetic dataset (12+34=46). This is a toy problem designed to demonstrate learning on a small scale. Data is loaded via a basic PyTorch Dataset and DataLoader.
• Original Code: Uses the official Qwen2Tokenizer from transformers, a sophisticated Byte-Pair Encoding (BPE) tokenizer trained on a massive multilingual corpus. The data loading is abstracted away by the llama-factory training framework, designed to handle terabyte-scale datasets.

C. Core Model Architecture & Implementation

1. Base Model Foundation:
   • Educational Script: Re-implements Qwen2 architectural components (RMSNorm, MLP, Attention) from scratch in pure PyTorch.
   • Original Code: Inherits directly from transformers.Qwen2Model and transformers.Qwen3Model, modifying their layers (Qwen2MemDecoderLayer, Qwen3MemDecoderLayer) to insert the AHN mechanism. This is a much more robust and feature-complete approach.
2. Attention Mechanism:
   • Educational Script: Uses a standard PyTorch implementation of scaled-dot-product attention.
   • Original Code: Leverages highly optimized attention kernels like FlashAttention and FlexAttention. This is a critical difference for performance and memory efficiency, especially for handling the long sequences that AHN is designed for.
3. AHN (RNN) Implementation:
   • Educational Script: The GatedDeltaNet is implemented as a Python for loop. While vectorized, this is computationally orders of magnitude slower than a compiled CUDA kernel. It correctly captures the logic of the recurrence but not the performance.
   • Original Code: The GatedDeltaNet and other RNN modules (DeltaNet, Mamba2) are thin wrappers around custom, high-performance CUDA kernels (chunk_gated_delta_rule, mamba_chunk_scan_combined, etc.) imported from libraries like fla and mamba_ssm.
4. Architectural Components Missing in the Educational Script:
   • AHNRouter: The original code includes a sophisticated gating mechanism (AHNRouter) to dynamically combine the outputs of the standard attention (h_local) and the AHN memory (h_mem). The educational script uses a simple addition (attn_output[:, win:] + ahn_output).
   • Complex Loss Functions: The original supports multiple loss functions for training, including Kullback-Leibler (kl) divergence and a distillation loss (distill) to force the AHN's output to match the full attention's output. The educational script only uses a standard cross-entropy loss.
   • Dynamic Sliding Window: The original can vary the sliding window size during training (e.g., 'linear' or 'random'), a technique for improving generalization. The educational version uses a fixed sliding_window.
   • Gradient Checkpointing: The original code is deeply integrated with gradient checkpointing to reduce memory usage during training on very long sequences, which is absent in the educational script.

D. Training and Inference Infrastructure

• Educational Script: A basic, single-GPU training loop with a simple evaluate function. Generation is a simple autoregressive loop.
• Original Code: Built for large-scale, distributed training using DeepSpeed (supporting ZeRO stages 0, 2, and 3 for memory optimization). The entire training process is managed by llama-factory. Inference is also highly optimized, capable of handling KV caching and the custom AHN state caching simultaneously.

Diagnosis: Successful Training!

• Converging Losses: Both the Training Loss (blue) and Validation Loss (orange) are decreasing steadily and moving together. This is the hallmark of a healthy training process. It shows the model is not just memorizing the training data but is learning generalizable patterns that also apply to the unseen validation data.
• No Overfitting: Unlike the previous run, the validation loss is not increasing. This confirms that the regularization techniques we introduced (increasing dropout and dataset size) were highly effective at preventing the model from simply memorizing the answers.

Evidence from the Output

• Prompt: '12+34='
• Generated: '12+34=46'

• 104=42+62 (Correct)
• 100=40+60 (Correct)
• 11+7=1 (Incorrect) - This is perfectly normal. The model isn't perfect yet, but its high success rate shows it is genuinely learning.
• The repeating 101=81+20 at the end is a common behavior in simple text generation (called "getting stuck in a loop"). This can be fixed with more advanced sampling methods, but for this educational exercise, the primary goal has been achieved.

Summary of What Worked

1. Increased Dropout: Adding dropout: 0.1 forced the model to learn more robust features instead of relying on specific connections, making it harder to memorize.
2. Larger, More Complex Dataset: Increasing the dataset_size and operand_digits gave the model more varied examples to learn from, pushing it to find general rules rather than memorizing a small set of problems.

1. Initial Training: Encountered a critical bug.
2. Debugging: Fixed the bug and got the model to run.
3. First Results: Diagnosed severe overfitting from the loss curves.
4. Regularization: Applied techniques to combat overfitting.
5. Final Results: Achieved successful training and generalization.

Diagnosis: Severe Overfitting

• Training Loss (Blue Line): This line goes down steadily. This means the model is getting better and better at predicting the next character on the data it is being trained on. It's successfully memorizing the training examples.
• Validation Loss (Orange Line): This line starts low but then increases dramatically. This is the crucial part. It shows that while the model is memorizing the training set, its ability to handle new, unseen data (the validation set) is getting worse.

Evidence from the Output

• Prompt: '12+34='
• Generated: '12+34=64...'

How to Fix Overfitting

1. Increase Dropout (Most Recommended First Step)

• Problem: dropout is set to 0.0 in the debug config, meaning it's completely disabled.
• Solution: Change this to a small non-zero value.

to:
```python
"dropout": 0.1,  # A good starting point

2. Make the Training Task Harder

• Problem: The debug dataset is very small (1000 samples) with simple 2-digit numbers.
• Solution: Increase the dataset size and the difficulty of the problems.