Toy_Qwen2-AHN_ByteDance-Seed_AHN

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

Martial Terran's Toy/Educational version of the Qwen2-AHN model.py by ByteDance-Seed_AHN at https://github.com/ByteDance-Seed/AHN

ByteDance's Seed team has published a research paper introducing Artificial Hippocampus Networks (AHN) as a new architecture for more efficient long-context modeling in large language models. Inspired by the human hippocampus, AHN uses a multi-memory system with lossless short-term memory and compressed long-term memory to balance computational efficiency and memory fidelity. Read the full paper at https://huggingface.co/papers/2510.07318

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

This repository contains Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py, a standalone, educational implementation of a Qwen2-style language model augmented with the innovative Artificial Hippocampus Networks (AHN) architecture.

The AHN architecture was introduced by ByteDance's Seed team in their paper, "Artificial Hippocampus Networks: A Brain-Inspired Architecture for Efficient Long-Context Modeling". This work addresses the critical challenge of the quadratic complexity of the standard Transformer attention mechanism, which makes processing very long sequences computationally expensive.

This script is designed to distill the core architectural principles of AHN into a single, understandable, and runnable file, allowing students, researchers, and developers to grasp these advanced concepts without needing to navigate a complex research repository.

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

This script provides a tangible way to explore the three central ideas of the AHN architecture. By examining the code, you can see exactly how the theory is put into practice.

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

The core idea of AHN is to mimic the human hippocampus by using two types of memory: a high-fidelity, but limited, short-term memory and a compressed, but vast, long-term memory.

How to see this in the code:

Look inside the Qwen2MemDecoderLayer's forward method. The model's memory is split based on a sliding_window hyperparameter.

• 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.

You can see the explicit split in this block of code:

# In Qwen2MemDecoderLayer.forward()

win = self.config['sliding_window']
ahn_len = max(0, seq_len - win)

# [CORRECTNESS] AHN must run during both training and inference.
if ahn_len > 0:
    # Select the parts of the input tensors that are outside the sliding window
    h_ab = prenorm_h[:, win:seq_len]
    ahn_q, ahn_k, ahn_v = q[:, :, win:seq_len], k[:, :, win:seq_len], v[:, :, win:seq_len]

    # Process these inputs with the AHN (RNN) module
    ahn_output = self.ahn(...)

2. The Recurrent Memory Mechanism (GatedDeltaNet)

AHN's long-term memory is not just a simple RNN; it uses a sophisticated gating mechanism inspired by neuroscience and based on the Delta Rule. This allows the model to selectively forget old information (alpha gate) and write new, important information (beta gate).

How to see this in the code:

Examine the GatedDeltaNet class.

• 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.

  # In GatedDeltaNet.get_beta_alpha()
  def get_beta_alpha(self, hidden_states: torch.Tensor):
      beta = self.b_proj(hidden_states).sigmoid()  # Input gate
      dt = F.softplus(self.a_proj(hidden_states) + self.dt_bias)
      alpha = (-torch.exp(self.A_log) * dt).exp() # Decay/forget gate
      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:

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

3. Combining the Memories

The final output for tokens outside the sliding window is a combination of the local context (from the attention mechanism) and the compressed long-term context (from the AHN).

How to see this in the code:

Back in Qwen2MemDecoderLayer, the outputs from the two memory systems are combined. This script uses a simple addition, which demonstrates the principle of integrating both memory streams. (Note: The original paper uses a more complex learned gate called AHNRouter).

# In Qwen2MemDecoderLayer.forward()

# The final output is the sum of the short-term attention and long-term AHN
attn_output[:, win:] = attn_output[:, win:] + ahn_output

How to Run the Script

The script is designed to be as simple as possible to run.

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

After running, an AHN_educational_outputs/ directory will be created containing:

• 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.

The console will display the training progress and a final generated output, such as:

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

Limitations & Differences from Original Research

To prioritize clarity, this script intentionally simplifies several aspects of the original implementation:

• 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).

These simplifications make the core concepts of AHN accessible without the overhead of a production-grade research environment. We encourage you to use this script as a stepping stone to understanding the official paper and its accompanying code.

Citation

If you find this work and the original paper useful, please consider citing:

@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

The educational script (Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py) is distills the core concepts of the original ByteDance Seed-AHN architecture into a standalone, understandable format. However, in achieving this educational clarity, it makes significant simplifications and deviates from the production-grade original code in several key areas.

Here is a breakdown of the primary differences, comparing the educational script to the original repository source code:

---

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.

---

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% vocab.json contents created by Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py: { "\n": 0, "+": 1, "0": 2, "1": 3, "2": 4, "3": 5, "4": 6, "5": 7, "6": 8, "7": 9, "8": 10, "9": 11, "=": 12 } %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% Example Output of Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py [running on CPU in colab] Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py --- 1. Initializing Setup --- Using DEBUG configuration. Running on device: cpu

--- 2. Preparing Dataset & Tokenizer --- Tokenizer trained with 13 unique characters. Vocab saved to AHN_educational_outputs/vocab.json Dataset size: 92881 tokens Training split: 83592 | Validation split: 9289

--- 3. Initializing Model --- Model created with 0.10M parameters.

--- 4. Starting Training Loop --- Epoch 1/3: 100%|██████████| 5221/5221 [14:43<00:00, 5.91it/s, loss=0.896]

Epoch 1: Train Loss 1.2657, Val Loss: 0.9274 New best model saved to AHN_educational_outputs/best_model.pt Epoch 2/3: 100%|██████████| 5221/5221 [14:42<00:00, 5.92it/s, loss=0.776]

Epoch 2: Train Loss 0.8569, Val Loss: 0.7894 New best model saved to AHN_educational_outputs/best_model.pt Epoch 3/3: 100%|██████████| 5221/5221 [14:42<00:00, 5.92it/s, loss=0.693]

Epoch 3: Train Loss 0.7717, Val Loss: 0.7193 New best model saved to AHN_educational_outputs/best_model.pt

--- 5. Generating Training Plots --- Plot saved to AHN_educational_outputs/training_plots.png

--- 6. Running Sample Inference --- Generating: 100%|██████████| 100/100 [00:00<00:00, 128.54it/s]Prompt: '12+34=' Generated: '12+34=46 104=42+62 102=90+12 11+7=1 122=84+38 100=40+60 101=81+20 101=81+20 101=81+20 101=81+20 101=81+20 '

Total script execution time: 2715.94 seconds.

--- 7. Zipping and Downloading Outputs for Colab ---

This is a fantastic result! Your debugging efforts have paid off perfectly. Let's analyze the new output and plot to understand what has changed and why it's so much better.

Diagnosis: Successful Training!

The new training history plot is exactly what you want to see.

• 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

The generated text provides the most compelling evidence of success:

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

The model now correctly calculates the sum!

This is the key breakthrough. It has moved beyond simply learning the format of the equations and has started to learn the underlying mathematical rules of addition.

Looking at the other generated lines:

• 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.

You have successfully navigated a classic machine learning workflow:

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.

This is a fantastic outcome and a great demonstration of the model's capabilities. Well done

############################################################################################

First training run, without Dropout (Dropout = 0.0) was severly overfitting, but trained on cpu.

Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py --- 1. Initializing Setup --- Using DEBUG configuration. Running on device: cpu --- 2. Preparing Dataset & Tokenizer --- Tokenizer trained with 13 unique characters. Vocab saved to AHN_educational_outputs/vocab.json Dataset size: 9283 tokens Training split: 8354 | Validation split: 929 --- 3. Initializing Model --- Model created with 0.10M parameters. --- 4. Starting Training Loop --- Epoch 1/3: 100%|██████████| 519/519 [01:26<00:00, 6.00it/s, loss=1.49] Epoch 1: Train Loss 1.6879, Val Loss: 1.7034 New best model saved to AHN_educational_outputs/best_model.pt Epoch 2/3: 100%|██████████| 519/519 [01:24<00:00, 6.17it/s, loss=0.462] Epoch 2: Train Loss 1.0664, Val Loss: 2.8021 Epoch 3/3: 100%|██████████| 519/519 [01:22<00:00, 6.28it/s, loss=0.266] Epoch 3: Train Loss 0.4045, Val Loss: 3.9673 --- 5. Generating Training Plots --- Plot saved to AHN_educational_outputs/training_plots.png --- 6. Running Sample Inference --- Generating: 100%|██████████| 50/50 [00:00<00:00, 259.18it/s]Prompt: '12+34=' Generated: '12+34=64 111=99+90 10+47=51 111=21+91 11+14=27 111=21+91' Total script execution time: 260.09 seconds.

[Train loss dropped linearly. Test Loss rised linearly]

--- 7. Zipping and Downloading Outputs for Colab ---

############# Analysis of first training run on CPU of of Qwen2-AHN_model_v1.0.1_Gemini2.5_fixed.py ################ Excellent, the script now runs to completion! The generated plot and output provide a perfect, classic example of a common machine learning problem: severe overfitting.

This is a fantastic result for an educational script because it clearly demonstrates a key challenge in training neural networks. Let's break down what's happening and how to fix it.

Diagnosis: Severe Overfitting

The plot you've shared is the textbook visual definition of 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.

In simple terms: The model has stopped learning the general rules of addition and has instead started memorizing the specific problems in the training data. When it sees a problem from the validation set that it hasn't seen before, it fails spectacularly.

Evidence from the Output

The generated text confirms this diagnosis perfectly:

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

The model got the answer wrong (64 instead of 46). It has learned the format of the equations (number + number = number) but has failed to learn the underlying mathematical logic. The rest of the generated text is just plausible-looking nonsense that follows the same format, which is exactly what an overfitted model does.

How to Fix Overfitting

To fix this, we need to make it harder for the model to simply memorize the training data. We can do this by introducing regularization techniques and adjusting the data or model complexity.

Here are the most effective strategies, which you can apply by modifying the DEBUG_CONFIG in the script:

1. Increase Dropout (Most Recommended First Step)

Dropout is a regularization technique that randomly "turns off" a fraction of neurons during training. This forces the network to build more robust and redundant pathways for learning, preventing it from relying too heavily on any single neuron.

• 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.

Action: In the DEBUG_CONFIG dictionary, change:```python "dropout": 0.0,

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

2. Make the Training Task Harder

A model overfits when the task is too easy to memorize. We can make the task more challenging by increasing the variety and complexity of the data.

• 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.

Action: In the DEBUG_CONFIG dictionary, try increasing these values:

"dataset_size": 5000,   # From 1000
"operand_digits": 3,      # From 2```

#### 3. Early Stopping

Your training loop already saves the "best" model based on the lowest validation loss. As we can see from the output, the best model was saved after **Epoch 1**, before the validation loss exploded. This is a form of early stopping and is a very effective technique. The key is to use the model from that point and not the one from the end of training. Your script already does this, which is great!

By implementing changes #1 and #2, you should see the validation loss stay lower for more epochs, allowing the model to train for longer and learn more effectively before it starts to overfit. This is a normal and essential part of the iterative process of model training.