OMiSO: Adaptive optimization of state-dependent brain stimulation to shape neural population states
Authors: Yuki Minai
Joana Soldado-Magraner
Byron M. Yu
Matthew A. Smith
Venue: NeurIPS 2025 poster ()
Keywords: closed-loop optimization
microstimulation
neural population activity
dimensionality reduction
latent variable models
reinforcement learning
behavior cloning
proximal policy optimization
Primary area: neuroscience_and_cognitive_science
TL;DR: OMiSO (Online MicroStimulation Optimization): a brain stimulation framework to drive neural population activity toward specified states by adaptively choosing stimulation parameters based on brain state information

Abstract: The coordinated activity of neural populations underlies myriad brain functions. Manipulating this activity using brain stimulation techniques has great potential for scientific and clinical applications, as they causally influence the nervous system. To improve the accuracy by which one can manipulate neural activity, it is important to (1) take into account the pre-stimulation brain state, which can influence the brain’s response to stimulation, and (2) adaptively update stimulation parameters over time to compensate for changes in the brain’s response to stimulation. In this work, we propose Online MicroStimulation Optimization (OMiSO), a brain stimulation framework that leverages brain state information to find stimulation parameters that can drive neural population activity toward specified states. OMiSO includes two key advances: i) training a stimulation-response model that leverages the pre-stimulation brain state, and inverting this model to choose the stimulation parameters, and ii) updating this inverse model online using newly-observed responses to stimulation. We tested OMiSO using intracortical microstimulation with a ``Utah'' array and found that it outperformed competing methods that do not incorporate these advances. Taken together, OMiSO provides greater accuracy in achieving specified activity states, thereby advancing neuromodulation technologies for understanding the brain and for treating brain disorders.

OMiSO: Adaptive optimization of state-dependent brain stimulation to shape neural population states
Yuki Minai Joana Soldado-Magraner Byron M. Yu 13,45,,55 Matthew A. Smith 133,5
NNurosccience Institute, Carnegie Mellon University
Machine Learning Department, Carnegie Mellon University
Center for the Neural Basis of Cognition
Department of Electrical and Computer Engineering, Carnegie Mellon University
Department of Biomedical Engineering, Carnegie Mellon University {yminai, , jsoldado, byronyu, msmi Saamith@@@aan Bandreew cmu. edu Denotes equal contribution.
Abstract
The coordinated activity of neural populations underlies myriad brain functions. Manipulating this activity using brain stimulation techniques has great potential for scientific and clinical applications, as they causally influence the nervous system. To improve the accuracy by which one can manipulate neural activity, it impor- tant to (I) take into account the pre-stimulation brain state, which can influence the brain's 's response to stimulation, and (2) adaptively update stimulation parameters over time to compensate for changes in the brain's 's response to stimulation. In this work, we propose Online MicroStimulation Optimization (OMiSO), brain stimulation framework that leverages brain state information to find stimulation parameters that can drive neural population activity toward specified states. OMiSO includes two key advances: i) training simulation-respons. model that leverages the pre-stimulation brain state, and inverting this model to choose the stimulation parameters, and ii) updating this inverse model online using newly-observed sponses to stimulation. We tested OMiSO using intracortical microstimulation with "Utah" array and found that it outperformed competing methods that do not incorporate these advances. Taken together, OMiSO provides greater accu- racy in achieving specified activity states, thereby advancing neuromodulation technologies for understanding the brain and for treating brain disorders.
Introduction
Most brain functions realized through the coordinated activity of neural populations 1 Causal perturbations of the brain, for example with electrical stimulation, can influence brain function by modulating neural population activity [5]. This activity is variable from moment to moment, and such variability in part reflects the current state of the brain [1 The brain's state can influence how neural populations respond incoming sensory stimuli 7 g and is therefore also important for understanding how the brain responds to causal perturbations.
Closed-loop brain stimulation tools iteratively test and adjust stimulation parameters based on observed neural responses to induce desired brain activity and/or behavior. Prior studies proposed closed-loop brain stimulation tools drive neural population activity toward specific states [9] 10] These methods learn and update the mapping between a wide range of potential brain stimulation parameters and brain responses during experiment, but they do not take into account the state of the brain when choosing the stimulation parameters (but see [11 for a study testing this idea in simulation). Deep brain stimulation (DBS) approaches have leveraged brain state information
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

to update stimulation parameters for improving treatment efficacy [12 [3] or reducing the energy consumption of the stimulation device [14. However, these methods were designed for manipulating low-dimensional (low-d) bio-markers (e.g., local field potentials [4] rather than the activity of
many simultaneously recorded neurons.
In this study, we develop Online MicroStimulation Optimization (OMiSO), a brain stimulation framework that manipulates neural population activity by incorporating the brain state to choose the stimulation parameters. OMiSO characterizes neural population activity in a low-d latent space identified using dimensionality reduction [15]. We refer to the low-d projection of the recorded population activity as a "brain state". The objective of OMiSO is to identify stimulation parameters that induce specified target brain state given the pre-stimulation brain state. Specifically, OMiSO first fits a stimulation- cimulation-response model, which predicts brain responses to different stimulation parameters given the pre-stimulation state. OMiSO then inverts the trained dimulation-responeee model, which updated in real time (i.e., online), to output stimulation parameters for inducing a target brain state.
We tested OMiSO using electrical microstimulation (uStim) in a macaque monkey implanted with a multi-electrode Utah array in the prefrontal cortex (PFC, area 8Ar). OMiSO optimized the location of five stimulated electrodes on each trial, enabling the generation of complex spatial patterns of electric fields across a wide range of electrodes in the array. We found that taking into account brain state prior to uStim (pre-uStim) was beneficial to improve the prediction of the brain's 's response to uStim (Section 3.1). An inverse model, which incorporates brain state, successfully identified the stimulation parameters (in this case, electrode locations) necessary to achieve the specified brain states (Section 3.2) By adaptively updating the inverse model using newly-observed stimulation-responee samples, OMiSO significantly improved the uStim parameter optimization performance (Section 3.3).
Methods
2.1 OMiSO overview
The goal of OMiSO is to find set of stimulation parameters (termed "stimulation pattern")) to drive neural population activity toward targeted state given the pre-stimulation state (Fig. [A, left). Specifically, OMiSO first collects and merges neural activity across multiple experimental sessions by using latent space alignment (Section 2.2) The merged sessions
used to fit stimulation- response model, which predicts the post-stimulation brain state given the stimulation pattern and the pre-stimulation brain state (Section 2.3) OMiSO then inverts the trained stimulation-response model to output a stimulation pattern that can create a target state (Section 2.4p. The inverse model is used to choose a stimulation pattern in a brain stimulation experiment. On each trial, OMiSO observes the rresstimulation brain state Zre and uses the inverse model T 2/ Pte x TTrr)), which is parameterized with 0, choose stimulation pattern & that minimizes the error between the target state 2 larg and the induced brain state N3 (Section 2.5)
argmin Tang za
(I)
Then, using the newly-observed PPr and zg. OMiSO adaptively updates the inverse model mo to improve the selection of (Section 2.6 Fig. [ right). The hyperparameters used in the model
fitting and brain stimulation experiments are summarized in Section ST
2.2 Latent space identification and alignment
When the stimulation parameter space is large, one can typically test only small fraction of all possible stimulation patterns within an experimental session. This requires merging neural activity across sessions to create a large enough set of stimulation-response samples to learn their relationship. To merge neural activity across multiple experimental sessions, OMiSO identifies a low-d latent space of the high-d population activity in each session and aligns the latent spaces across sessions. Each session consists of two types of trials: "stimulation trials". in which we applied stimulation, and ito-stimllttion trials", in which we did not apply stimulation. Following [10], for each session, we use Factor Analysis (FA) to identify a latent space that captures activity co-variations between neurons using only no-stimulation trials. To align the identified latent spaces across sessions, we solve the Procrustes problem to find an orthogonal transformation matrix that maximizes the alignment between two FA latent spaces [16].

---

A
Compusation usitm
PostuBil
Pro-uSti
Early in a session
Late in a session
Electrooes
Nauron
96
Target brain state
Target brain state
Memory period
Pre-stim, state z
Zi
arey
Target activity
Latent dim 2
state
Error
Pre-uStim state
Neuron
Neuron
Post-uStim
Pre-stim state 1
Pre-slim state
uStim
state
Neuron
Neuron
Latent
Time
Figure 1: Goal of OMiSO and experimental paradigm. (A) OMiSO finds a stimulation pattern (grid representing an electrode array, where the stimulated electrodes are indicated by orange cells) to drive neural population activity toward target state (pink star) given pre-stimulation state (black dots). For different pre-stimulation states (Pre-stim state vs. Pre-stim state 2), OMiSO selects different stimulation patterns (green bolt yellow bolt) achieve the same target state. For Pre-stim state 1, OMiSO updates the selected stimulation pattern from early (left panel, green bolt) to late (right panel, orange bolt) in experimental session to achieve the target state. This is necessitated by the fact that the stimulation pattern that achieved the target state early in the session (left panel, green bolt) longer does late in the session (right panel, green bolt). (B) Experimental setup. Top: The monkey performed a memory-guide recorded from multi-electrode array implanted in PFC. Pre-uStim neural activity (gray bars) was recorded over five consecutive 50 ms bins. During the 50 ms computation period (blue unfilled bar), the system chose uStim pattern, which was then applied for 40 ms (orange bar). Post-uStim neural activity (pink bar) was measured starting 10 ms after stimulation offset. Bottom: neural population activity was analyzed in a low-d (in our experiments, 5-d) latent space.
saccade task (Section S3) while spiking activity was
After collecting data over R experimental sessions, for the ith session, OMiSO extracts list of n, "usable" electrodes e; E IR,,, where each element of e; is an integer index an electrode. An electrode is deemed "usable" if it satisfies criteria involving its mean firing rate, Fano factor, and coincident spiking with other electrodes (see Section S2). In our multi-electrode array stimulation set-up with 96 implanted electrodes (indexed 96), e; represents the subset of electrode indices identified as usable in the ith session (e.g., [1, 10, .2....,93) 93)).
For each time bin indexed by 1.....JJ Js (where J, is the total number of time bins used to fit FA for the ith session), OMiSO takes spike counts on each usable electrode Ii,5 E I and fits the following FA model using the EM algorithm:
Zij NOO, LiglZ1a N(Az + p4, V.)
where ži,j E R.m (m < ni)) is the low-d brain state for the jth time bin, A, E RR,xm is the loading matrix whose columns define the low-d latent space, 14, E RR contains the mean spike counts for each electrode, and V, E RT xR, a diagonal matrix capturing the independent variance of the spike counts for each electrode. In this work, the latent dimensionality the optimal dimensionality separately for each session based on cross- volos-valdd data likelihoods, then finding the mode of the distribution of optimal dimensionalities across sessions (m 5 in our implementation, Section ST).
chosen by first computing
OMiSO then defines the latent space from one of the R sessions reference latent space Ag E me xm and aligns the latent space of the i-th session A, to Ag. contrast to our previous work (10), which aligns the latent spaces using the common "usable" electrodes, OMiSO introduces extra step to identify the common "stable" electrodes among usable electrodes to make the alignment more robust neural recording instabilities across sessions (see Algorithm in [16). Concretely, for the
3

---

ith session, OMiSO identifies an orthogonal transformation matrix O, Rmxm that fulfills:
Õ argmin [[ :) (e(Eaable: :) 0:00r=1
where Estable is list of nstable stable electrodes that are common to the reference session and the ith session, and || || is the Frobenius Norm. This optimization can be solved in closed-form [17]. The O, found is applied to A; to obtain the aligned latent space Å; E I=
A A,O
(4)
The brain states in the pre-stimulation (Fig. TB. bottom panel, gray dots) and post-stimulation periods (Fig. 3 bottom panel, pink dot) are estimated as the posterior mean from the FA model (Eq. using the loading matrix Ag:
E| E [و,دع = ( m4)
(5)
where 8, = !! + +V,)-1 By using A, the induced brain state ([2t,]aa,))) resides in a common latent space across all sessions.
2.3 Stimulation-responss model fitting
Using the merged neural activity across sessions, OMiSO fits a stimulation respulaaattio---sssseeee model to predict a post-stimulation brain state given stimulation pattern and pre-stimulation brain state. In this work, we used combination of CNN and a LSTT (Fig. [ST). To train the cinnulaatio-response model, OMiSO uses the stimulation trials of R experimental sessions collected in Section 2.2 which involves stimulating with randomly chosen patterns among all possible stimulation patterns defined by the user. OMiSO creates dataset that comprises the pre-stimulation and post-stimulation brain states (where the latent space is defined using latent space alignment, Section 2.2). and the stimulation patterns tested across the R sessions. For the keth trial in the ith session, the pre-stimulation brain state over / time bins 7 E Rmxt and induced post-stimulation brain state 2 E RT are computed using Eq. 5 Note that Z denotes pre-stimulation brain states across multiple timesteps (a matrix), whereas z denotes a post-stimulation brain state at single timestep (a vector). The entries of z and i,k corresponding to the user-defined target dimensions within the dimensional aligned
latent space (1 < m) are then subselected and collected in new vectors 2 T and &PPs RR (see Section S3 for how / target dimensions are chosen). Given a planar grid of electrodes of size h x u (h, u E R), the stimulation pattern St,e R tested in the kth trial during the ith session is encoded using value of 1 for stimulated electrodes and for all other electrodes.
To train the stimulation-respont response model f( Post [St,,, _. which maps the stimulation pattern S,, and pre-stimulation state _ to the predicted post-stimulation state 2j,k E IR', OMiSO minimizes the mean squared error (MSE) between the predicted Post and observed &Pot post-stimulation brain states. OMiSO performs R-fold cross validation, where each session is used as a validation set once (Section S4h. The final prediction i is obtained by averaging the predictions of the R models.
Post
Post
2.4 Inverting the stimulation-responee model
Using the trained stimulation-responsee model, OMiSO estimates a stimulation- ressonse inverse model. The goal of the inverse model is to identify a stimulation pattern that can reach target brain state given the pre-stimulation brain state. More specifically, a stimulation-response inverse model [o Tare, Z parameterized by 6, receives the target brain state Targ e IR' and pre-stimulation state 2ht E Rixt_ Then, it returns u-dimensional vector of stimulation suitabilities for each electrode $1,k E R", where u is the total number of candidate electrodes to be stimulated whose value determined by the user and < h (Section S3). Each entry in & is a number between 0 and I representing the suitability of stimulating each electrode for achieving the target brain state, where a value closer to I indicates greater suitability. Note that the vector is normalized to sum to one.
OMiSO generates synthetic data using the trained ssimuoa model to train an inverse model we (Fig. SI). The stimulation-responn model outputs the average post-stimulation state over R models a Post
R for any given combination of a stimulation pattern S hxx and pre-stimulation

---

state the ERixt, By considering the predicted post-stimulation state 2 &ost as target state Tur, OMiSO constructs synthetic data tuples (0 T 8). where 8 E {0, 1 is a vector representation of S including only the candidate electrodes for stimulation. The value for each candidate electrode in & is if stimulated and O if not. The generated synthetic data is used to train the inverse model by minimizing the binary cross entropy over multiple epochs with respect to 0 (Section (SS). In this study, we used Multi Layered Perceptron (MLP, Fig. (SI) as the stimulation-respel ressonse inverse model.
2.5 Stimulation pattern selection on each trial
To find the stimulation pattern for achieving target brain state with the inverse model, each brain stimulation session starts with latent space identification trials, where no stimulation is applied. OMiSO extracts a list of usable electrodes ei. which includes electrodes, using these trials. The observed spike count vectors æijj E R for all time bins j across these trials are used to fit the FA parameters A, E Ra." p4 R" and V E R The identified latent space A, is aligned to the reference latent space Ag using the methods described in Section 2.2, yielding A, E R To start the optimization, OMiSO loads the user defined target brain state 2 Tan E R', the trained inverse model Tey and the number of electrodes used for each stimulation pattern ssm Z. Note that out of h x u total electrodes (used for the stimulation pattern representation S). not all may be suitable for stimulation (e.g., some may have no effect on the neural population activity, Section [S3). This results in u ( h x r) candidate electrodes (used for the stimulation pattern representation 8). Of those, OMiSO chooses Tistim electrodes to form each stimulation pattern, where nsmm S
On each trial, OMiSO selects the stimulation pattern to perform given the observed pre- pre-stimulation brain state using the inverse model πe. On the kth trial, OMiSO estimates the pre-stimulation brain state 2 E IR mg in real time using Eq. 5 The entries of i.k corresponding to the target dimensions Xt
are subselected as Z Rixt, and 2 and = Targ are passed to we to get u-dimensional vector & of stimulation suitabilities for each candidate electrode. Using $t,,, OMiSO chooses a stimulation pattern with epsilon greedy algorithm. With probability 1 Ei,k, where 0 < Ei, S 1, it chooses the Ssm electrodes with the highest predicted suitabilities in & With probability Es it chooses nsšm electrodes stochastically following the softmax transformed probability yºA, R computed for each electrode indexed by w 1, ..., as:
exp
PLA
(6)
exp(
where 65;, E IR is the predicted stimulation suitability for the wth electrode. In this way, electrodes with low predicted stimulation suitability can be occasionally chosen to induce exploration, but the electrodes with higher stimulation suitabilities are chosen more often. OMiSO uses a time-varying eie to initially encourage exploration and gradually shift to exploitation:
Ei,A max (Eiiit -7*, E(toor)
where 0 < mii S 1 is an initial value of Ei 0 S 7 S 1 is a discount factor, 0 S Enoor S is the smallest allowable value of 6i (Etoor S Einit, and & is the current trial index. OMiSO uses Elour to maintain some amount of exploration throughout a session. Note that the inverse model Ty estimates the suitability of each electrode independently, without explicitly modeling potential interactions between electrodes stimulated simultaneously. Consequently, the selection of the top Visim electrodes is based on their individual contributions producing the target state. This design choice was motivated by our empirical observation that, using our stimulation parameters (25 μA, 50 Hz, 40 ms). stimlating a single electrode primarily affected nearby recording electrodes, with minimal interaction with other concurrently stimulated electrodes.
Once the stimulation pattern is selected with the epsilon greedy algorithm, OMiSO stimulates the selected nsmm electrodes and measures post-stimulation spike counts 2 E "s to compute
Post
Post iik
E RR (Eq. 5p. The set of =Tauy Z PPsst and the tested stimulation pattern &,,k E 1)" are i,k
stored update the inverse model batched manner, as explained in the next section.
2.6 Adaptive updating of the stimulation-responne inverse model
During the brain stimulation experiment, OMiSO adaptively updates the ddimulation-responee inverse model ne using the recent stored observations. In this study, we implemented OMiSO's adaptive
5

---

model update using the clipped policy gradient objective function used in Proximal Policy Optimiza- tion (PPO) 18 PPO is widely-used RL method to update an action selection strategy. OMisO uses PPO-inspired objective function to update the stimulation-respons inverse model to improve the stimulation pattern selection strategy. Compared to other policy update algorithms such as REINFORCE [99] and Actor Critic (20), PPO offers more stable policy update by using clipped objective function, which might be particularly well-suited to handle trial-by-trial variance in neural activity.
At each update, OMiSO increases the suitability of stimulating electrodes that produced a post- stimulation state closer to the target state and decreases it for electrodes that did not. Specifically, for the kth stored trial in batch update of the inverse model (a batch five trials in our implementation, Section ST), OMiSO computes the ad advantage A E IR for the stimulation pattern used in the trial:
A = -|| 1| Tang PPsl 1|2 + b
(8)
where b E IR is fixed baseline value. In contrast to the original PPO implementation, which uses a separate model (often called a value network) to estimate the expected outcome at each state, OMiSO uses fixed baseline b shared across all pre-uStim states. By avoiding the simultaneous online updating of two interacting models (the inverse model and the value network), this design reduces the inverse model update instability especially when only limited number of observations are available. The value of was chosen by running a simulation (Section (6). With the computed advantage A, OMiSO updates the parameters l in the stimulation- respooose inverse model me by performing gradient ascent over multiple epochs on the following PPO-inspired objective function:
Lrro Eesr min (1), clip
(ri, (0), Eclip, 1 Eclip) Ai,
1T (T 111 2
where rf (0)
Th (8) 1E 2.1)"
e = 1, .... Astim is the index of the electrode among Rss electrodes used for the stimulation pattern 8t,ks and the clipping parameter 0 S Sclip < 1 defines upper and lower bounds of x (0) through the clip function, limiting the magnitude of the gradient update to avoid taking a too large update step (Section s1p. To (8 I= 2 represents the predicted stimulation suitability for the eth stimulated
Tury
electrode under the updated aiimulation-respons inverse model. Similarly, Tew, ( ( t is the predicted suitability under the original stimulation-response inverse model before updating O. The ratio (0) is the change in suitability for selecting the eth electrode. Computing this ratio separately for each electrode rather than for an entire stimulation pattern is a design choice, motivated by the same empirical observation as in Section 2.3 for the selection of which electrodes to stimulate. To maximize this objective function, the model updates 0 to increase rf (f) when the advantage is positive and decrease it when the advantage is negative. By updating 6, OMiSO not only updates the predictions of the stimulation patterns tested in the stored observations, it also updates the predictions of stimulation patterns that were not tested.
2.7 Experimental details
We tested OMiSO using uStim in a rhesus macaque monkey with a 96-electrode Utah array implanted in the PFC. In each experimental session, the monkey performed memory-guided saccade task (Section S3p. Each trial began when the monkey fixated a central dot, followed by a memory target briefly appearing at a peripheral location. After the memory target disappeared, the monkey needed to rememher the target location during a "memory period" (Fig. B). The disappearance of the central dot served as the "go cue". after which the monkey made a saccade to the remembered location. On about 80% of trials, we applied uStim for 40 ms during the memory period (Fig. [BB. The stimulation patterns tested in this study consisted of all possible T' electrode spatial patterns chosen from
= 20 candidate electrodes (15,504 possible patterns, Fig. S2 see Section S3 for the selection criteria of the candidate electrodes). Each experimental session started with 120 no-uStim trials for latent space identification. We identified m = 5 latent dimensions by applying FA to spike counts taken in 50 ms bins. Among them, we chose = 2 target dimensions (Section (S). Post-uStim spike
counts (50 ms bin) were calculated starting 10 ms after uStim offset (Fig. [F,g.11. pink bar), which we empirically verified to be sufficient to recover from stimulation artifacts induced by our uStim set-up.
The number of sessions used in each analysis is summarized in Section S4

---

A
Pre-uStim states
Post-uStim states
B
Beter
Pre-uStim state (dim 2)
135
uStim
Post-uSüm state (dim 2)
uStim pattern only
315"
pattern 1
uStim pattern
Pre-stim1 (MLP)-
Avg over triwhs
uStim pattern
Pre-stim3 (MLP)-
0
135
uStim pattern
Pre-stimS (MLP)-
uStim pattern Pre-stimS (LSTM)-
One triel
Pre-stim5 only (LSTM)
-2
0.178 0:200 0.225
Pre-uStim state (dim 1)
Post-uStim state (dim 1)
5-fold CV MSE
Figure The pre-uStim brain state affects the post-uStim brain state. (A) An example electrode uStim pattern. The orange cells in the array map indicate the location of the stimulated electrodes. Each small circle (135% memory condition) or cross (3159 memory condition) indicates the brain state single trial, and the large black circle and cross indicate the mean across trials within each memory condition. (B) Post-uStim state prediction performance across models (see Fig. [T]for for model architectures). Error bars indicate 1 standard error across cross-validation folds (Section 2.3
3
Results
3.1
The brain's response to stimulation depends on the pre-uStim brain state
Does neural population activity prior to brain stimulation affect the brain's response to stimulation? To assess this, we measured the impact of the pre-uStim state on the uStim response in behavioral experiment where different memory conditions (Section S3 were used to create distinct pre-uStim brain states in low-d space identified using FA (Fig. 2A, left). We observed that uStim shifted neural activity within the low-d space (Fig. 2A. right; see Fig. S3 for additional examples) away from the location where the activity would have been without uStim (cf. Fig. S3 leftmost column, bottom panel). This demonstrates that uStim perturbed the neural population activity. Furthermore, the induced post-uStim activity states appeared to depend on the memory condition (Fig. (Fig. right; see Fig. S3 for additional examples). This suggests that brain states induced by uStim depend on the pre-uStim brain state.
To quantify whether knowledge of this state dependency helps improve the prediction of brain responses to uStim, we compared the post-uStim state prediction performance of the stimulation- response model with and without pre-uStim state information (Fig. ST]. For a model without pre-uStim state information, we used CNN which predicted post-uStim states only as a function of the uStim patterns applied (as in (TO), termed "uStim pattern only". For models with pre-uStim state information, tested combination of CNN and a MLP [21] with up 5 pre-uStim time bins, termed "uStim pattern + Pre-stimm 1,3,5) (MLP)", and a combination of a CNN and an LSTM with pre-uStim time bins, termed "uStim pattern + Pre-stim5 (LSTM)". To collect training data to fit these 5-electrode uStim patterns in 6 experimental sessions.
models, we applied randomly-chose
The models with pre-uStim state information (Fig. 23 four middle bars in light blue) outperformed the model without this information (Fig. 2B, top dark blue bar, p < 0.05 for all four models, one- tailed (-test). Among the models with pre-uStim state information, uStim pattern Pre-stimš (LSTM) achieved the best prediction performance by incorporating time bins of pre-uStim state information. A model using only pre-uStim state information, termed "Pre-stimS only (LSTM)" (Fig. 2 bottom dark blue bar), performed worse than uStim pattern + Pre-stim5 (LSTM) (p < 0.05, one-tailed t-test). Thus, both the uStim pattern and pre-uStim state information were important for prediction. The pre-uStim state affects where the activity evolves after uStim, and this can be leveraged to improve the prediction of brain response uStim.
3.2 A stimulation-respons inverse model selects uStim patterns based on the pre-uStim state
OMiSO's 's goal is to choose a uStim pattern that can create a targeted neural population state. While the stimulation-response model trained in Section 3.1 can predict brain responses to any input combination of a uStim pattern and pre-uStim state (Fig. 3 top), they can not do the inverse start with desired brain state and find uStim pattern that will drive neural activity to that state. One
7

---

A
Input
Model
Output
B
0.30
Süimulation-response sample generation
(Disannee to the targat targat stat))
100
Absolute error relative Static
Inversion error
0.25
Post-
Pre-uStim
uStim
states
Stimulaton- response model
state
0.20
90-
Latent dim
Latent
Time
0.15
Static
Static with pre-stim
Stimulation-response inverse model training
0:10
Benter
5000
10000
atent dim 2
Train epoch
Pro-uStim
Target
state
Stimulation
C
Method
Usea pre-stim state
Performe adaptive undate
response
Static
No
No
Lateni cim
inverse
model
Static with pre-stim
Yes
No
Figure 3: Inverting the stimulation-respone model to choose uStim pattern based the pre-uStim state. (A) Schematic of the inversion procedure. (B) Inversion performance. The error shading indicates the standard deviation across 5 inverse models. (C) Methods used for performance comparison to evaluate the model inversion quality. (D) Absolute error relative to the "Static" method computed in 2-d target latent space. Error bars indicate +1 standard error across sessions.
potential approach is to construct lookup table of predicted post-uStim states for all possible input combinations and find the uStim pattern expected to produce the closest post-uStim state the target state. However, this approach is infeasible with continuum of pre-uStim states. Another approach is mathematically invert the commulation-respoone model choose a uStim pattern. However, this is challenging because the stimulation-respone model could map multiple combinations of uStim patterns and pre-uStim states to the same post-uStim state (i.e., the mapping is not one-to-one).
Instead, we approximated the inverse mapping of the stimulation- response model by generating synthetic data (Section 2.4. More specifically, we generated synthetic data using the stimulation- response model, where we considered the model- predicted post-uStim states as potential target states (Section (S)). The generated synthetic data allowed to train the inverse model, which receives pre-uStim state and a potential target state and returns a suitability of stimulating each candidate electrode (Fig. 33,. bottom). For the stimulation- coimulation-rsponsee model, we used uStim pattern + Pre-stim5 (LSTM), which achieved the best prediction performance (Fig. (2 For the dimulation-respoon inverse model, we used an MLP (Fig. ST). since computationally fast and could complete all necessary computations within 50 ms (Fig. BB, computation period) following the observation of pre-uStim states, allowing uStim to be delivered immediately afterward.
To assess inversion quality, we evaluated whether the inverse model could output uStim patterns that induce the appropriate post-uStim neural states (i.c., the post-uStim state predictions). We found that with more training epochs (each of which consists of 15,504 synthetic samples), the inverse model improved choosing uStim patterns that induced post- iist-uStim states close to the target state (Fig. 3 Manhattan distance in the 2-d latent space, Eq. [T]. This improvement saturated around 100,000 epochs with distance (i.e., inversion error) of around 0.13 a.u. (error for perfect inversion is 0). For comparison, the trial-by-trial standard deviation in brain states (in no-uStim trials within the same memory condition) was 0.7 a.u.
To understand the impact of imperfect model inversion on achieving target brain states, we experimen- tally compared the performance of two methods, one with perfect inversion and another that could only be approximately inverted (Fig. 3C). The first method (termed the "Static" method) constructed a lookup table of predicted post-uStim states for all possible uStim patterns (15,504 patterns) using the uStim pattern only model, which does not consider pre-uStim state information. It then found the uStim pattern expected to produce the closest post-uStim state to the target state. Having a complete lookup table implies perfect inversion of the uStim pattern-only model. The second method (termed the "Static with pre-stim" method) used the inverse model of the uStim pattern + Pre-stim5 (LSTM) model trained on the generated synthetic data (which only achieves an imperfect inversion). Neither the prediction lookup table nor the inverse model were updated during the experiments, so we refer to them as static methods. To compare the two methods, we performed brain stimulation experiments

---

A
Trat
Triet 26
Triat
Tra 76
[s,Z
Z
Pro
Post
stimutation
response model
Stimuiation inverse
uStim suitability
Synthetic data
modil
Candidate electrode
Load
Stimulated
electoodes
26
1. Anallyz
pro-uStim state
ProuSSim Latwri
50 Trial
75
100
Latent
D
Static with pre-stim
MISO
Accumulate
Request
Choose uStim
Target dim
online observations
parameters
Stimutation
Pre-uStim
states
state
Torget stane
modal
-2
Update
Respond
-2.5
0.0
-2.6
0.0
Target dim
Target dim
Evaluate post-uStim state
3. Perform uStim
E
Worse
asosions
OMISO
Laterg
Absolute error relative to
Stalii with pre- pre-stim (%)
Post--uiitm
state stsile
reror
Target dim
90-
B
Uses pre-stim state
Performe adaptive update
Beteer
with siim
MSO
OMB0
Stalic pre-slim
No
-255
0.0
Static
MISO
Yes
OMIOO
Yes
Target dim
Figure 4: Performance OMiSO in brain stimulation experiments. (A) Schematic of OMiSO.
OMiSO trains stimulation-responee and stimulation-respons inverse models using previously col- lected atimulationrespone samples (top box). each uStim trial, OMiSO performs five steps: (1) analyzes the pre-uStim brain state across five time points (gray dots), (2) inputs these pre-uStim states and the target state into the inverse model to choose uStim pattern (represented as a grid), (3) applies the chosen uStim pattern, (4) compares the resulting post-uStim state (dot) to the target state (star), and (5) updates the inverse model every five trials using the new stimulation-eespons samples. (B) Methods used for performance comparison. (C) Examples of stimulation suitability for each electrode (top) and selected uStim patterns (bottom) by OMiSO during an example session using the same target state across trials. (D) Induced post-uStim states by each method during example session (the same session as in panel C). Different colors (gray: static with pre-stim; blue: MiSO; pink: OMiSO) indicate the method used select uStim patterns. Each small dot shows the post-uStim brain state from a single trial, whereas each large hollow dot shows the mean post-uStim state across trials. The red dot in each panel is the target brain state, shared across all three panels and trials, (E) Absolute error of different methods relative to "Static with pre-stim", computed in the 2-d target latent space. Error bars indicate standard error across sessions.
which we applied uStim patterns chosen by each method in an attempt achieve 2-d target brain state. If the effect of the inversion error on the uStim pattern selection performance is small, the "Static with pre-stim" model, which leverages pre-uStim information, should outperform the "Static" method. We found that the mean error to the target state of "Static with pre-stim" across sessions was smaller than for the "Static" method, although the difference was not statistically significant (Fig. BD. N 14 sessions, p 0.3, Wilcoxon signed-rank test). This provides support for incorporating the inverse model into an adaptive brain stimulation framework to choose uStim patterns based on pre-uStim state information.
3.3 OMiSO adaptively updates the inverse model during brain stimulation experiment
During each experimental session, new stimulation-responsee samples become available. Leveraging these samples may help us adaptively update the stimulation-response inverse model to account for inversion errors, as well as changes in the stimulation-responsee relationship within and across sessions. To test this, we developed adaptive brain stimulation framework, OMiSO (Fig. 4A, Section 2.1

---

where we used a PPO-inspired objective function to adaptively update the stimulation stimulatio-resppose inverse model (Section 2.6). We tested OMiSO during brain stimulation experiments and compared its performance against an inverse model that was not adaptively updated ("Static with pre-stim", tested in Section (.22) (Fig.BB. OMiSO had limited observations (104 trials on average after the latent state identification trials within session in 14 experimental sessions) to update the inverse model (after every 5 trials), which could put it at a disadvantage due to overfitting compared to the "Static with pre-stim" method. We also compared OMiSO's performance against our previously proposed method, MiSO [TO], which adaptively updates its predictions but does not leverage pre-uStim states. OMiSO adaptively updated the stimulation suitability for each electrode over time (Fig. 4C, top; example session), changing the uStim pattern applied for achieving the same target state (Fig. IC. bottom). Changes in pre-uStim brain states could also result in different uStim patterns being selected over time, even in a non-adaptive method. However, the induced brain states produced by the adaptive method were closer the target state than those produced by the "Static with pre-stim" method (for which the inverse model was not adaptively updated) (Fig. [4)). adaptive model updates were essential for improving the accuracy in inducing the target state. Across multiple sessions, OMiSO drove neural activity closer to the target state than the two alternative methods (Fig. 4E, N 14 sessions, p 0.05 for Static with pre-stim and p = 0.01 for MiSO, Wilcoxon signed-rank test) by taking into account the pre-uStim state and by adaptively updating the cimulation-responsee inverse model.
response
(Fig.4D)). This indicates that the
4
Discussion
We propose brain stimulation framework, OMiSO, which involves two key methodological advances for shaping brain states: incorporation of the pre-stimulation brain state and adaptive updating of the model used to choose stimulation pattern. In brain stimulation experiments, OMiSO outperformed methods that do not incorporate these advances. While the neural activity used in this study consists of spiking responses recorded with planar grid of electrodes implanted in the brain, and each stimulation pattern specifies the location of the stimulated electrodes within the grid, OMiSO can be readily applied to other stimulation and recording protocols (e.g., holographic optogenetics [2) or other uStim parameters such as frequency and amplitude).
In comparison to MiSO [10]. OMiSO introduces several key technical advances. First, MiSO does not account for the pre-stimulation brain state when predicting brain responses, whereas OMiSO explicitly incorporates it, enabling more accurate prediction of brain responses. Second, during online adaptation, MiSO relies on a tabular approach that stores the predicted post-stimulation state for each candidate stimulation pattern and updates only the entry corresponding to the tested pattern. This approach is sample inefficient as it does not propagate the knowledge from a new stimulation-respont response observation to improve the predicted response of other stimulation patterns. OMiSO instead employs an inverse model that updates its parameters online, enabling learning beyond the tested stimulation pattern. Finally, the tabular method in MiSO becomes impractical the size of the stimulation parameter space grows. OMiSO's bodel-based approach enables better scaling to larger parameter spaces, expanding the search from 4,560 ("96 choose electrodes") stimulation patterns [0O] to 15,504 ("20 choose 5 electrodes"") stimulation patterns, with improved performance over MiSO.
Although we experimentally demonstrated the advantages of OMiSO, there are scenarios where it might be less beneficial. First, in response to strong stimulation, the brain might be driven to a similar post-stimulation state regardless of the pre-stimulation activity. Second, in clinical applications, neural states can exhibit substantial non-stationarity. For example, in individuals with epilepsy, the neural state distributions might be quite different during periods with scizures (when neurons tend to fire in bursts) versus periods with no seizures 23 When using OMiSO, one should ensure that the neural state distributions are similar between the training period and online use.
OMiSO can be adapted to both clinical and basic neuroscience applications, in which the goal is to drive neural activity to target brain state (Eq. D. In clinical applications (e.g.. DBS), the target could be defined as neural activity corresponding to a healthy brain state. OMiSO can then identify stimulation patterns that move neural activity towards this desired state. Alternatively, OMiSO could be extended to drive neural activity away from unhealthy brain states. For basic neuroscience, OMiSO can increase the accuracy with which one perturbs neural activity to elucidate the neural mechanisms underlying sensory, cognitive, and motor function (55 24
10

---

Acknowledgments
This work was supported by Japan Student Services Organization fellowship, NIH CRCNS RO1 MH1 18929, NSF NCS DRL 2124066, and Simons Foundation NC-GB- BE--B--CULM-0000224-00000 The authors are inventors on pending International Patent Application No. PCT/US25/47445, which relates to the methods developed in this paper. We are grateful to Samantha Schmitt for assistance with data collection, and to our animal care staff.
References
[1] Bruno Averbeck, Peter E Latham, and Alexandre Pouget. Neural correlations, population coding and computation. Nature reviews neuroscience, 7(5):358 33),658-666, 2006.
[2] Shreya Saxena and John P Cunningham. Towards the neural population doctrine. Current opinion neurobiology, 55:103-111, 2019.
[3] Saurabh Vyas, Matthew Golub, David Sussillo, and Krishna Shenoy. Computation through neural population dynamics. Annual review of neuroscience, 43(1):249-22) 275, 2020.
[4] R Becket Ebitz and Benjamin Y Hayden. The population doctrine in cognitive neuroscience. Neuron, 109(19)33055 30055536---68, 2021.
[5] Marlene R Cohen and William T Newsome. What electrical microstimulation has revealed about the neural basis of cognition. Current opinion neurobiology,
2004.
[6] Marlene R Cohen and Adam Kohn. Measuring and interpreting neuronal correlations. Nature neuroscience, (4(7):811-815 2011.
[7] John H Reynolds and Leonardo Chelazzi. Attentional modulation of visual processing. Annu. Rev. Neurosci., (((1):611-647) 2004.
[8] David A MeCormick, Dennis B Nestvogel, and Biyu J He. Neuromodulation of brain state and behavior. Annual review of neuroscience, 43(1):391-4155 2020.
[9] Sina Tafazoli, Camden J MacDowell, Zongda Che, Katherine C Letai, Cynthia R Steinhardt, and Timothy J Buschman, Learning to control the brain through adaptive closed-loop patterned stimulation. Journal of Neural Engineering, 1755):056007, 2020.
[10] Yuki Minai, Joana Soldado-Magranes Matthew Smith, and Byron M Yu. MiSO: Optimizing brain stimulation to create neural activity states. Advances in Neural Information Processing Systems, 37.24126-24149 2024.
[11] Yuxiao Yang, Shaoyu Qiao, Omid Sani, Isaac Sedillo, Breonna Ferrentino, Bijan Pesaran, and Maryam M Shanechi. Modelling and prediction of the dynamic responses of large-scale brain networks during direct electrical stimulation. Nature biomedical engineering. 5(4):324 345, 2021.
[12] Arthur Guez, Robert D Vincent, Massimo Avoli, and Joelle Pincau. Adaptive treatment of epilepsy via batch-mode reinforcement learning. In AAAI, volume 8, pages 1671-1678, 2008,
[13] Simon Little, Alex Pogosyan, Spencer Neal, Baltazar Zavala, Ludvic Zrinzo, Marwan Hariz, Thomas Foltynie, Patricia Limousin, Keyoumars Ashkan, James FitzGerald, et al. Adaptive deep brain stimulation in advanced parkinson disease. Annals of neurology, (4(3):449-457 2013.
[14] Qitong Gao, Stephen L Schmidt, Afsana Chowdhury, Guangyu Feng, Jennifer Peters, Kather- ine Genty, Warren M Grill, Dennis A Turner, and Miroslav Pajic. Offline learning of closed-loop deep brain stimulation controllers for parkinson disease treatment. In Proceedings of the ACM/IEEE 14th International Conference on Cyber-Physical Systems (with CPS-loT Week 2023), pages 44 54-55, 2023.
[15] John Cunningham and Byron M Yu. Dimensionality reduction for large-scale neural record- ings. Nature neuroscience, 1500 2014.

---

[16] Alan D Degenhart, William E Bishop, Emily R Oby, Elizabeth C Kabara,, Steven M Chase, Aaron P Batista, and Byron M Yu. Stabilization of brain- computer interface via the aligmment of low-dimensional spaces of neural activity. Nature biomedical engineering, 4(7):672-665, 2020.
[17] Peter H Schönemann. A generalized solution of the orthogonal procrustes problem. Psychome- trika, 31(1):1-10, 1966.
[18] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06444 2017.
[19] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256, 1992.
[20] Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
[21] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back- propagarpgg errors. Nature, (23 1986.
[22] Hillel Adesnik and Lamiae Abdeladim. Probing neural codes with two-photon holographic optogenetics. Nature neuroscience, 24(10):155-1166, 2021.
[23] Roland D Thijs, Rainer Surges, Terence O' O'Brien, and Josemir W Sander. Epilepsy in adults. The lancet, 393(10112))689-701,, 2019.
[24] Mehrdad Jazayeri and Arash Afraz. Navigating the neural space in search of the neural code. Neuron, 1005 2017.
[25] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna:
A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2019.
[26] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetooo, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, hirh-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8024-8035, 2019.
12

---

Supplementary material
S1 Summary of hyperparameters
The table below summarizes the hyperparameter values used in this study. The latent dimensionality m for the FA model was chosen by first computing the optimal dimensionality separately for each session by maximizing the cross- validaadd data likelihood, then finding the mode of the distribution of optimal dimensionalities across sessions. The number of electrodes used for latent space alignment was chosen to be smaller (by to electrodes) than the number of common usable electrodes to increase the robustness of alignment. For the stimulation-response model training, the learning rate, weight decay, and the number of training epochs (i.e., number of complete passes through the entire training dataset) were chosen based grid search. For the model inversion, we did not perform extensive tuning of the hyperparameters. such as a max epoch, since the model yielded similar performance (Eq. [1] even with different hyperparameters. Batch size per epoch refers to the number of synthetic data samples generated for each epoch. The parameters for uStim pattern selection and the PPO-inspired adaptive model update were chosen by running simulations (Section (66). uStim parameters (amplitude, frequency. and duration) were set to avoid causing overt behavioral changes or strong post- activity inhibition across the entire array.
Method
Parameter name
Description
Value
Latent space identification & aligmment
m Thsable
FA latent dimensionality
5
Num. of stable electrodes for alignment
40
Stimulation-resppe response model fitting
R
Num. of training experimental sessions
5
t
Num. of pre-uStim state time bins
1-5
Target state dimensionality Optimizer
2
Adamn
Learning rate
0.0006
Weight decay
0.001
Training epoch
20
Model inversion
Batch size per epoch
15.504
Max epoch
100,000
Num. of validation samples
1,000
Validation frequency (epoch)
100
Farly stop patience (epoch)
20
uStim pattern selection
Num. of latent space identification trials
120
Einit
Initial value of Ei,
0.5
6Boor
Minimum value of Ei,k
0.1
?
Discount factor of fil
0.95, 0.96
Adaptive mudel update
Eclip b
Clip value of model update
0.15, 0.2
Baseline for advantage computation
0.6, 0.7
Batch size Optimizer
5 Adam
Learning rate
10-5 10-4
Weight decay
0.001
uStim parameters
Num. of uStim candidate electrodes
20
7iszim
Num. of uStim electrodes
5
Total num, of possible uStim patterns
15,504
h
Horizontal size of electrode grid
10
V
Vertical size of electrode grid
10
uStim amplitude (uA)
25
uStim frequency (Hz)
50
uStim duration (ms)
40
S2 Spiking activity preprocessing
To identify latent dimensions of the neural population activity using FA, we computed binned spike counts during the 1.5 s memory period of the first 120 trials recorded in each session (termed the "latent space identification trials"). We used 50 ms bins, yielding 30 bins per trial. The total number of time bins used to fit the FA model each session was 3,600 (120 trials x 30 bins). The same
13

---

trials were used to extract a list of usable electrodes e; for the ith session based on the following three criteria: mean firing rate >1 Hz, Fano factor <8, and <20% coincident spiking with each of the other electrodes. The FA model was fitted using only the usable electrodes.
S3 Details of uStim experimental paradigm
Experimental procedures were approved by the Institutional Animal Care and Use Committee of Carnegie Mellon University. In each experimental session, the monkey performed memory-guided saccade task. On each trial, the monkey first fixated on a dot at the center of the screen. After establishing fixation, target appeared on the screen for 100 ms. This was followed by a memory period, after which the center dot turned off (go cue) and the monkey performed an eye movement to the remembered target location to receive liquid reward. The fixation duration (i.e., the time from when the monkey acquired fixation to when the go cue was given) was set at either 1.65 or 1.95 seconds for the latent space identification trials and 1.25 or 1.55 seconds for the other trials (randomly chosen, with probabilities 0.8 and 0.2 for the shorter and longer durations, respectively). The location of the target was chosen from four peripheral targets ([45", 135%, 225". 315 [0 70", 135%, 270")) on the latent space identification trials and from two peripheral targets (1135%, 315 or [0?, 70%D on the other trials. These target directions were chosen based on the mapped receptive fields of the recorded PFC neurons so that diverse brain states could be induced.
There were two types of memory--gideed saccade trials: "uStim trials", in which we applied uStim, and "no-uStim trials", in which we did not apply uStim. The experimental system randomly chose which type to perform on each trial. On uStim trials, we applied uStim for 40 ms during the memory period. The stimulation was biphasic with each square pulse in the biphasic pair being 250 js in duration. We set the current amplitude low enough not to induce any eye movements (25 μU for each electrode we stimulated). On each trial, we changed the locations of R'šim 5 stimulated electrodes selected among u 20 candidate electrodes (Fig. S2). while keeping other parameter values such as current amplitude (25 μA) and frequency (50 Hz) fixed. The candidate electrodes for stimulation were selected based three criteria: (1) they were spatially distributed across the array to prevent excessive current concentration when stimulating with multiple electrodes, (2) they were positioned near electrodes with consistent recordings of spiking activity across sessions to ensure reliable evaluation of brain responses to stimulation, and (3) they induced changes in brain activity when stimulated. We selected the candidate electrodes from the full of electrodes, rather than restricting selection to the Taaall stable electrodes, to allow greater spatial diversity in stimulation configurations.
Although we could have defined the target state in the m = 5 dimensional latent space, we instead defined the target state within subset (I (I=32 2) of those dimensions. We made this choice because not all latent dimensions are equally suitable for reliable perturbation of neural activity. For example, the activity in some dimensions might not be easily modifiable by uStim, and some dimensions might not be well-aligned across sessions. Therefore, we chose the two target dimensions based on two criteria: 1) diverse brain states could be induced along these dimensions with uStim (evaluated by visualizing the induced brain states across different dimensions), and 2) they showed good alignment across multiple sessions (as determined by the similarity of the aligned dimensions across sessions, cf. Eq. 3 The dimensions that passed the criteria were not necessarily the top FA dimensions that explained the greatest covariance among the neurons.
To select which uStim pattern to apply, we chose the 7iştim electrodes with the highest predicted suitabilities, with 1 Ei,A probability (Section 2.5). For simplicity, stimulated using the same number of electrodes (Nsmm = 5) when evaluating the performance of OMiSO as during the training sessions. Alternativly, one suitability threshold (e.g.. 0.8) and stimulating all electrodes whose predicted suitability exceeds that threshold.
other electrode selection strategies, for example applying
For the adaptive model update, OMiSO only used new observations without using a replay buffer (i.c., it performed on-policy model updates), in [18].
14

---

S4 Summary of experimental sessions
We conducted 31 brain stimulation experimental sessions in total. Of these sessions, 7 sessions used for offline analyses ("Offline analysis sessions"), 10 sessions with randomly selected uStim patterns were used to train the stimulation-responsee model and the stimulation model ("Training sessions"), and 14 sessions were used to test the performance of different methods ("Test sessions"). Among the 14 test sessions, 6 of them were obtained by running test sessions per day for 3 days. The table below summarizes the number of sessions used for each result reported in the manuscript.
inverse
Figure
Offline analysis sessions Training sessions
Test sessions
Fig.
session
6 sessions
g
Fig.
10 sessions
14 sessions
S5 Details of stimulation-responsse model inversion
To train the stimulation-response inverse model (Fig. SIP) with synthetic data, OMiSO iterates the data generation process and model parameter updates over multiple epochs. In each epoch, OMiSO generates 15,504 synthetic samples, each of which consists of a pre-uStim state the E | uStim pattern s {0, 133, and a predicted post-uStim state averaged across R stimulation-response models, which is used the target state =TTagg E R'. The number of synthetic samples need not match the total number of uStim patterns, although we chose that as our hyperparameter value. To generate one synthetic sample, OMiSO selects one of the experimentally observed pre-uStim states - from the
ciimulatioo-responss model training data as 2nd and pairs with uStim pattern chosen randomly among the 15,504 possible patterns. OMiSO then uses the paired pre-uStim state and uStim pattern to generate post-uStim state prediction that can be used as the target TTary With these synthetic data, OMiSO trains diimulatio-responee inverse model (
||| = Tang z TTe by minimizing the binary cross entropy loss between the chosen uStim patterns s and the predicted patterns with respect to
LECE = - [s, log(š.) + (1 - (z) log(1 kz)
(10)
where and {0, 13", and s, and 8, are the wth entry of these vectors, respectively. Then the next epoch begins, until the stopping criterion is reached.
To evaluate the inversion quality for the stopping criterion, OMiSO uses 1,000 validation samples, which are separately generated using the R coimulation-reponaa models. The kth validation sample consists of pre-uStim state ze E R uStim pattern Sy E {0, 13", and the target state &Tag R. OMiSO first obtains the predicted stimulation suitability of each candidate electrode using the inverse model, To | Tang =hth ). then chooses ssum electrodes with the highest predicted suitabilities in šy. Using a grid format representation of the selected uStim pattern S₁ E IR hxu (locations of the stimulated electrodes indicated with a 1, and all other electrodes have value of 0), OMiSO obtains the predicted post-uStim state using the siimulatoo-respponee model (( z/th) where & Post E g'. Finally, OMiSO evaluates the error between the predicted and target states across all validation samples:
Lin 1000
1000
Targ
Post 2g
(11)
k=
If the inversion with the synthetic data were perfect, this error would equal O. OMiSO performs the validation every 100 epochs (instead of every epoch) to reduce the computational cost of inverse model training. This validation frequency is a tunable hyperparameter that can be adjusted based on the available computing resources. This validation performance was used for the early stopping of the model training to avoid overfitting. OMiSO stops the inverse model training when Eq. TT shows no improvement for 20 consecutive validation performance evaluations.
15

---

S6 Simulations to determine the hyperparameter values for the brain stimulation experiments
To set the values of the hyperparameters for the adaptive brain stimulation experiments (listed in the Adaptive model update section in Section (ST). we ran simulations by using the R trained stimulation- response models (Section 2.3) and the stimulation-respone simulation, we generated 1O0 uStim trials with the trained stimulation- response inverse model and one of the R stimulation- be used in the subsequent experiment. On each trial, the simulator randomly sampled one pre-uStim state from the I experimental session datasets and passed it to the stimulation-response inverse model together with the target state. The stimulatoon-responee inverse model returned the stimulation suitabilities for each candidate electrode, and we chose the uStim pattern to "apply" using the method described in Section 2.5 The sampled pre-uStim state and selected uStim pattern were then passed to the stimulation response model to obtain the post-uStim state. Finally, the pre-uStim state, the resppatniionreeoeeeeee
inverse model (Section 2.4). In each
models. We set the target brain state to be one of the target states to
selected uStim pattern, and the post-uStim state were accumulated to perform the PPO-inspired adaptive update of the stimulation-respons inverse model, as described in Section 2.6
For each combination of hyperparameters, we performed R x T simulations (where each of the diimulationressonns models was used to produce post-uStim states for each of IT' different target states) and computed the average euclidian distance I-dimensional space between the predicted post-uStim state and the target state over all simulations. We performed Bayesian Optimization (using Optuna (25)) with this error objective and selected the hyperparameter values that achieved the smallest error. The ranges of hyperparameter values explored during Bayesian Optimization are summarized in the table below.
Method
Parameter
Description
Range of values
uStim pattern selection
iii
Initial value of 6j.k
0.3-0.7
ffoor
Minimum value of 6i,k
0.05-0.15
y
Discount factor of E
0.95-0.99
Adaptive model update
Eclip
Clip value of model update
0,1-0.3
Baseline for advantage computation
0,5-1.5
Batch size Learning rate
3-10
0.00001-0.001
S7 Details of model fitting and computing resources
The model architectures used in this study are depicted in Fig. ST Hyperparameters are listed in Sec All models were implemented in PyTorch (6 and fit using the AdamW optimizer. We trained all models on a local computing cluster using 4 NVIDIA GeForce RTX GPUs and 11GB of RAM. The same local computing cluster was used to run PPO-inspired adaptive model updates the stimulation-response inverse model during uStim experiments. Our experiments also involved three additional computers dedicated to specific tasks: one for experimental trial control (Intel Core i5-4590 @ 3.30 GHz x 4, 15.5 GB RAM, Intel HD Graphics 4600, Ubuntu 20.04.2 LTS), one for neural data recording (Intel Core 17-7700 Windows 10 Pro), and one for real-time silltimmeeee neural data processing and uStim pattern selection (11th Gen Intel Core i5- 1150 @ 2.70 GHz x 12, 15.4 GB RAM, AMD Radeon Pro WX 3200 + Intel Graphics, OMiSO implementation, the dimulltion-responss inverse model was
3.60 GHz, 4 GB AMD Radeon R7 450 + Intel HD Graphics 630,
Ubuntu 20.04.4 LTS). In adaptively updated every 5 trials using a PPO-inspired objective function on a local computing cluster, which different from the machines used run the experimental code. In this way, the adaptive model updates did not interfere with the stimulation pattern selection process, which is time-sensitive.
S8
Code availability
Python code for OMiSO is available on GitHub at http:///iiiibbooomyyyiiiiis/oSSs
16

---

A
uStim pattern only (CNN)
uStim pattern 10x10x1
C1 Relu 8x8x32
C2 Relu
Flattened 2,306
FC+ Relu 256
FC+ Relu
Output
6x6x64
Relu
64
B
uStim pattern Pre-stim5 (MLP)
uStim pattern
C1 + Relu
C2 + Relu 6x6x84
Flattened 2,306
FC+ Relu 256
FC+
FC+
10x10xx
8x8x32
Relu
Relu
64
32
Concatenated
FC 10
Output
Pre-uStim state 5x2
Flattened 10
FC Relu 16
FC Relu 32
C
uStim pattern + Pre-stim5 (LSTM)
uStim pattern 10x1 0x1
Relu 8x8x32
C2 Relu 6x6x64
Flattened 2,306
FC+
FC+
FC+
Relu
Relu
Relu 32
256
64
Concatenated 64
FC 32
FC 10
Output
Pre-uStim state 5x2
LSTM cell
Features 32
D
Stimulus-responss inverse model (MLP)
Pre-uStim state
10
FC + Relu
FC + Relu
FC + Relu
FC + Relu
Sigmoid 32
Concatenated 12
32
64
128
64
Output 20
Target state
2
Figure S1: Model architectures. (A) Stimulation-responss model architecture. uStim pattern only [10 (Fig. 23. Cl and C2 denote convolutional layers, and FC denotes fully connected layer, with ReLU activations applied throughout. (B) Stimulation-respone model architecture, uStim pattern + Pre-stim5 (MLP) (Fig. 2B.). (C) Stimulation ressoose model architecture, uStim pattern + Pre-stim5 inverse model architecture, MLP (Fig, 3 Fig. 4. The model architectures in (A)-CC) involve the same GNNN to process the uStim pattern input.
(LSTM) (Fig. 2 3) (D) Stimulation-respone

---

Medial
Anterior
Posterior
Lateral
Figure S2: Candidate uStim electrodes. The orange cells indicate the locations the 20 candidate electrodes chosen based on the criteria described in Section S3
Pre-uStim state (dim 2)
135'
316*
One trial
-2
Pre-uStim state (dim1)
Pre-uStim state (dim1)
Pre-uStim state (dim1)
Pre-uStim state (dim1)
No uStim
uStim pattern
uStim pattern
uStim
pattern
Post-uStim state (dim 2)
135
318'
-2
-2
-2
Post-uStim state (dim 1)
Past-uStim state (dim 1) Post-uStim state (dim 1) Post-uStim state (dim 1)
Figure S3: Pre-uStim states and subsequent post-uStim states induced by example uStim patterns. Pre-uStim and post-uStim states with additional uStim patterns tested during the same experimental session in Fig. FAg The row shows pre-uStim states, and the bottom row shows post-uStim states. The leftmost column shows brain states from no-uStim trials, while the remaining three columns show brain states from uStim trials with three distinct uStim patterns. Each uStim pattern represented as grid, with the stimulated electrodes indicated by orange cells. Same conventions in Fig. A As in Fig. 2A uStim shifted neural activity within the low-d space, and the induced post-uStim activity states depended on the memory condition with the degree varying by the uStim pattern.
18

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide a summary of the current literature and knowledge gap, and explain the contributions of our work in the abstract and introduction. Our claims are supported by the experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We discuss the limitations of this work in the third paragraph the discussion section.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer:
[NA]
Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We disclose the details of the proposed framework, experiments, data analyses, and modeling approaches needed to reproduce the main results. We also have released the code to reproduce the tested framework github.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
20

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have published the code and sample data needed to reproduce the tested framework github.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ https/ccpppblligggddeeeeeeeeeeee Coaee//dddeSebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The training and test details are provided in Section 2.3 and Section S7
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We include error bars in all figures and provide the outcomes of statistical tests for the main results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
21

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative еттоr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide details for the compute resources in Section S7
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't It make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:///wwrrppppcccppppplc/eeeee
Answer: [Yes]
Justification: We reviewed the NeurIPS Code of Ethics and our work is in compliance. All experimental procedures involved in this work were conducted in accordance with the United States National Research Council's 's Guide for the Care and Use of Laboratory Animals, and were approved by the Institutional Animal Care and Use Committee of Carnegie Mellon University.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdictio)..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the potential societal impacts in the last paragraph of the discussion section.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
22

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The data used in this study do not pose a risk for release as was not acquired from human subjects and does not contain images or identifiable information.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We cite all assets used in this paper.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
23

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The methods section provides the detail necessary to reproduce the proposed framework. We have also published the code and sample data on github.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This work does not involve human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The core components of this study do not involve the use of LLMs.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy pettts://neuriae /0tps///ececoooeeeeeee/2/2////// for what should or should be described.
25


Reviews:
[{"review_id": "TQMndMCak5", "reviewer": "Reviewer_1S96", "summary": "This paper introduces Online MicroStimulation Optimization (OMiSO), a closed-loop brain stimulation framework designed to modulate neural population activity more precisely by leveraging the brain’s ongoing state. Unlike prior methods, OMiSO conditions stimulation parameter selection on pre-stimulation neural activity and continuously adapts its model based on observed responses. The framework builds a stimulation-response model online and inverts it to choose stimulation parameters that drive neural activity toward specified target states. The approach was validated in a non-human primate experiment using intracortical microstimulation, where it outperformed baseline methods lacking brain state awareness or adaptive learning. While the method shows strong empirical results and is broadly applicable to different stimulation modalities, its effectiveness may vary depending on the stability of neural dynamics and consistency between training and deployment conditions. Overall, OMiSO presents a principled and flexible framework that enhances the precision of neural control, with promising implications for both neuroscience research and clinical neuromodulation.", "questions": "1. In Section 2.2, you mention that due to the large stimulation parameter space, only a small subset of possible patterns can be tested within a session. To address this, OMiSO first identifies a low-dimensional latent space of the high-dimensional neural population activity. Given this, could you clarify why you chose to model the latent space from spiking activity rather than directly leveraging low-dimensional biomarkers such as local field potentials (LFPs), which are commonly used in clinical settings? Additionally, do you have any experimental comparisons showing that OMiSO outperforms LFP-based stimulation methods?\n2. Could you elaborate on the key technical and application-level differences between OMiSO and previous methods such as MiSO? A clear comparison would help readers better understand the novelty and specific contributions of your approach.\n3. I would appreciate clarification on the connection between your monkey experiments and clinical deep brain stimulation (DBS). My understanding is that DBS is typically used in disease treatment contexts (e.g., Parkinson’s disease). How does your framework relate to that setting? Specifically, is the goal of OMiSO to generalize toward clinical applications, or is it primarily intended as a neuroscience tool for experimental control (e.g., guiding the monkey to desired neural states)?\n4. Could you comment on whether OMiSO aligns with model-based reinforcement learning paradigms?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## strengths\n1. This paper evaluates its proposed method on the real animal experiment, bridging the gap between machine learning framework and the real world application\n2. This work is well-motivated and application-oriented solving the challenge in neuromodulation\n\n## weakness\n1. Although authors try to provide a nice illustration in figure 1 and 2.1 overview, I still think it is not easily following for pure machine learning researchers.\n2. Limited novelty: Core components (Factor Analysis, behavior cloning, PPO) are not new and their integration is not clearly justified as novel.", "was_revised": true, "final_justification": "My original main concern is about their contribution. I think authors fully address my concern with technical and application-level differences. I would like to keep my positive recommendation.", "rebuttal": "Thank you to the reviewer for the constructive comments. Our responses are in the text below:\n\n> 1. Although authors try to provide a nice illustration in figure 1 and 2.1 overview, I still think it is not easily following for pure machine learning researchers.\n\nAgreed, we will make the following improvements in the revision:\n1. Add a mathematical problem statement to Section 2.1 to clearly frame the overall objective in machine learning terms:\n\nGiven a target brain state $z^{Targ}$, find a stimulation configuration $s$ that induces a brain state $z_s$ minimizing the error to the target:\n$$\n\\underset{s}{\\operatorname{argmin}} ||z^{Targ}-z_s||_1 \\\\\n$$\n\n2. Incorporate the notation introduced in the problem statement in Fig. 1 to clarify how the machine learning formulation maps to our specific neuroscience application.\n　\n> 2. Limited novelty: Core components (Factor Analysis, behavior cloning, PPO) are not new and their integration is not clearly justified as novel.\n\nOur contribution lies in the novel integration of these components and demonstrating that this combination is effective for neuroscience applications. Designing a framework that works reliably in neuroscience settings is non-trivial. It requires addressing domain-specific challenges such as inter-trial and inter-session neural variability (due to shifting brain state over time, recording instabilities, and changes in neural circuits), a limited number of neural recording samples, and interfacing with neural recording and stimulation devices in real time. Each component of OMiSO was selected after extensive experimentation to satisfy these constraints. FA was chosen based on its success in BCI applications by stably representing neural population activity across sessions. BC and PPO were selected for their robustness under real-time adaptive stimulation experiments even with a limited number of observations. While the individual components are not novel, our contribution is in the careful integration of them to satisfy the unique requirements in neuroscience applications.\n　　\n> 1. In Section 2.2, you mention that due to the large stimulation parameter space, only a small subset of possible patterns can be tested within a session. To address this, OMiSO first identifies a low-dimensional latent space of the high-dimensional neural population activity. Given this, could you clarify why you chose to model the latent space from spiking activity rather than directly leveraging low-dimensional biomarkers such as local field potentials (LFPs), which are commonly used in clinical settings? Additionally, do you have any experimental comparisons showing that OMiSO outperforms LFP-based stimulation methods?\n\nOur current work focuses on spiking activity because it provides richer and higher-resolution neural signals than LFPs. Compared to LFPs, spiking activity tends to be less spatially correlated across nearby electrodes and can capture more detailed information about underlying neural computations (Kajikawa et al., 2011).\n\nAt the same time, we agree that LFPs are highly relevant, especially in clinical settings. While our current experiments focus on spiking data, the OMiSO framework is generalizable and could be applied to control LFP activity as well. Specifically, one could train a stimulation-response model to predict evoked LFPs from different stimulation patterns, and then use the same behavior cloning and PPO-based online update strategy to perform adaptive stimulation experiments.\n　　\n> 2. Could you elaborate on the key technical and application-level differences between OMiSO and previous methods such as MiSO? A clear comparison would help readers better understand the novelty and specific contributions of your approach.\n\nWe summarize here the key technical and application-level differences between OMiSO and MiSO below, which we will bring out more clearly in our paper.\n\n**Technical differences**:\n1. *Incorporation of brain state*: MiSO does not consider the brain state when predicting the brain’s response to stimulation. In contrast, OMiSO explicitly models the brain state, allowing it to capture trial-to-trial differences in pre-stimulation brain state. This results in more precise prediction and control of neural population activity.\n2. *Sample efficient online model*: MiSO uses a tabular approach that stores the predicted post-stimulation state for each candidate stimulation pattern and updates only the entry corresponding to the tested pattern. This approach is sample inefficient as it does not propagate what can be learned from a new observation across patterns. OMiSO leverages an inverse model and updates its parameters online, enabling learning beyond the tested stimulation pattern. Consequently, OMiSO learns more efficiently and can adapt quickly with fewer observations.\n3. *Scalability*: The tabular method in MiSO becomes impractical as the stimulation parameter space grows, due to the increase in table size. OMiSO’s model-based approach scales more effectively to richer parameter spaces. This scalability allowed us to extend stimulation parameter space to the 5-electrode stimulation, compared to the 2-electrode stimulation explored in MiSO.\n\n**Application-level differences**:\n1. *Dependency on the brain state*: In applications which use strong stimulation (e.g., large currents), the brain responses to stimulation might not depend on the pre-stimulation brain state (e.g., all the neurons are excited or suppressed). In such cases, MiSO may be sufficient. However, when stimulation is used to modulate neural activity, which is likely needed in most applications, the brain’s response to stimulation will depend on the pre-stimulation brain state, as shown in Fig. 2. In this case, OMiSO would work better than MiSO.\n2. *Experimental session duration*: If the experimental session is short (e.g., a few minutes) and the brain state remains stable throughout, MiSO may be sufficient. On the other hand, if the experimental session is longer than a few minutes (likely for most applications) and there are shifts in brain state, OMiSO may be necessary to compensate for the shift.\n3. *Computation speed requirements*: For applications that require selecting a stimulation pattern within a very short time window, MiSO may be a good choice. However, if longer computation times are acceptable or a fast GPU infrastructure is available, OMiSO can provide more precise control of neural population activity.\n\n> 3. I would appreciate clarification on the connection between your monkey experiments and clinical deep brain stimulation (DBS). My understanding is that DBS is typically used in disease treatment contexts (e.g., Parkinson’s disease). How does your framework relate to that setting? Specifically, is the goal of OMiSO to generalize toward clinical applications, or is it primarily intended as a neuroscience tool for experimental control (e.g., guiding the monkey to desired neural states)?\n\nOur goal includes both clinical and neuroscience research applications. We designed OMiSO as a framework that can be adapted to both contexts by modifying the target brain state and by defining the appropriateness of a given stimulation pattern for achieving the target state (eq. 7).\n \nIn clinical DBS applications, for example, the target could be defined as a healthy brain state. OMiSO can then identify stimulation patterns that move neural activity towards this desired state. OMiSO could also be used to stay away from unhealthy brain states by maximizing the distance from unhealthy brain states.\n\nIn neuroscience research applications, OMiSO can be used as a tool for different stimulation modalities (electrical microstimulation, optogenetic stimulation, etc) to study sensory, cognitive, and motor function. For instance, researchers can perturb neural activity away from one working memory state to another working memory state to investigate how the neural population activity evolves after such perturbation.\n\n> 4. Could you comment on whether OMiSO aligns with model-based reinforcement learning paradigms?\n\nOMiSO does not use a model of the environment (i.e., the brain), and thus, it is not a model-based reinforcement learning approach. However, this is a great point, and we are working on extending OMiSO toward a model-based reinforcement learning framework. The current version of OMiSO focuses on a single time-step stimulation to achieve a target brain state at a single time point. However, many behaviors and brain states involve neural activity evolving over time. Thus, extending OMiSO to control neural population activity over time may be an important direction to affect such behaviors or brain states. Such an extension requires a model that incorporates stimulation inputs and predicts future brain states, thereby guiding the selection of stimulation patterns over multiple time steps."}, {"review_id": "wVIKdrY0kY", "reviewer": "Reviewer_wb75", "summary": "The authors introduce an algorithm to identify desired stimulations in order to move brain activity to a target state, in a closed loop. \nThe key ingredients involve learning an inverse mapping that outputs the desired stimulation pattern required to reach a target state from the previous state of neural activity, and then updating this mapping online as more data is collected during the experiment. They show that taking into account the current state of the brain in this pipeline results in far better performance in terms of reaching the target state relative to existing methods such as MiSo. In particular, results from stimulations in macaque PFC recordings using Utah array show successful reaching of target state in a low-dimensional space.", "questions": "Aside from some of my questions listed in the weakness section, I have the following questions / comments:\n\n1. In the introduction, the authors don't make it clear what \"brain state\" means. Later on it is clear that the authors are referring to a low-d projection of neural activity, but through the introduction it is unclear whether this refers to a high-level cognitive state etc, and it would help the readers a lot to know this upfront.\n\n2. Forward model: In line 118, what kind of function does $f$ represent? The supplement suggests these are deep networks, which would be helpful to clarify here.\n\n3. What if a reaching a target brain state requires multiple steps of stimulations? How would this algorithm handle such a case given intermediate target states may be unknown in that setting?\n\nMinor comments / typos:\n1. Lines 36-39: this sentence needs to be broken into two for readability.\n2. Line 83: I am assuming this should be $e^i \\in \\\\{1,...n_i\\\\}$.", "limitations": "yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written, key claims and results are easily to follow.\n2. It attempts to solve an impactful and translatable problem, and is extremely significant to research at the cross roads of machine learning and neuroscience. \n3. The methodology overall makes sense, the high-level pipeline is novel, and the experiments are thoroughly done on real macaques.\n4.  I also appreciate the supplement discussion re how the online updates are done without hindering the generation of stimulation patterns themselves. \n\nWeaknesses:\n1. My primary concern is that reinforcement learning terminology seems to be used very loosely here. While that doesn't overshadow the main contributions of the paper, this being a machine learning venue does require either further explaining or rewriting. For example, the inverse model training is really  a supervised problem on synthetic data generated from the forward model. Calling it behavior cloning is confusing in my opinion, because it is unclear what the markov decision process looks like here. In its current form \"state\" in the MDP sense includes past + current ($z_{pre}$) and future ($z_{post}$) information. If we think of the action as the stimulation, what should the next markovian state look like? Is there a sequential problem underlying this process at all that I might be missing? Relatedly, PPO also seems to be loosely used, where the advantage function typically is defined as $A(s,a) = Q(s,a)-V(s)$. I am completely fine with the authors saying that they update the model using a PPO-inspired loss function so that the inverse model doesn't deviate much between consecutive trials. Alternatively, I would like a concrete definition of the RL problem if the authors wants to stick to referring to these approaches as behavior cloning and PPO. \n\n2. Some of the choices made during the algorithm are unclear from the main text / methodology, making it hard to fully wrap one's head around it. Crucially, after learning the forward model on low-dimensional projection of the data, the authors further select $l$ of those dimensions to learn the inverse model. The rationale behind this is unclear to me. Is this is a computational bottleneck? Presumably, factor analysis already compresses neural data to a low-dimensional space, is this further compression necessary as it presumably results in not being able to control certain neural dimensions?", "was_revised": true, "final_justification": "I am keeping my original score. The authors have agreed to address my concern regarding RL terminology, otherwise the paper presents an interesting and important approach.", "rebuttal": "Thank you to the reviewer for the constructive comments. Our responses are in the text below:\n\n> 1. My primary concern is that reinforcement learning terminology seems to be used very loosely here. While that doesn't overshadow the main contributions of the paper, this being a machine learning venue does require either further explaining or rewriting. …..I am completely fine with the authors saying that they update the model using a PPO-inspired loss function so that the inverse model doesn't deviate much between consecutive trials. Alternatively, I would like a concrete definition of the RL problem if the authors want to stick to referring to these approaches as behavior cloning and PPO.\n\nWe appreciate the reviewer highlighting this point. While the design of OMiSO was inspired by RL, we agree that RL terminology was used loosely in the writing. In the revision, we will explain OMiSO without relying on RL concepts. Specifically, we will describe the \"behavioral cloning\" component as a supervised learning problem trained on synthetic data, as suggested by the reviewer. For the online update procedure, we will refer to the objective as a PPO-inspired loss function, as the reviewer suggests. These edits should provide a more accurate description of the method while avoiding confusion related to RL terminology.\n\n> 2. Some of the choices made during the algorithm are unclear from the main text / methodology, making it hard to fully wrap one's head around it. Crucially, after learning the forward model on low-dimensional projection of the data, the authors further select of those dimensions to learn the inverse model. The rationale behind this is unclear to me. Is this is a computational bottleneck? Presumably, factor analysis already compresses neural data to a low-dimensional space, is this further compression necessary as it presumably results in not being able to control certain neural dimensions?\n\nIn principle, OMiSO can operate by defining the target state in all $m$ dimensions of the latent space. However, there are two key reasons why users might prefer to define the target in a lower-dimensional subspace, as we did in our experiments.\n \nFirstly, not all latent dimensions are equally suitable for reliable perturbation of neural activity. For example, some dimensions may be aligned poorly across sessions because they are not consistently present in every session, while others may not be directly modifiable by stimulation. Thus, in our work, to evaluate the optimization performance of OMiSO, we selected target latent dimensions satisfying two conditions: 1) they are well-aligned across multiple sessions, and 2) uStim can induce diverse latent activity states along these dimensions (Section S3, line 457). For 1), we computed the cosine similarity of each column vector of $\\tilde{\\Lambda_i}$ (eq. 3) between a reference session and all other training sessions, which quantifies the consistency of each aligned latent dimension across sessions. For 2), we visualized the post-uStim state (eq. 4) induced in each single trial during the training sessions to assess the diversity of the induced latent activity.\n\nSecondly, some scientific or clinical applications may require only a subset of dimensions. Previous studies investigated output-potent and output-null dimensions in the latent space, where only certain dimensions directly influence behavior or downstream processes (e.g., Kaufman et al., 2014). For instance, previous studies on working memory suggest that specific latent dimensions stably encode memory states, even though overall neural activity continues to evolve dynamically (Murray et al., 2017; Parthasarathy et al., 2019; Soldado-Magraner et al., 2025). While five latent dimensions may optimally capture overall neural population activity, certain brain states such as working memory states may primarily rely on just one or two specific dimensions. In these cases, we may want to focus on those relevant dimensions to affect the brain state. \n\nWe will revise our paper to include these points in the main text.\n\n> 1. In the introduction, the authors don't make it clear what \"brain state\" means. Later on it is clear that the authors are referring to a low-d projection of neural activity, but through the introduction it is unclear whether this refers to a high-level cognitive state etc, and it would help the readers a lot to know this upfront.\n\nAgreed, we will add a clarification that “brain state” refers to low-d projection of recorded neural population activity in this study.\n\n> 2. Forward model: In line 118, what kind of function does represent? The supplement suggests these are deep networks, which would be helpful to clarify here.\n\nIn our animal experiments, we used a combination of a CNN and a RNN (Section S1) as the function $f$. We kept the description in the main text general (i.e., any function $f$) to emphasize the flexibility of our framework because the stimulation-response model architecture could, in principle, be any architecture depending on the application and data characteristics. However, we agree that specifying the architecture used in our experiments would improve clarity. We will revise the main text to explicitly state the architecture we used in this study.\n\n> 3. What if a reaching a target brain state requires multiple steps of stimulations? How would this algorithm handle such a case given intermediate target states may be unknown in that setting?\n\nExtending OMiSO to handle multi-step stimulations is an important direction, especially given the dynamic nature of neural activity, and we are actively exploring this extension. OMiSO can accommodate multi-step stimulation scenarios in two ways.\n \nOne way is to iterate stimulation and observation within an experimental trial. In this case, OMiSO can apply a stimulation pattern, observe the resulting brain state, and select the next stimulation pattern based on the induced state. The induced brain state from the previous stimulation pattern will be the pre-stimulation brain state to guide selection of the next stimulation pattern.\n\nAnother way is to optimize a sequence of different stimulation patterns. In this case, the parameter space becomes the space of stimulation sequences rather than single patterns. OMiSO can be applied to this expanded parameter space by obtaining training data with such sequences. To encode these sequences, we could update the stimulation-response model architecture (e.g., by using multiple input channels with CNN) to predict the cumulative effect of a stimulation sequence. However, the reviewer is right that the inability to access the brain state during stimulation due to stimulation artifacts may limit the performance of OMiSO. This could be addressed by using stimulation methods that allow simultaneous recording and stimulation, such as optogenetic stimulation or spatially separating the stimulation and recording sites (instead of using the same electrode array for both stimulation and recording).\n\n> Lines 36-39: this sentence needs to be broken into two for readability.\n\nAgreed, we will update it.\n\n> Line 83: I am assuming this should be $e^i$ ∈ {1,...,$n_i$}\n\n$e_i$ is a list of $n_i$ usable electrodes rather than an integer. In this project, we implanted 96 electrodes, each of which is assigned an index from 1 to 96. Among them, we identified a subset of usable electrodes for each session. We denoted this subset as $e_i$, which is a list of indices of $n_i$ usable electrodes (e.g., [1,4,10,32,..,93]).  We see how this could be confusing, and we will make it more explicit."}, {"review_id": "IuwBP3AwKM", "reviewer": "Reviewer_moJh", "summary": "The authors developed a method called OMiSO, which, given the current neural population state and the desired state, outputs an optimal stimulation pattern that will move the state to the desired state. In this study, the authors focus on not the full neural state space, but the low-dimensional latent state obtained from factor analysis. Factor analysis is performed for each session, and the latent spaces are aligned, following the same approach as in MiSO. The authors then fit a stimulation-response model that predicts the latent state after stimulation, given the latent state before stimulation and the stimulation pattern. This model is then used to generate enough data to fit an RL model via behavioral cloning. The RL model’s job is to predict, given the pre-stimulation latent state and the target latent state, what the stimulus pattern should be. After these steps, which are so far done offline, the authors further refine the RL model online via PPO as new data comes in. The authors test their method on a real monkey electrophysiology dataset and find superior performance compared to baselines and a recent method, MiSO.", "questions": "Please see above for questions.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Quality**: The authors’ claims are well supported by experiments. \n\n1. The authors claim that leveraging pre-stimulation neural activity for predicting optimal stimulation patterns can be helpful, and show experimental evidence that the latent states of the monkey before and after stimulation for two different stimuli are distinguishable. I agree that providing the model with the pre-stimulation activity (more data) is helpful, and wonder if, for their setup, providing it with experimental stimuli (the location of the target dots) can make it even better. What do the authors think about including “optional data” (e.g., target location in this case, but could be some other data in general) as one of the features of the model?\n\n2. The authors show evidence that their stimulation-response model works well, by experimenting with different neural network architectures. The authors report the hyperparameters in the Supplementary Material. The authors also justify their stimulation-response inverse model by comparing it to a “perfect” inversion model, which constructs a table of all possible combinations of stimulation patterns and their corresponding states post-stimulation. This perfect inversion model predicts the optimal stimulation patterns from the post-stimulation states using this table (a table including pre-stimulation combinations may be infeasible due to the latent space being continuous and large). Their inverse model (trained by BC) performs well compared to the perfect inversion model.\n\n3. Lastly, the authors show evidence that refining their predictions online via PPO helps in actual brain stimulation experiments. 14 sessions were used to make comparisons between MiSO, OMiSO, and their perfect inversion model without pre-stimulation states. The authors show that OMiSO performs better than the baseline methods in this experiment.\n\n\n**Clarity**: The work is overall clear and well-organized. There were a few places where I thought would benefit from clarification.\n\n1. Equations (1-4) are exactly the same as Equations (1-4) of the MiSO paper, and as far as I understand, there is no conceptual difference between the latent alignment procedure performed by MiSO and that performed in this work. I think the authors could instead cite MiSO for this Section, with just Eq. 4 in the paper, as this one is used throughout the later Sections.\n\n2. It wasn’t clear from the main text how the function f (the stimulation-response model) was modeled when it was first introduced in L118. Adding a sentence early on that the network architectures in Fig. 2B are used for their function f, and that the specific architectures are shown in Fig. S1, could be helpful. This is also true for \\phi_\\theta, first introduced in L128. Perhaps mentioning that this is an MLP early on could be helpful.\n\n3. It wasn’t clear in the beginning why Z is a matrix with t, whereas z is a vector without t? If I understand correctly, this is because for Z, the stimulation-response model f can use 1 to 5 pre-uStim time bins. I think this information could be added early on to make it clearer.\n\n4. The authors mention that “OMiSO performs R-fold cross validation where every session is used as a validation set once (see Section S4)”. S4 didn’t specify this, and perhaps mentioning S6 in S4 makes sense.\n\n5. It wasn’t clear from the main text how the experiments were run -- for each trial, how did they choose which method to use? For Fig. 4D, the authors show the average predicted post-stimulation state and target state, along with single-trial predictions of post-stimulation states. Is the plot in Fig. 4E made by first taking the Euclidean distance between the average predicted state and the target state for each session, and then computing the mean and standard error across the 14 sessions? For L312, should this be p=0.05 and p=0.01, and not p<0.05 and p<0.01?\n\n6. L157: The number of stimulus patterns is not in real space but in integer space.\n\n7. L195: What does it mean for the baseline b to be used “to evaluate the induced post-uStim activity”? How does this baseline b help here?\n\n8. After doing dimensionality reduction to m=5 dimensions, how does the user select l from m? How m and l were selected was not explained in S3 but in S1 (only for m). If l components are selected from m, I think the same state in this l-dimensional latent space doesn’t represent the same space in neural space or the m-dimensional latent space. Would results be different if the authors used l=m=5, or l=m=2? I think a justification for subselection would be helpful. How does increasing or decreasing l or m affect model performance (both accuracy and speed)? \n\n9. Consistent terminology could help: for example, sometimes the stimulation-response model is referred to as the statistical model, and pre-uStim is referred to as pre-stimulation.\n\n10. There is no Fig. 2C (L268).\n\n11. MiSO also performs online inference, but this method is called Online MiSO. Is the online-ness the most important factor that distinguishes between this approach and MiSO? A few sentences discussing in-depth how this approach differs from MiSO would be helpful.\n\n**Significance**: Optimal online control of the neural population state from one state to another is a key challenge in BCI, and OMiSO suggests a new approach to solve this problem. The authors apply their method to real brain stimulation experiments in monkeys and demonstrate its efficacy in comparison to a previous method, MiSO. Their demonstration of their method and baselines in actual experimental settings is a big strength of this paper. The method also consists of modular steps, and each step of the method can be replaced without affecting the other steps. For example, in the paper, factor analysis and the latent alignment method via Procrustes optimization were used in one of the steps, and this may be replaced by other fast dimensionality reduction methods and alignment methods in the future if one finds improvement over the current one. This modularity allows accessing each step of the model and what features most affect its performance.\n\n**Originality**: While there is previous work incorporating pre-stimulation neural state for choosing optimal stimulation patterns, the authors suggest that this is the first to incorporate pre-stimulation activity for manipulating high-dimensional neural population state. Because this work focuses on the prediction of the low-dimensional latent state, I think it would help to show prediction in the neural population state (i.e., predictions on firing rates of individual neurons) are close to the observed high-dimensional neural population state. Currently, the authors focus on the learned low-dimensional state, and it is not clear whether the learned low-dimensional state accurately captures the high-dimensional neural population state. As mentioned in Clarity point 8 above, it may also be the case that if the users subselect arbitrary dimensions in the m-dimensional latent space, it would be hard to recover the full high-dimensional neural population state. A discussion on how this can be addressed would be helpful. That using PPO to update/refine the predictions in real-time is better than the method used in MiSO is also an interesting finding. Is there a more direct way to incorporate pre-stimulation neural state into the MiSO framework and compare that to OMiSO?", "was_revised": true, "final_justification": "The authors' response properly addressed my questions/concerns. Thus, I would like to keep my score.", "rebuttal": "Thank you to the reviewer for the constructive comments and careful reading. Our responses are in the text below:\n> 1. The authors claim...\n\nWe appreciate the suggestion. Although we did not incorporate the location of the target dots as features, we agree that incorporating it could improve model prediction. More broadly, this direction opens up an interesting opportunity to combine electrical stimulation with, for example, visual stimuli to not only improve the prediction precision but also increase the diversity of achievable neural population activity states. The use of visual stimuli to induce diverse neural population activity states has been explored in recent work (e.g., Cowley et al., 2017; Bashivan et al., 2019; Walker et al., 2019), which demonstrated that optimization of visual stimuli can shape neural population activity. Incorporating such stimuli into our framework is an interesting direction for future research.\n> 1. Equations (1-4) are...\n\nThe reviewer is right that there is no conceptual difference in the latent alignment procedure between MiSO and OMiSO. We will make this section more concise by citing MiSO paper.\n> 2. It wasn’t clear...\n\nThis is a great point. We will add a sentence in Section 2.3 to explain the model and architecture we used for our validation. We will also add a similar sentence in Section 2.4 for the stimulus-response inverse model.\n> 3. It wasn’t clear in the...\n\nYes, the reviewer’s understanding is correct. While $Z$ contains multiple timesteps of latent activity making it a matrix, $\\mathbf{z}$ is a vector containing only one timestep of latent activity. In revision, we will clarify the difference between them when introducing the $Z$ matrix.\n> 4. The authors mention...\n\nThank you for pointing this out. We will cite S6 instead of S4.\n> 5. It wasn’t clear...\n\nFor each trial, our experimental system randomly selected which method to use. In Fig. 4D, each dot represents the _actual_ neural activity induced by the selected stimulation pattern by the selected method on a single trial. The circled dot in each panel indicates the average induced activity across trials for that method within an example session.\n \nTo generate Fig. 4E, for each session, we first computed the Euclidean distance between the target state and the induced neural activity for each trial (i.e., each dot in Fig. 4D), and then averaged these distances across trials separately for each method. For each method, we then computed the mean and standard deviation of these session-level averages across the 14 sessions.\n> For L312, should...\n\nYes, they are p=0.05 and p=0.01. More specifically, they are p=0.048 and p=0.011.\n> 6. L157: The number...\n\nWe will change the notation from $\\\\mathbb{R}$ to $\\\\mathbb{Z}$.\n> 7. L195: What does...\n\nThe baseline b ensures the advantage can take both positive and negative values. Without it (i.e., b=0), the advantage (eq. 7) becomes always negative. This would cause the model parameter updates to consistently reduce the predicted suitability of all stimulation patterns, leading to unstable learning. With baseline b, the advantage $A_{c,k}$ becomes positive for stimulation patterns with smaller errors than b and negative for those with larger errors. This encourages increasing the probability of better-performing patterns and decreasing that of worse ones, helping the model converge more stably and quickly toward small-error stimulation patterns.\n> 8. After doing...\n\nWe selected _l_ target latent dimensions satisfying two conditions: 1) they aligned well across multiple sessions, and 2) diverse latent activity states could be induced along these dimensions with uStim (S3, line 457). We made this choice because not all latent dimensions are equally suitable for reliable perturbation of neural activity. For example, some dimensions may be aligned poorly across sessions because they are not consistently present in every session, while others may not be directly modifiable by uStim.\n\nIf we set l=m=5, the optimization problem becomes more challenging because it may require optimizing over a larger parameter space (e.g., optimizing stimulation amplitude, in addition to stimulation locations) to find the stimulation pattern which can create a target state along all 5 dimensions. We expect OMiSO to still outperform MiSO in this scenario because OMiSO is designed to explore a larger parameter space to produce more diverse brain states and leverage brain state information to improve prediction accuracy.\n \nIf we set l=m=2, the alignment across sessions may deteriorate, leading to lower performance for both MiSO and OMiSO. 5 dimensions were identified as the optimal latent dimensionality that captures variance shared among the recorded neurons consistently across sessions. Limiting to 2 dimensions could result in inconsistent identification of the latent dimensions across sessions, thereby limiting alignment quality.\n\nSome applications may require only a subset of dimensions. For instance, previous studies on working memory suggest that specific latent dimensions stably encode memory states, even though overall neural activity continues to evolve dynamically (e.g., Murray et al., 2017). While five latent dimensions may optimally capture overall neural population activity, certain brain states such as working memory states may primarily rely on just one or two specific dimensions. In these cases, we may want to focus on those relevant dimensions to affect the brain state. We will revise our paper to include these points in the main text.\n> 9. Consistent terminology...\n\nAgreed about consistent terminology. For the stimulation-response model, we will define the stimulation-response model in the introduction and consistently use it throughout the paper. \nFor the use of pre-stimulation and pre-uStim, we used pre-uStim when referring specifically to electrical microstimulation, and pre-stimulation when including other types of stimulation (e.g., optogenetics).\n> 10. There is no Fig. 2C...\n\nThanks, we will update it to Fig. 2B.\n> 11. MiSO also performs...\n\nWe summarize the differences between OMiSO and MiSO below, which we will bring out more clearly in our paper.\n\n1. MiSO does not consider the brain state when predicting the brain’s response to stimulation. In contrast, OMiSO explicitly models the brain state, allowing it to capture trial-to-trial differences in pre-stimulation brain state. This results in more precise prediction and control of neural population activity.\n2. MiSO uses a tabular approach that stores the predicted post-stimulation state for each candidate stimulation pattern and updates only the entry corresponding to the tested pattern. This approach is sample inefficient as it does not propagate what can be learned from a new observation across patterns. OMiSO leverages an inverse model and updates its parameters online, enabling learning beyond the tested stimulation pattern. Consequently, OMiSO learns more efficiently and can adapt quickly with fewer observations.\n3. The tabular method in MiSO becomes impractical as the stimulation parameter space grows, due to the increase in table size. OMiSO’s model-based approach scales more effectively to richer parameter spaces. This scalability allowed us to extend uStim parameter space to the 5-electrode uStim, compared to the 2-electrode uStim explored in MiSO.\n\nWe will update the paper to clarify these innovations. While OMiSO introduces multiple innovations, we chose to highlight the online model component in its name, as it played the most critical role in both incorporating brain state information and scaling the method to richer stimulation spaces.\n> Originality: While there...\n\nThe reviewer raises a great point. We will update the paper to clarify that, while previous deep brain stimulation (DBS) studies have focused on manipulating low-dimensional biomarkers such as LFPs, our work introduces a novel approach to manipulate low-dimensional representations of multi-dimensional neural population activity. Although both approaches operate in low-dimensional spaces, our method makes a significant advancement because compared to LFPs, spiking activity tends to be less spatially correlated across nearby electrodes and can capture more detailed information about underlying neural computations. This enables us to achieve more precise control of neural population activity.\n\nOur choice to focus on the low-d representation of neural population activity is motivated by many recent studies showing that the essential changes in neural activity underlying sensory, cognitive, and motor processes can be captured by low-d spaces (see Cunningham & Yu 2014, Saxena & Cunningham 2019, Ebitz & Hayden 2021). Furthermore, in brain-computer interfaces, it is possible to accurately decode movement intent from low-d representations of population activity (e.g., Degenhart et al., 2020). These insights inspired the design of OMiSO to operate within a low-d space.  While we show OMiSO can shape neural activity within the latent space, future work can test whether doing so produces meaningful changes in behavior. The success in this direction could be a major step forward for both clinical and scientific applications. \n> That using PPO...\n\nOne way to incorporate pre-stimulation neural states in MiSO is by discretization. For example, we could divide the latent neural activity state along each latent dimension into five discrete bins. This approach may work when the latent dimensionality is small. However, it becomes quickly infeasible as the dimensionality increases. For instance, with 15,504 stimulation patterns and a 5-d latent space discretized into 5 bins per dimension, we would need to optimize a table with 387,600 entries (15,504x5×5). Optimizing such a large table with only a few hundred experimental trials is extremely challenging. In contrast, OMiSO can efficiently handle continuous pre-stimulation latent states, offering a more scalable method."}, {"review_id": "4XblmgCJPV", "reviewer": "Reviewer_oMHT", "summary": "This paper introduces OMiSO (Online MicroStimulation Optimization), a brain stimulation framework that optimizes stimulation parameters by explicitly modeling the pre-stimulation neural state to achieve precise control over population-level neural activity. This work proposes a closed-loop, state-aware optimization approach, which dynamically adapts stimulation protocols based on real-time neural recordings.", "questions": "See Weaknesses.", "limitations": "See Weaknesses.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The framework introduces optimal parameter selection conditioned on pre-stimulation brain states—a well-established control-theoretic principle previously unrealized in MiSO systems. This enables context-aware stimulation tailored to baseline neural activity.\n\n2. By continuously refining parameters using newly observed stimulation responses, OMiSO overcomes the limitation of MiSO's restricted parameter-space search. The proposed reinforcement learning algorithm dynamically optimizes stimulation protocols, achieving superior neural control.\n\nWeaknesses:\n1. While the paper improves upon MiSO, it inherits fundamental limitations—notably, the low-dimensional stimulation parameter space (constrained by dual-electrode setups) inherently restricts the spatial scope of neuromodulation. The work focuses on algorithmic refinements rather than addressing these hardware-imposed boundaries, leaving the core challenge unresolved.\n\n2. The chosen metrics (e.g., accuracy in reaching target states) may not fully capture the biological efficacy of neural control, especially given the dynamic nature of brain states.\n\n3. Due to inter-session variability in neural activity, the claimed optimization benefits require validation across totally unseen sessions/timepoints to demonstrate robustness. The authors should discuss strategies to mitigate session-dependent confounds.", "was_revised": true, "final_justification": "Although the authors‘ explanation couldn't fully convince me, this paper doesn't have major flaws.", "rebuttal": "Thank you to the reviewer for the constructive comments. Our responses are in the text below:\n\n> 1. While the paper improves upon MiSO, it inherits fundamental limitations—notably, the low-dimensional stimulation parameter space (constrained by dual-electrode setups) inherently restricts the spatial scope of neuromodulation. The work focuses on algorithmic refinements rather than addressing these hardware-imposed boundaries, leaving the core challenge unresolved.\n\nA substantial advancement of this work over MiSO is that we move beyond the dual-electrode uStim setup. In this study, we applied uStim using up to 5 electrodes. In principle, our uStim device is capable of supporting much more flexible stimulation configurations, including arbitrary combinations of electrodes (up to the 96 available with our hardware), as well as different amplitude, frequency, and waveform choices. The main constraints on our stimulation parameter space stem from biological safety limits, specifically, the total current that can be safely delivered, rather than limitations imposed by hardware. We agree with the reviewer that future developments in stimulation hardware (e.g., devices that allow precise control of current return paths or significantly lower latency for reconfiguration) is another critical field to expand the capability of neuromodulation technology. We believe that OMiSO is well-suited to take advantage of such advances, as our framework is generalizable to a wide range of stimulation parameter spaces.\n\n> 2. The chosen metrics (e.g., accuracy in reaching target states) may not fully capture the biological efficacy of neural control, especially given the dynamic nature of brain states.\n\nWe agree. To manipulate behavior or brain state, it may be important to create a sequence of neural activity evolving over time (i.e., a target neural trajectory), rather than achieving a single static target state because many of the brain's computations are believed to rely on particular time courses of neural activity (e.g., decision making). Taking into account the time courses of neural activity is precisely where we are heading next, and our ongoing goal is to extend OMiSO to account for the temporal dynamics of brain activity. This could be achieved, for instance, by incorporating a dynamical model of the brain that predicts how neural activity will evolve over time. The prediction from this model guides the stimulation pattern selection process by suggesting what neural activity state to achieve at each time point to achieve a desired evolution over time.\n\n> 3. Due to inter-session variability in neural activity, the claimed optimization benefits require validation across totally unseen sessions/timepoints to demonstrate robustness. The authors should discuss strategies to mitigate session-dependent confounds.\n\nWe agree that to demonstrate robustness, we need strategies to mitigate session-dependent confounds to generalize to a new session. This is particularly important in neuroscience applications, where the recorded neural activity can be non-stationary across days due to small movements of the recording electrodes relative to the brain and changes in the neural circuits.  As a result, a model that performs well in one session may fail in another. Stimulation frameworks need to be able to overcome this variability. To address this challenge and demonstrate the robustness of OMiSO, we designed our experiments and analyses in the following ways:\n\n1. We trained the stimulation-response model using all but one experimental session and evaluated it on the held-out session (Fig. 2B, Section S6). This ensured that the model was evaluated on a totally unseen session.\n2. We trained an inverse model using data from training sessions and tested its performance in totally unseen sessions (Fig. 3D, Section S6).\n3. We tested the robustness of OMiSO by running adaptive stimulation optimization in totally unseen sessions not used during the inverse model training (Fig. 4E, Section S6).\n\nIn all these evaluations, we ensured that the test sessions were entirely held out from training. The consistent performance across these unseen sessions provided strong evidence that OMiSO can generalize across sessions and is robust to inter-session variability."}]

Author rebuttal:

Meta-review: This paper introduces a new approach for online experimental design for directed brain stimulation. It's a clear and well-specified problem, in a challenging domain; all reviewers were impressed by the strength of the validation, which includes exploring different model architectures, and also real-world lab experiments that demonstrate the approach actually working in a real-world scenario. There was a consensus to accept among all reviewers, though borderline for most.

The major concerns were around clarity and scope of the machine learning contributions; this has been well-discussed in the author rebuttal period. For the final version of the paper, it should incorporate aspects discussed here, including making connections to RL and PPO clearer, making it more accessible to machine learning researchers (e.g. by clarifying the representation of the brain state early on), and making clear the connections to MiSO in the background sections. In particular all the "clarity" issues from reviewer moJh should be addressed.

One reviewer had a more fundamental concern regarding the hardware limitations of the approach used here, but not strongly enough to discourage acceptance.
Decision: Accept (poster)