MRDS-JAX

Code for the paper Modeling state-dependent communication between brain regions with switching nonlinear dynamical systems, ICLR 2024.

Overview

A hybrid generative model for multi-region neural data that combines structured SSM dynamics with a Transformer-based inference network. The model infers latent states and communication signals across brain regions, trained end-to-end using a structured variational inference scheme.

Structure

• model.py — top-level model definition
• generative_models.py — SSM generative components
• inference.py — Transformer-based inference network
• networks.py — shared network components
• run_training.py / run_analysis.py — training and analysis entry points
• run-configs/ — experiment configuration files

Dependencies

Python 3.x, JAX, TensorFlow (data pipeline)