LongMem

Environment Setup

• torch: Please follow torch official installation guide. We recommend torch>=1.8.0. Please select the torch-gpu version which is consistent with your cuda driver version.
• Faiss-GPU: For Nvidia V100 GPUs, simply install via pip install faiss-gpu. For Nvidia A100, A6000 GPUs, please run conda install faiss-gpu cudatoolkit=11.0 -c pytorch. The A100 GPU is not officially supported by faiss-gpu, sometimes it will lead to errors, you can refer to this git issue of faiss for help.
• fairseq: pip install --editable ./fairseq Then the revised fairseq and dependency packages will be installed. We strongly recommend you to use python 3.8 for stability.
• other packages: pip install -r requirements.txt

Project Structure

• Pre-trained LLM Class (L24, E1024, Alibi positional embedding): fairseq/fairseq/models/newgpt.py
• Transformer Decoder with SideNetwork (L12, E1024): fairseq/fairseq/models/sidenet/transformer_decoder_sidenet.py
• Transformer Language Model with SideNetwork Class: fairseq/fairseq/models/transformer_lm_sidenet.py
• Memory Bank and Retrieval: fairseq/fairseq/modules/dynamic_memory_with_chunk.py
• Joint Attention for Memory Fusion: fairseq/fairseq/modules/joint_multihead_attention_sum.py

Memory-Augmented Adaptation Training

Data collection and Preprocessing

Evaluation

Memory-Augmented In-Context Learning

# Evaluate gpt2 baseline
python eval_scripts/eval_longmem_icl.py --path /path/to/gpt2_pretrained_model
# Evaluate LongMem model
python eval_scripts/eval_longmem_icl.py --path /path/to/longmem_model --pretrained-model-path /path/to/gpt2_pretrained_model

Credits