DOCSEARCHbyAlgolia
Submit Docs

Artificial Hippocampus Networks

repository·main·Indexed 41 days ago

https://github.com/bytedance-seed/ahn

A framework for efficient long-context modeling that integrates lossless attention-based memory with compressed RNN-style memory. AHN supports multiple memory module implementations including GatedDeltaNet, DeltaNet, and Mamba2, and provides tools for merging weights with base models like Qwen2.5, training with custom datasets, and evaluating on benchmarks such as LongBench, LV-Eval, and InfiniteBench.

Tokens
16.6K
Snippets
39
Records
65
Agent score
49%

What's inside AHN

  1. How AHN works: Lossless vs. Compressed Memory

    main

    Artificial Hippocampus Networks (AHNs) are designed for efficient long-context modeling by combining two types of memory:

    1. Lossless Memory: Typically the attention mechanism's Key-Value (KV) cache. It stores exact input information but grows linearly with sequence length, making it computationally expensive for very long contexts.
    2. Compressed Memory: Typically an RNN-style hidden state. It maintains a constant size and fixed computational cost per token, but it suffers from information loss.

    The AHN Mechanism: AHNs continually convert information from the lossless memory (located outside the current sliding attention window) into a compact, compressed form. The model then uses both the lossless information within the sliding window and the compressed memory to make predictions, allowing for efficient long-context processing.

  2. Specify and mix multiple datasets during training

    main

    When running training scripts, you can specify multiple datasets and control how they are mixed using the following methods:

    1. Concatenation: Use a comma (,) within the --dataset flag to list multiple datasets (e.g., chatqa2,my_dataset).
    2. Sampling Probabilities: Use the --interleave_probs flag to control the sampling ratio for each dataset. Provide a comma-separated list of floats corresponding to the order of datasets specified in --dataset (e.g., --interleave_probs 0.1,0.5).
    # Example: Training with two datasets and specific sampling probabilities
python train.py --dataset chatqa2,my_dataset --interleave_probs 0.1,0.9

  3. Merge AHN weights with a base model

    main

    AHN weights are often released as AHN-only checkpoints. To use them, you must merge the AHN weights into a base model (e.g., Qwen2.5) using the merge_weights.py script. This produces a single directory containing both the base and AHN parameters.

    # Base model (repo_id or local path)
BASE_MODEL=Qwen/Qwen2.5-3B-Instruct

# AHN-only weights (repo_id or local path)
AHN_PATH=ByteDance-Seed/AHN-GDN-for-Qwen-2.5-Instruct-3B

# Output directory for the merged model
MERGED_MODEL_PATH=./merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN

python ./examples/scripts/utils/merge_weights.py \
  --base-model $BASE_MODEL \
  --ahn-path $AHN_PATH \
  --output-path $MERGED_MODEL_PATH

  4. Run LV-Eval inference and evaluation together

    main

    To run the LV-Eval pipeline, execute the prediction script followed by the evaluation script.

    Note on Output Directory: The evaluation script expects an input directory following the pattern: ${OUTPUT_DIR}/${MODEL_NAME}_${METHOD}_as${NUM_ATTENTION_SINK}_sw${SLIDING_WINDOW}.

    cd eval/lveval

MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
METHOD=ahn
MAX_INPUT_LENGTH=256000
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=32640
OUTPUT_DIR=./eval_results

python3 pred.py \
  --model-name $MODEL_NAME \
  --model-path $MERGED_MODEL_PATH \
  --method $METHOD \
  --model-max-len $MAX_INPUT_LENGTH \
  --start_size $NUM_ATTENTION_SINK \
  --recent_size $SLIDING_WINDOW \
  --output-dir $OUTPUT_DIR

python3 eval.py \
  --input-dir ${OUTPUT_DIR}/${MODEL_NAME}_${METHOD}_as${NUM_ATTENTION_SINK}_sw${SLIDING_WINDOW}

  5. Run LongBench inference and evaluation together

    main

    You can run the full LongBench pipeline (inference and evaluation) using a single shell script.

    Required Environment Variables:

    • MODEL_NAME: Model identifier (used for naming output files).
    • MERGED_MODEL_PATH: Path to the weights of the base model and AHN.
    • METHOD: Inference method (e.g., ahn).
    • MAX_INPUT_LENGTH: Maximum input tokens.
    • NUM_ATTENTION_SINK: Number of attention sink tokens.
    • SLIDING_WINDOW: Sliding window size.
    • OUTPUT_DIR: Directory to save predictions.
    cd eval/longbench

MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
METHOD=ahn
MAX_INPUT_LENGTH=32000
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=8064
OUTPUT_DIR=./eval_results

bash eval.sh $MODEL_NAME $MERGED_MODEL_PATH $METHOD $MAX_INPUT_LENGTH $NUM_ATTENTION_SINK $SLIDING_WINDOW $OUTPUT_DIR

  6. Run LongBench inference and evaluation separately

    main

    If you need to decouple the processes, follow these steps:

    1. Inference

    Run pred.py with the following arguments:

    • --model_name: Model identifier.
    • --model_path: Path to weights.
    • --max_length: Maximum input tokens.
    • --method: Inference method.
    • --dataset: One of "dureader", "hotpotqa", "musique", "narrativeqa", "qmsum", or "triviaqa".
    • --attention_sink: Number of attention sink tokens.
    • --sliding_window: Sliding window size.
    • --results_dir: Output directory.

    2. Evaluation

    Run eval.py using the results generated by the inference step:

    • --model: Model identifier.
    • --results_path: Path to the directory containing results (typically ${OUTPUT_DIR}/${MODEL_NAME}_${MAX_INPUT_LENGTH}).
    # Inference
cd eval/longbench
MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
MAX_INPUT_LENGTH=32000
METHOD=ahn
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=8064
DATASET=hotpotqa
OUTPUT_DIR=./eval_results

CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name $MODEL_NAME \
  --model_path $MERGED_MODEL_PATH \
  --max_length $MAX_INPUT_LENGTH \
  --method $METHOD \
  --dataset $DATASET \
  --attention_sink $NUM_ATTENTION_SINK \
  --sliding_window $SLIDING_WINDOW \
  --results_dir $OUTPUT_DIR

# Evaluation
python eval.py \
  --model $MODEL_NAME \
  --results_path $OUTPUT_DIR/${MODEL_NAME}_${MAX_INPUT_LENGTH}

  9. Install AHN for training

    main

    To set up the AHN environment for training, ensure you are using Python 3.11, CUDA 12.4, and PyTorch 2.5.1+cu124. Follow these steps:

    1. Clone the repository and enter the directory.
    2. Install required forked libraries (flash-linear-attention and llamafactory).
    3. (Optional) If using AHN-Mamba2, install the forked Mamba version.
    4. Install AHN in editable mode with training extras.
    # 1. Clone the AHN repository and move into it
git clone https://github.com/ByteDance-Seed/AHN.git
cd AHN

# 2. Install required forked libraries
pip install "git+https://github.com/Seerkfang/flash-linear-attention.git@main#egg=flash-linear-attention"
pip install "git+https://github.com/Seerkfang/LLaMA-Factory.git@main#egg=llamafactory"

# (Optional) Install the forked Mamba version if you plan to use AHN-Mamba2
# MAMBA_FORCE_BUILD=TRUE pip install "git+https://github.com/yuweihao/mamba.git"

# 3. Install AHN in editable mode with training extras
pip install -e ".[train]"

  10. Run InfiniteBench inference and evaluation separately

    main

    To run InfiniteBench separately, use pred.py for inference and eval.py for evaluation.

    1. Inference

    Arguments for pred.py:

    • --model_name: Model identifier.
    • --model_path: Path to weights.
    • --max_length: Maximum input tokens.
    • --method: Inference method.
    • --dataset: Set to infinitebench.
    • --split: Either longbook_qa_eng or longbook_qa_chn.
    • --attention_sink: Number of attention sink tokens.
    • --sliding_window: Sliding window size.
    • --results_dir: Output directory.

    2. Evaluation

    Run eval.py with:

    • --model: Model identifier.
    • --results_path: Path to results (typically ${OUTPUT_DIR}/${MODEL_NAME}_${MAX_INPUT_LENGTH}).
    # Inference
cd eval/longbench
MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
MAX_INPUT_LENGTH=128000
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=32640
DATASET=infinitebench
SPLIT=longbook_qa_eng
OUTPUT_DIR=./eval_results

CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name $MODEL_NAME \
  --model_path $MERGED_MODEL_PATH \
  --max_length $MAX_INPUT_LENGTH \
  --method $ahn \
  --dataset $DATASET \
  --split $SPLIT \
  --attention_sink $NUM_ATTENTION_SINK \
  --sliding_window $SLIDING_WINDOW \
  --results_dir $OUTPUT_DIR

# Evaluation
python eval.py \
  --model $MODEL_NAME \
  --results_path $OUTPUT_DIR/${MODEL_NAME}_${MAX_INPUT_LENGTH}

  11. Configure AHN training arguments

    main

    When training AHN modules, use the following arguments to control the training process:

    • --model_name_or_path: The base model to use (e.g., Qwen/Qwen2.5-3B-Instruct).
    • --dataset: The training dataset (default: chatqa2).
    • --loss_type: The loss function to use. Options: kl, ce.
    • --ahn_implementation: The AHN memory module type. Options: GatedDeltaNet, DeltaNet, Mamba2.
    • --ahn_position: Where to insert AHN modules. Options: prefix, random.
    • --sliding_window_type: The sliding window strategy. Options: fixed, random.
    • --save_ahn_only: Boolean flag to determine if only AHN weights should be saved. Options: True, False.
    • --filter_len: Filter training data by sequence length (e.g., keep samples longer than 288).
    • --cutoff_len: Set the cutoff sequence length (e.g., 24576).

  12. Manage AHN cache with BaseAHN

    main

    The BaseAHN class provides a mechanism to manage the recurrent memory (gate values $\alpha$ and $\beta$) used in Artificial Hippocampus Networks.

    Key Methods:

    • update_cache(hidden_states): Extracts and appends new $\beta$ and $\alpha$ gate values from the current hidden_states using the underlying AHN function (self.fn).
    • query_cache(num_attn_sinks, num_cached_tokens): Retrieves specific slices of the $\beta$ and $\alpha$ cache, skipping a specified number of attention sinks.
    • trim_cache(num_attn_sinks, num_cached_tokens): Removes a segment of the cache from the middle (after the sinks and the specified cached tokens).
    • reset_cache(): Clears all cached gate values and resets the token count.

    Note: During forward calls, if the model is not in training mode, self.num_cached_tokens is automatically incremented by the sequence length of the output.

    # Conceptual usage of BaseAHN cache management
# Assuming 'model' is an instance of BaseAHN

# 1. Update cache with new states
model.update_cache(hidden_states)

# 2. Query specific parts of the cache
# num_attn_sinks: number of tokens to skip at the start
# num_cached_tokens: number of tokens to retrieve
gates = model.query_cache(num_attn_sinks=2, num_cached_tokens=10)

# 3. Reset for a new sequence
model.reset_cache()