Inferring how internal brain state shapes neural responses with state-dependent diffusion models

Neural responses to the same stimulus can vary dramatically with fluctuations in the brain's internal state and the organism's behavioral state. Yet studying how brain state shapes neural responses is difficult: limited experimental recordings sparsely sample this joint space, leaving many biologically valid state-stimulus pairings unobserved. Existing methods either capture state-dependent modulation with restricted model assumptions or synthesize neural activity without explicitly controlling internal brain state. We introduce BrainStateDiff, a conditional diffusion framework that models the spatiotemporal distribution of future neural activity and behavior given external stimuli and recent brain states. Motivated by two gaps in finite recordings, missing state-stimulus combinations and limited repeats of observed conditions, BrainStateDiff introduces two state-aware sampling strategies: cross-state sampling, which recombines observed states and stimuli to synthesize plausible neural responses from unobserved pairings, and matched-state sampling, which resamples from observed pairings to capture local trial-to-trial variability. We evaluate BrainStateDiff on our newly collected wide-field calcium imaging data in mouse V1, electrophysiology data in monkey V4, and two-photon calcium imaging data in mouse striatum. Our results show that BrainStateDiff generates biologically meaningful samples that preserve the state-dependent structure of real recordings and expand neural response space coverage beyond the original recordings. BrainStateDiff establishes state-conditioned generative modeling as a practical framework for studying how internal brain state shapes neural responses beyond the state-stimulus combinations directly observed in the experiments.