Behavior-Reinforced Latent Alignment for Generating Functional Neural Spike Patterns ^*

Cognitive neural prostheses offer a potential therapy for neural pathway damage by modulating neural activities of downstream cortical regions. Previous studies using reinforcement learning (RL) frameworks have made significant progress in generating spike patterns that can effectively induce desired behaviors. However, these approaches have struggled to generate realistic neural activities that align with natural firing patterns. In this work, we propose a novel method that aligns the latent dynamics of transregional neural population activities through behavioral reinforcement, achieving both natural neural encoding and optimal behavioral performance. Specifically, we train two Transformer-based variational autoencoders (tVAEs) on upstream and downstream spike trains through self-supervised training. The latent dynamics inferred from the upstream tVAE encoder are aligned with the downstream latent dynamics. Functional spike patterns are then generated from the aligned latent representations via the downstream tVAE decoder, with the alignment weights tuned toward maximal behavioral rewards by policy gradient. We validate our method using neural recordings from the medial prefrontal cortex (mPFC) and the primary motor cortex (M1) of a Sprague Dawley rat performing a two-lever discrimination task. Preliminary results demonstrate that our method successfully generates neural activities that drive task success while preserving natural firing patterns at the neural population level. Moreover, our method outperforms previous RL-based models at both neural and behavioral levels, highlighting its potential in future real-time stimulation experiments with cognitive neural prostheses.