BLEND: Behavior-guided Neural Population Dynamics Modeling via Privileged Knowledge Distillation
Zhengrui Guo, Fangxu Zhou, Wei Wu, Qichen Sun, Lishuang Feng, Jinzhuo Wang*, and Hao Chen*

[ArXiv]

• 2025.04.22: The model and training/evaluation codes have been released!
• 2025.01.23: Our paper is accepted by ICLR 2025! 🎉

This codebase is based on Neural-Data-Transformer and nlb_tools. Many thanks to the authors of these great projects!

Please refer to Neural-Data-Transformer and nlb_tools for detailed installation and environment requirements.

Please download NLB'21 Benchmark from the official website here. The MC-Maze Subset from NLB'21 Benchmark will be used to exemplify our method in the following sections.

This repository is based on the Pytorch version of the BLEND implementation.

We have provided the model implementation, training, and evaluation code, with detailed instructions shown as follows:

Run the following script for original neural-data-transformer or LFADS model training on MC-Maze Dataset:

git clone https://github.com/dddavid4real/BLEND
cd BLEND
bash scripts/train.sh

For evaluation on the test set, please run:

bash scripts/eval.sh

Run the following script for training hard-label distillation on MC-Maze Dataset:

bash scripts/train_lupi.sh

For evaluation on the test set, please run:

bash scripts/eval_lupi.sh

Before running the training or evluation scripts for these three types of distillation frameworks, please specifify which distillation strategy you intend to use by setting the indicator to True in configs/mc_maze_LUPI_soft.yaml:

TRAIN:
  ...
  SOFT_DISTILLATION: True #* Soft Label Distillation
  # FEATURE_DISTILLATION: True #* Feature Distillation
  # CORRELATION_DISTILLATION: True #* Correlation Distillation
  ...

After this, run the following script for training soft-label/feature/correlation distillation on MC-Maze Dataset:

bash scripts/train_lupi_soft.sh

For evaluation on the test set, please run:

bash scripts/eval_lupi_soft.sh

• Please open new threads or report issues directly (for urgent blockers) to zguobc@connect.ust.hk
• Immediate response to minor issues may not be available

If you find our work useful in your research, please consider citing our paper at::

@inproceedings{guo2024blend,
  title={BLEND: Behavior-guided Neural Population Dynamics Modeling via Privileged Knowledge Distillation},
  author={Guo, Zhengrui and Zhou, Fangxu and Wu, Wei and Sun, Qichen and Feng, Lishuang and Wang, Jinzhuo and Chen, Hao},
  booktitle={The Thirteenth International Conference on Learning Representations,
                  {ICLR} 2025, Singapore, April 24-28, 2025},
  publisher={OpenReview.net},
  year={2025},
}