Neural-DISCO: Source-Conditioned Counterfactual Editing of Neural Population Activity

# Neural-DISCO: Source-Conditioned Counterfactual Editing of Neural Population Activity
Source: [https://arxiv.org/html/2610.33145](https://arxiv.org/html/2610.33145)
James Ho11footnotemark:1Joseph Soo11footnotemark:1Xin Tang††thanks:Corresponding author:xin\.tang@msl\.ubc\.ca\.Affiliation:Michael Smith Laboratories and Department of Computer ScienceAffiliation:University of British ColumbiaAffiliation:Vancouver, BC, Canada

###### Abstract

Understanding how neural activity encodes behavioral and sensory variables is a key challenge in systems neuroscience\. This requires models that go beyond predicting activity from labels and instead support source\-conditioned counterfactual editing: given a recorded source response, we should be able to edit one measured behavioral or stimulus label and predict how neural activity would change while preserving other observed labels and source\-specific variability not explained by the labels\. Here, we introduceNeural\-DISCO, a disentangled representation learning framework for counterfactual modeling of neural activity\. Neural\-DISCO partitions the latent space into label\-specific components, each associated with an observed behavioral or stimulus label, together with a residual latent that captures variability not supplied through observed labels\. Across simulated and real neural recording datasets, we show that this structure enables selective editing of measured labels to generate counterfactual neural responses while preserving the remaining content of the original activity\. Further, to gain insight into how behavioral and stimulus labels are encoded at single\-neuron resolution, we pair this disentangled framework with feature attribution methods that identify which neurons are most important for each label\. We find that disentanglement enhances the recovery of known functional cell roles in both synthetic and real data\. Together, these results support Neural\-DISCO as a framework for counterfactual modeling of neural activity and provide a practical interface for probing how neural populations encode behavior and sensory information\. Code is available at[https://github\.com/TangLab\-UBC/Neural\-DISCO\-neurips](https://github.com/TangLab-UBC/Neural-DISCO-neurips)\.

## 1Introduction

Modern neural recording technologies record hundreds to thousands of neurons while animals move, make decisions, and receive structured sensory input\. These data make it possible to relate neural population activity to simultaneously recorded behavioral and stimulus labels, a central goal in systems neuroscience\[[20](https://arxiv.org/html/2610.33145#bib.bib11),[52](https://arxiv.org/html/2610.33145#bib.bib10),[60](https://arxiv.org/html/2610.33145#bib.bib6),[25](https://arxiv.org/html/2610.33145#bib.bib12),[31](https://arxiv.org/html/2610.33145#bib.bib13),[32](https://arxiv.org/html/2610.33145#bib.bib53),[11](https://arxiv.org/html/2610.33145#bib.bib54),[55](https://arxiv.org/html/2610.33145#bib.bib55),[3](https://arxiv.org/html/2610.33145#bib.bib16)\]\. Many questions, however, are not purely predictive\. We often want to ask how a neural population response would change if a measured variable were altered—for example, if the animal occupied a different position, reached in another direction, or ran at a different speed—while keeping the remaining source\-specific state as intact as possible\.

![Refer to caption](https://arxiv.org/html/2610.33145v1/method.png)Fig\. 1:Neural\-DISCO overview\.\(a\) Neural\-DISCO reconstructs source spike counts from supervised label latents and a bottlenecked residual latent\. \(b\) Counterfactual editing replaces a selected label latent while holding the remaining label blocks and the source residual fixed\. \(c\) The same block\-structured generator supports label\-specific neuron attribution\.We call this problem*source\-conditioned counterfactual editing*, borrowing counterfactual terminology from causal inference while using it operationally rather than claiming causal identification from observational data alone\[[16](https://arxiv.org/html/2610.33145#bib.bib56)\]\. Given a source responsexxand observed label values\(y1,…,yK\)\(y\_\{1\},\\ldots,y\_\{K\}\), the goal is to replace one label valueyay\_\{a\}byy^a\\hat\{y\}\_\{a\}and generate neural activity for the edited condition while preserving the other observed label valuesy−ay\_\{\-a\}, and residual structure not captured by the labels\. This differs from ordinary label\-conditioned generation\[[51](https://arxiv.org/html/2610.33145#bib.bib21),[37](https://arxiv.org/html/2610.33145#bib.bib22)\]\. A model ofp⁡\(x′∣y^a,y−a\)p\(x^\{\\prime\}\\mid\\hat\{y\}\_\{a\},y\_\{\-a\}\)can sample plausible activity consistent with the target labels, but unless it is conditioned on the source response, it generally marginalizes over internal state, spontaneous movement, and unmeasured task variables\[[53](https://arxiv.org/html/2610.33145#bib.bib19),[38](https://arxiv.org/html/2610.33145#bib.bib20),[60](https://arxiv.org/html/2610.33145#bib.bib6)\]\. This source\-conditioned counterfactual setting instead requires representations that retain residual content of the source sample\. We note that only measured labels are editable; unmeasured variation may be preserved in the source residual but is neither identified nor directly editable\.

We introduceNeural\-DISCO, aDISentangled learning model for source\-conditionedCOunterfactual editing of neural population activity \(Fig\.[1](https://arxiv.org/html/2610.33145#S1.F1)\)\. Each measured behavioral or stimulus label is encoded by its own embedding network into corresponding label\-specific latents, while a separate residual encoder maps the source spike counts to a residual latent\. A generator then decodes the concatenated latents into firing rates\. To reduce observed\-label leakage into the residual latent, we corrupt the residual with Gaussian noise during training and penalize its activation norm\. This bottleneck biases the model to route observed factors through their designated label latents while preserving residual source\-specific variability in the residual latent\.

At test time, Neural\-DISCO encodes a source response, replaces one label\-specific latent with the embedding of a target value, keeps the residual and other label\-specific latents fixed, and decodes the edited representation into firing rates\. The same factorization also supports neuron\-level interpretation\. Because labels enter the generator through label\-specific blocks, we compute integrated gradients to obtain label\-specific attribution scores for each neuron\[[54](https://arxiv.org/html/2610.33145#bib.bib31)\]\.

We evaluate Neural\-DISCO in three settings that probe complementary failure modes: a rat neural activity simulator with known functional cell classes and ground\-truth edited rates; macaque reaching electrophysiology, where edited activity is evaluated by held\-out behavioral decoding; and a mouse visual encoding Neuropixels dataset, where running\-speed and grating stimulus orientation edits are evaluated jointly with stimulus preservation and independent unit\-level tuning metrics\. Across these settings, Neural\-DISCO improves the balance between edit fidelity and residual preservation relative to classical Generalized Linear Model \(GLM\) baselines\. Further, on traditional conditional\-generation tasks Neural\-DISCO performs competitively with strong neural encoding baselines\.

Our contribution is a source\-conditioned framework for neural activity editing and label\-specific attribution\. Neural\-DISCO combines label\-specific embeddings with a noisy residual bottleneck characterized by a capacity bound, and experiments on controlled simulation and two real neural recording datasets support selective editing, source preservation, and label\-specific neuron attribution\.

## 2Related work

Neural latent\-variable models and behavior\-conditioned generation\.Latent\-variable methods such as GPFA, LFADS, and demixed PCA summarize population activity in low\-dimensional representations\[[64](https://arxiv.org/html/2610.33145#bib.bib15),[42](https://arxiv.org/html/2610.33145#bib.bib57),[24](https://arxiv.org/html/2610.33145#bib.bib18)\]\. Recent generative models, including LDNS, NEDS, pi\-VAE, and GNOCCHI, model neural activity conditionally on behavioral labels or learned neural codes\[[21](https://arxiv.org/html/2610.33145#bib.bib28),[65](https://arxiv.org/html/2610.33145#bib.bib38),[66](https://arxiv.org/html/2610.33145#bib.bib35),[33](https://arxiv.org/html/2610.33145#bib.bib29)\]\. Neural\-DISCO instead targets source\-conditioned editing: the edited sample should encode a target label while retaining residual content from a particular source response\.

Residual neural structure and behaviorally targeted dynamics\.Single\-trial neural variability reflects internal state, spontaneous movement, unmeasured task variables, and recording noise\[[53](https://arxiv.org/html/2610.33145#bib.bib19),[38](https://arxiv.org/html/2610.33145#bib.bib20),[60](https://arxiv.org/html/2610.33145#bib.bib6),[52](https://arxiv.org/html/2610.33145#bib.bib10),[3](https://arxiv.org/html/2610.33145#bib.bib16)\]\. Related methods identify task\-specific subspaces\[[4](https://arxiv.org/html/2610.33145#bib.bib61)\], separate behaviorally relevant from other neural dynamics\[[45](https://arxiv.org/html/2610.33145#bib.bib33),[19](https://arxiv.org/html/2610.33145#bib.bib34),[61](https://arxiv.org/html/2610.33145#bib.bib40),[46](https://arxiv.org/html/2610.33145#bib.bib41)\], or distinguish shared and private neural–behavioral representations\[[12](https://arxiv.org/html/2610.33145#bib.bib42),[63](https://arxiv.org/html/2610.33145#bib.bib62)\]\. Neural\-DISCO is complementary: it provides an editable factorization that makes label encodings directly addressable while preserving residual source\-specific structure\.

Disentangled representation learning and counterfactual editing\.Disentangled representation learning traditionally includes approaches based on constrained VAEs, mutual\-information objectives, and factorial latent priors\[[15](https://arxiv.org/html/2610.33145#bib.bib47),[6](https://arxiv.org/html/2610.33145#bib.bib48),[23](https://arxiv.org/html/2610.33145#bib.bib49)\]\. However, unsupervised disentanglement is not identifiable without inductive biases or supervision\[[29](https://arxiv.org/html/2610.33145#bib.bib23)\], motivating our supervised factorization\. Counterfactual or attribute editing in generative models often aims to vary selected attributes or labels while preserving source content\[[27](https://arxiv.org/html/2610.33145#bib.bib50),[17](https://arxiv.org/html/2610.33145#bib.bib51),[28](https://arxiv.org/html/2610.33145#bib.bib52)\]\. Our bottleneck design is closest in spirit to LORD and biolord, which use asymmetric noise regularization to separate observed labels from residual latents\[[7](https://arxiv.org/html/2610.33145#bib.bib30),[44](https://arxiv.org/html/2610.33145#bib.bib32)\]\. Neural\-DISCO adapts this principle to neural spike counts with a Poisson observation model, multiple behavioral or stimulus labels, source\-conditioned editing, and neuron\-level attribution\.

Neuron\-level interpretation\.Classical neuroscience interprets neurons through tuning and modulation analyses relating firing rates to stimuli or behavioral state\[[18](https://arxiv.org/html/2610.33145#bib.bib24),[9](https://arxiv.org/html/2610.33145#bib.bib25),[10](https://arxiv.org/html/2610.33145#bib.bib26),[39](https://arxiv.org/html/2610.33145#bib.bib27)\]\. Recently, xCEBRA provides a complementary approach to identifiable attribution maps for neural time series\[[48](https://arxiv.org/html/2610.33145#bib.bib37)\]\. BrainDiVE synthesizes images predicted to activate cortical regions\[[30](https://arxiv.org/html/2610.33145#bib.bib3),[14](https://arxiv.org/html/2610.33145#bib.bib4)\], while MIG\-Vis visualizes the stimulus semantics of neural latent groups through diffusion\[[62](https://arxiv.org/html/2610.33145#bib.bib5)\]\. These methods produce stimulus images; Neural\-DISCO predicts edited neural responses while retaining source residual structure\.

## 3Methods

### 3\.1Neural\-DISCO model

Data and notation\.Let𝒟=\(X,Y\)\\mathcal\{D\}=\(X,Y\)be paired neural and behavioral or stimulus observations\. Each source responsex∈X⊆ℝnx\\in X\\subseteq\\mathbb\{R\}^\{n\}is a vector of binned spike counts fromnnneurons, and the paired label vectory=\(y1,…,yK\)y=\(y\_\{1\},\\ldots,y\_\{K\}\)contains theKKmeasured labels recorded in the same bin\. A label may be scalar, vector\-valued, or categorical\. The editing task is: given a source responsexxand labelsyy, replace one or more selected labels with user\-specified target values while preserving unedited labels and source\-specific variability not represented by the observed labels\. For an edit to labelaa,y^a\\hat\{y\}\_\{a\}denotes its target value andy−a=\(yb\)b≠ay\_\{\-a\}=\(y\_\{b\}\)\_\{b\\neq a\}denotes all other labels\.

Factor\-specific and residual representations\.For each observed labela∈\{1,…,K\}a\\in\\\{1,\\ldots,K\\\}, Neural\-DISCO uses a separate embedding network

za=ϕa​\(ya\),za∈ℝda,z\_\{a\}=\\phi\_\{a\}\(y\_\{a\}\),\\qquad z\_\{a\}\\in\\mathbb\{R\}^\{d\_\{a\}\},\(1\)whereϕa:𝒴a→ℝda\\phi\_\{a\}:\\mathcal\{Y\}\_\{a\}\\to\\mathbb\{R\}^\{d\_\{a\}\}is learned\. Becausezaz\_\{a\}is computed only fromyay\_\{a\}, it can carry only information available through that observed label, including variables statistically correlated withyay\_\{a\}\[[29](https://arxiv.org/html/2610.33145#bib.bib23)\]\. Idiosyncratic source content not determined by observed labels must pass through the residual path\.

A residual encoder maps the source response to the source residual latentu=ϕres​\(x\),u∈ℝmu=\\phi\_\{\\mathrm\{res\}\}\(x\),u\\in\\mathbb\{R\}^\{m\}\. The clean representation is

z=\(z1,…,zK,u\),dtot=∑a=1Kda\+m\.z=\(z\_\{1\},\\ldots,z\_\{K\},u\),\\qquad d\_\{\\mathrm\{tot\}\}=\\sum\_\{a=1\}^\{K\}d\_\{a\}\+m\.\(2\)During training only, the residual latent is corrupted by Gaussian noise,

u~=u\+ϵ,ϵ∼𝒩⁡\(0,σ2​Im\),z~=\(z1,…,zK,u~\)\.\\tilde\{u\}=u\+\\epsilon,\\qquad\\epsilon\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}I\_\{m\}\),\\qquad\\tilde\{z\}=\(z\_\{1\},\\ldots,z\_\{K\},\\tilde\{u\}\)\.At test time and for counterfactual editing, we use the clean residualu=ϕres​\(x\)u=\\phi\_\{\\mathrm\{res\}\}\(x\); the noisy residualu~\\tilde\{u\}is used only to regularize the residual channel during training\.

Observation model\.The label\-specific encodersϕa\\phi\_\{a\}, residual encoderϕres\\phi\_\{\\mathrm\{res\}\}, and generatorGθG\_\{\\theta\}are implemented as feed\-forward multilayer perceptrons \(MLPs\), with dataset\-specific architectures summarized in Appendix Table[12](https://arxiv.org/html/2610.33145#A1.T12)\. The generatorGθ:ℝdtot→ℝ\>0nG\_\{\\theta\}:\\mathbb\{R\}^\{d\_\{\\mathrm\{tot\}\}\}\\to\\mathbb\{R\}\_\{\>0\}^\{n\}maps the concatenated representation to rates,

rθ​\(z~\)=Gθ​\(z~\),r\_\{\\theta\}\(\\tilde\{z\}\)=G\_\{\\theta\}\(\\tilde\{z\}\),\(3\)with positivity enforced by a softplus output\. Conditional onz~\\tilde\{z\}, neurons are modeled with an independent Poisson likelihood\[[42](https://arxiv.org/html/2610.33145#bib.bib57),[58](https://arxiv.org/html/2610.33145#bib.bib58)\],

pθ​\(x∣z~\)=∏ℓ=1nPoisson⁡\(xℓ∣rθ,ℓ​\(z~\)\)\.p\_\{\\theta\}\(x\\mid\\tilde\{z\}\)=\\prod\_\{\\ell=1\}^\{n\}\\mathrm\{Poisson\}\\big\(x\_\{\\ell\}\\mid r\_\{\\theta,\\ell\}\(\\tilde\{z\}\)\\big\)\.\(4\)
Training objective\.The model minimizes the Poisson negative log\-likelihood plus a residual activation penalty, dropping the additive constants∑ℓ=1nlog⁡\(xℓ\(s\)\!\)\\sum\_\{\\ell=1\}^\{n\}\\log\(x\_\{\\ell\}^\{\(s\)\}\!\)that do not depend on model parameters\. For sampless, define

u\(s\)=ϕres​\(x\(s\)\),u~\(s\)=u\(s\)\+ϵ\(s\),z~\(s\)=\(ϕ1​\(y1\(s\)\),…,ϕK​\(yK\(s\)\),u~\(s\)\),u^\{\(s\)\}=\\phi\_\{\\mathrm\{res\}\}\(x^\{\(s\)\}\),\\qquad\\tilde\{u\}^\{\(s\)\}=u^\{\(s\)\}\+\\epsilon^\{\(s\)\},\\qquad\\tilde\{z\}^\{\(s\)\}=\\big\(\\phi\_\{1\}\(y\_\{1\}^\{\(s\)\}\),\\ldots,\\phi\_\{K\}\(y\_\{K\}^\{\(s\)\}\),\\tilde\{u\}^\{\(s\)\}\\big\),andrθ\(s\)=Gθ​\(z~\(s\)\)r\_\{\\theta\}^\{\(s\)\}=G\_\{\\theta\}\(\\tilde\{z\}^\{\(s\)\}\)\. The model minimizes

ℒ⁡\(θ,\{ϕa\},ϕres\)=\\displaystyle\\mathcal\{L\}\(\\theta,\\\{\\phi\_\{a\}\\\},\\phi\_\{\\mathrm\{res\}\}\)=1N​∑s=1N𝔼ϵ\(s\)​\[∑ℓ=1n\(rθ,ℓ\(s\)−xℓ\(s\)​log⁡rθ,ℓ\(s\)\)\]\+λ​1N​∑s=1N‖u\(s\)‖22\.\\displaystyle\\frac\{1\}\{N\}\\sum\_\{s=1\}^\{N\}\\mathbb\{E\}\_\{\\epsilon^\{\(s\)\}\}\\left\[\\sum\_\{\\ell=1\}^\{n\}\\left\(r\_\{\\theta,\\ell\}^\{\(s\)\}\-x\_\{\\ell\}^\{\(s\)\}\\log r\_\{\\theta,\\ell\}^\{\(s\)\}\\right\)\\right\]\+\\lambda\\frac\{1\}\{N\}\\sum\_\{s=1\}^\{N\}\\\|u^\{\(s\)\}\\\|\_\{2\}^\{2\}\.\(5\)The expectation over noise is estimated with one Monte Carlo sample ofϵ\(s\)\\epsilon^\{\(s\)\}per training example in each minibatch\. Theλ\\lambda\-selection criterion favors high validation edit fidelity and low observed\-label decodability from the clean residual, averaged across ten validation folds\. Dataset\-specific settings are reported in Appendix[A](https://arxiv.org/html/2610.33145#A1)\.

##### Capacity interpretation of the residual bottleneck\.

The residual should preserve source information needed for reconstruction while avoiding redundant storage of the observed labels\. We adapt the asymmetric\-noise idea in\[[7](https://arxiv.org/html/2610.33145#bib.bib30),[57](https://arxiv.org/html/2610.33145#bib.bib59),[1](https://arxiv.org/html/2610.33145#bib.bib60)\]: observed labels enter deterministically through label\-specific latents, whereas the residual path is made noisy and norm\-penalized\.

###### Proposition 1\.

Capacity of the noisy training\-time residual channel\.Fix a trained residual encoderϕres\\phi\_\{\\mathrm\{res\}\}and assumeσ\>0\\sigma\>0\. LetS∼Unif​\{1,…,N\}S\\sim\\mathrm\{Unif\}\\\{1,\\ldots,N\\\}be the training\-sample index, and define

XS=x\(S\),Ya,S=ya\(S\),US=ϕres​\(XS\)∈ℝm\.X\_\{S\}=x^\{\(S\)\},\\qquad Y\_\{a,S\}=y\_\{a\}^\{\(S\)\},\\qquad U\_\{S\}=\\phi\_\{\\mathrm\{res\}\}\(X\_\{S\}\)\\in\\mathbb\{R\}^\{m\}\.Assumem≥1m\\geq 1and letϵ∼𝒩⁡\(0,σ2​Im\)\\epsilon\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}I\_\{m\}\)be independent ofSS\. Define the noisy training\-time residual and empirical residual power by

U~=US\+ϵ,P^=1N​∑s=1N‖ϕres​\(x\(s\)\)‖22\.\\tilde\{U\}=U\_\{S\}\+\\epsilon,\\qquad\\hat\{P\}=\\frac\{1\}\{N\}\\sum\_\{s=1\}^\{N\}\\\|\\phi\_\{\\mathrm\{res\}\}\(x^\{\(s\)\}\)\\\|\_\{2\}^\{2\}\.Then, using natural logarithms so that mutual information is measured in nats,

I⁡\(S,U~\)≤m2​log⁡\(1\+P^m​σ2\)\.I\(S;\\tilde\{U\}\)\\leq\\frac\{m\}\{2\}\\log\\left\(1\+\\frac\{\\hat\{P\}\}\{m\\sigma^\{2\}\}\\right\)\.\(6\)Moreover, for any observed factorYa,SY\_\{a,S\}and deterministic embeddingZa,S=ϕa​\(Ya,S\)Z\_\{a,S\}=\\phi\_\{a\}\(Y\_\{a,S\}\),

I⁡\(Ya,S,U~\)≤I⁡\(S,U~\),I⁡\(Za,S,U~\)≤I⁡\(Ya,S,U~\)\.I\(Y\_\{a,S\};\\tilde\{U\}\)\\leq I\(S;\\tilde\{U\}\),\\qquad I\(Z\_\{a,S\};\\tilde\{U\}\)\\leq I\(Y\_\{a,S\};\\tilde\{U\}\)\.\(7\)

The proof is given in Appendix[B](https://arxiv.org/html/2610.33145#A2)\. Proposition[1](https://arxiv.org/html/2610.33145#Thmproposition1)applies only to the randomized residual channelS↦U~S\\mapsto\\tilde\{U\}used during training, and does not directly bound the clean residualS↦USS\\mapsto U\_\{S\}which is used at test time\. The result does not guarantee disentanglement, but formalizes why training with the information bottleneck reduces the information capacity available for residual leakage\. Decoding probes show that each label is less decodable from the residual when supplied to the model rather than when omitted as intended \(Appendix Table[6](https://arxiv.org/html/2610.33145#A1.T6)\)\. Identifiability results rely on structured assumptions, such as auxiliary\-variable\-conditioned factorial priors\[[22](https://arxiv.org/html/2610.33145#bib.bib1)\]or paired observations before and after interventions\[[5](https://arxiv.org/html/2610.33145#bib.bib2)\]\. The asymmetric\-noise bottleneck\[[7](https://arxiv.org/html/2610.33145#bib.bib30)\]controls a noisy training channel without providing such identification guarantees\.

### 3\.2Source\-conditioned counterfactual editing

Given a source pair\(x,y\)\(x,y\), Neural\-DISCO computes

za=ϕa\(ya\),a=1,…,K,u=ϕres\(x\)\.z\_\{a\}=\\phi\_\{a\}\(y\_\{a\}\),\\quad a=1,\\ldots,K,\\qquad u=\\phi\_\{\\mathrm\{res\}\}\(x\)\.\(8\)To edit labelaa, only that label\-specific latent is replaced,

z^a=ϕa​\(y^a\),z^=\(z1,…,za−1,z^a,za\+1,…,zK,u\)\.\\hat\{z\}\_\{a\}=\\phi\_\{a\}\(\\hat\{y\}\_\{a\}\),\\quad\\hat\{z\}=\(z\_\{1\},\\ldots,z\_\{a\-1\},\\hat\{z\}\_\{a\},z\_\{a\+1\},\\ldots,z\_\{K\},u\)\.\(9\)The counterfactual rate vector is

r^=Gθ​\(z^\)\.\\hat\{r\}=G\_\{\\theta\}\(\\hat\{z\}\)\.\(10\)When sampled spikes are required,

x^ℓ∼Poisson\(r^ℓ\),ℓ=1,…,n\.\\hat\{x\}\_\{\\ell\}\\sim\\mathrm\{Poisson\}\(\\hat\{r\}\_\{\\ell\}\),\\qquad\\ell=1,\\ldots,n\.\(11\)
##### Counterfactual semantics\.

We use “counterfactual” operationally\. Outputs should be interpreted as model\-based edits, not as identified causal effects from observational recordings\. Poisson sampling after label latent replacement yields the corresponding interventional distribution under three assumptions: the true Poisson spike generator factorizes as in Appendix[A\.5](https://arxiv.org/html/2610.33145#A1.SS5); interventions leave the unedited labels and source state fixed; and the model recovers the true label maps, generator, and source\-state encoding\.

### 3\.3Neuron\-level label attribution

We quantify label\-specific contributions to each neuron’s pre\-softplus generator output using integrated gradients\[[54](https://arxiv.org/html/2610.33145#bib.bib31)\], along a straight path from zeroed observed labels and a zero residual to the source input\. We first average signed attributions across 32 randomly selected source samples, then take the mean absolute value over each label’s input dimensions or the residual dimensions\. Appendix[A\.4](https://arxiv.org/html/2610.33145#A1.SS4)gives the precise definition\.

### 3\.4Baselines and ablations

We compare with a no regularization baseline \(w/o reg\.\), a zero residual baseline, pi\-VAE\[[66](https://arxiv.org/html/2610.33145#bib.bib35)\], and NEDS\[[65](https://arxiv.org/html/2610.33145#bib.bib38)\], and Gaussian and Poisson GLMs\. The GLMs linearly project the observed labels to predict neural activity under their corresponding distributions\. The Gaussian GLM is included as a linear editing baseline; we use it to edit by replacing the projection of the source label with the projection of the target label\. The zero residual control removes Neural\-DISCO’s source residual during editing\. The no regularization baseline is an architecture\-matched ablation without any information bottleneck in the residual \(λ=σ=0\\lambda=\\sigma=0\)\. pi\-VAE, NEDS and the Poisson GLM do not expose the same source\-conditioned residual\-editing interface\. For counterfactual evaluation, we therefore use their target\-conditioned generation interface with the edited labels\. These comparisons test whether source\-conditioned residual preservation provides benefits beyond strong target\-conditioned baselines; they should not be read as claiming that these baselines were designed for the same editing interface\.

For the Allen Visual Coding evaluation, speed editors share the training\-fold mean\-speed target; methods within each task share decoders trained on real training activity\. Checkpoints are selected using validation folds\. Scores from the held\-out test epoch are averaged across folds, then across sessions for the 32\-session analysis\.

## 4Experiments

We evaluate whether Neural\-DISCO supports source\-conditioned counterfactual editing, as well as standard label\-conditioned generation\. The experiments address three questions\. First, when the data\-generating factors and counterfactual firing rates are known, does the model change the requested label while preserving non\-target and unobserved structure? Second, in real motor cortical recordings without ground\-truth counterfactual spikes, does generated activity decode the target behavioral label rather than remain source\-like? Third, in a larger visual cortex dataset, can the same factorization average out a specified behavioral label or edit the encoding of a specific stimulus label, while preserving held\-out stimulus and behavioral information, and producing unit\-level attributions aligned with independent biological metrics?

For our non\-simulated datasets, there is no ground\-truth neural response for performed edits\. We therefore use an indirect decoding\-based evaluation strategy\. We use decoders trained only on real neural activity data to evaluate whether generated activity encodes both the target label, as well as labels unobserved by the model which should be conserved through the residual latent\.

Unless otherwise stated, values are means across ten folds;±\\pmvalues and error bars denote the standard error of the mean \(SEM\)\. For our Allen Visual Coding dataset editing, we report the mean and SEM across 32 sessions\. Appendix[A](https://arxiv.org/html/2610.33145#A1)gives preprocessing and training settings; Appendices[A\.12](https://arxiv.org/html/2610.33145#A1.SS12)and[A\.13](https://arxiv.org/html/2610.33145#A1.SS13)report mixed\-selectivity and sampling controls\.

![Refer to caption](https://arxiv.org/html/2610.33145v1/fig2.png)Fig\. 2:Controlled RatInABox evaluation\.\(a\) Simulated movement with place, head\-direction and speed cells\. \(b\) Cell\-type recovery acrossλ\\lambda; speed\-cell recovery uses residual attribution because speed is withheld\. \(c\) Residual PCA colored by speed\. \(d,e\) Position shift while retaining head direction and speed\. \(f,g\) Head\-direction reversal while retaining position and speed\. Generated neural activity is scored against simulator targets\.### 4\.1RatInABox simulation

We first evaluate Neural\-DISCO on a controlled RatInABox simulation dataset, where an agent moves in a2×12\\times 1meter enclosed environment and both functional cell classes and ground\-truth counterfactual firing rates are known\[[8](https://arxiv.org/html/2610.33145#bib.bib39)\]\. The population contains 100 place cells\[[40](https://arxiv.org/html/2610.33145#bib.bib9)\], 100 head\-direction cells\[[56](https://arxiv.org/html/2610.33145#bib.bib8)\], and 100 speed cells\[[26](https://arxiv.org/html/2610.33145#bib.bib7)\], whose rates depend on position, head\-direction and speed respectively\. Neural\-DISCO receives position and head direction as observed labels, while speed is withheld and can only be represented through the residual latent\. This creates a direct test of both label\-specific editing and residual preservation\.

Attribution recovers simulated cell classes\.For each label, we use the corresponding neuron\-level attribution score to classify the matching cell type and report area under the ROC curve \(AUC\)\. Position attribution should identify place cells, head\-direction attribution should identify head\-direction cells, and residual attribution should identify speed cells as speed is unobserved\. Cell\-type recovery is higher at the strongest bottleneck than atλ=0\\lambda=0\(Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)b\)\. The residual representation is also organized by held\-out speed \(Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)c\), consistent with the residual pathway preserving source\-specific structure not supplied by the observed labels\.

Counterfactual edited rates match ground\-truth targets\.The simulator allows direct comparison between generated rates and the true rates under an edited behavioral label state\. We evaluate two edits\. In the position\-shift edit, the source trajectory is confined to the left half of the environment and the position label is shifted to the right half, while head direction and speed are unchanged\. The desired output should match target place\-cell rates, while preserving source speed\-cell and head\-direction cell rates\. Neural\-DISCO follows the shifted place\-cell target while preserving speed\-cell activity, whose corresponding behavior is unobserved, through the residual latent \(Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)d,e\)\. In the head\-direction edit, head direction is reversed while position and speed are unchanged\. The model should change head\-direction\-cell activity, preserve place\-cell activity, and preserve speed\-cell activity through the residual\. Neural\-DISCO again tracks the edited target rates for the manipulated label while retaining non\-target and held\-out components \(Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)f,g\)\. The no regularization model without the bottleneck shows the expected failure mode: the residual pathway can carry label information, so replacing an observed\-label embedding does not reliably make the decoded activity target\-like\. As an additional diagnostic, we train variants in which each behavioral variable is withheld in turn and must be represented through the residual\. These diagnostic tests yield the same conclusion\. Appendix[A\.1](https://arxiv.org/html/2610.33145#A1.SS1)reports these withheld\-factor tests; Appendix[A\.16\.1](https://arxiv.org/html/2610.33145#A1.SS16.SSS1)tests generalization when position–heading combinations are withheld during training\.

![Refer to caption](https://arxiv.org/html/2610.33145v1/fig3.png)Fig\. 3:Counterfactual editing in macaque reaching electrophysiology\.\(a\) Editing replaces hand\-position, hand\-velocity and eye\-position blocks while retaining the source residual\. \(b\) Trajectories decoded from source spikes, target spikes and Neural\-DISCO counterfactual activity\. \(c,d\) Synthetic constant\-rate controls test whether attribution distinguishes recorded neurons from controls\. \(e\) Target\-decodingR2R^\{2\}comparisons; dotted lines show the empirical target\-spike decoding reference\.
### 4\.2Macaque reaching electrophysiology

We next evaluate Neural\-DISCO on real motor and premotor cortical recordings from a center\-out reaching task in the Neural Latents Benchmark\[[43](https://arxiv.org/html/2610.33145#bib.bib17)\]\. The dataset contains 137 recorded units\. Unlike the simulator, the recordings contain unobserved variability, recording noise, and correlated behavioral structure, and there is no ground\-truth neural response for a hypothetical edited reach\. We therefore evaluate edited activity through behavioral decoding rather than direct rate error\. We also benchmark label\-conditioned generators and use inconsistent behavioral targets to check whether decoding simply tracks the requested labels \(Appendix[A\.9](https://arxiv.org/html/2610.33145#A1.SS9)\)\.

Edited activity decodes as the target reach\.For each held\-out source trial, we select a target trial with matched task structure and a different reach trajectory\. Here we jointly replace the hand\-position, hand\-velocity, and eye\-position latents with target\-trial label latents to test whether the generator can produce a target\-like behavioral episode while retaining the source residual\. A decoder based on CEBRA\[[49](https://arxiv.org/html/2610.33145#bib.bib36)\]trained only on real neural spikes is applied to source spikes, target spikes, and generated activity\. We compute target\-decodingR2R^\{2\}against the target behavior\. Source\-spike decoding is a negative control, and target\-spike decoding is an empirical reference for the performance achieved by this decoder on real data\.

As a result, Neural\-DISCO counterfactuals decode close to target trajectories, withR2=0\.895±0\.002R^\{2\}=0\.895\\pm 0\.002for hand position,0\.772±0\.0030\.772\\pm 0\.003for hand velocity, and0\.630±0\.0040\.630\\pm 0\.004for eye position\. Target\-spike decoding reaches0\.9090\.909,0\.8450\.845, and0\.6560\.656, respectively, whereas source spikes decode poorly against the target trajectories\. The unregularized ablation remains source\-like and yields negative target\-decodingR2R^\{2\}, consistent with label information leaking through the residual pathway \(Fig\.[12](https://arxiv.org/html/2610.33145#A3.F12)\)\. Neural\-DISCO substantially improves over the unregularized control and performs competitively with pi\-VAE and NEDS across decoded variables \(Fig\.[3](https://arxiv.org/html/2610.33145#S4.F3)e\)\. The decoded examples in Fig\.[3](https://arxiv.org/html/2610.33145#S4.F3)b show the same qualitative behavior: generated trajectories follow the target reach rather than the source reach, while the unregularized ablation remains source\-like \(Fig\.[13](https://arxiv.org/html/2610.33145#A3.F13)\)\.

Synthetic control neurons provide an attribution sanity check\.This dataset does not provide ground\-truth behavioral importance labels for individual units\. We therefore add synthetic constant\-rate control neurons as a negative\-control population\. These controls are Poisson spike trains with constant behavior\-independent firing rates matched to the dataset mean firing rate\. They should receive low attribution to behavior, whereas recorded motor and premotor units are more likely to contain movement\-related structure\[[9](https://arxiv.org/html/2610.33145#bib.bib25),[10](https://arxiv.org/html/2610.33145#bib.bib26)\]\. We use the total attribution to all three behaviors in our analyses\. Real\-neuron recovery AUC increases withλ\\lambdaover the positive range \(Fig\.[3](https://arxiv.org/html/2610.33145#S4.F3)d\)\. This pattern is consistent with the counterfactual results: bottlenecking the residual discourages leakage of observed behavioral information and makes label\-specific attribution more accurate\. This is a sanity check rather than biological ground truth, because not every recorded unit is expected to be modulated by the observed labels\.

![Refer to caption](https://arxiv.org/html/2610.33145v1/fig4.png)Fig\. 4:Allen Visual Coding evaluation\.\(a,g\) Speed\- and orientation\-edit schematics\. \(b,c\) Speed\-edit example and speed\-target RMSE\. \(d,e\) Temporal\-frequency preservation example and balanced accuracy after speed editing\. \(h,i\) Orientation\-edit example and target\-orientation balanced accuracy\. \(j,k\) Running\-speed preservation example and speed RMSE after orientation editing\. Editing summaries are mean±\\pmSEM across 32 session means, each averaging ten folds\. \(f,l\) Running\-modulation AUC and orientation\-selectivity correlation acrossλ\\lambda\. In summary bars, Source denotes decoding of recorded activity\.
### 4\.3Allen Visual Coding

Finally, we evaluate Neural\-DISCO on 32 recording sessions from the Allen Visual Coding Neuropixels dataset\[[50](https://arxiv.org/html/2610.33145#bib.bib43),[2](https://arxiv.org/html/2610.33145#bib.bib46)\]\. These recordings sample multiple visual and related brain areas while mice view visual stimuli and run on a wheel, a behavioral variable known to modulate visual cortical responses\[[39](https://arxiv.org/html/2610.33145#bib.bib27)\]\. We focus on drifting\-grating stimulus epochs\.

Speed editing and stimulus preservation\.We use running speed and grating orientation as observed labels\. The temporal frequency \(TF\) of the gratings is intentionally withheld from the observed\-label embeddings, so temporal\-frequency information must be preserved through the residual pathway\. We edit running speed to its mean in each training fold, using decoders trained only on real training activity\. Across 32 sessions, Neural\-DISCO reduces speed\-target RMSE to3\.629±\.5673\.629\\pm\.567cm/s, compared with6\.017±1\.0446\.017\\pm 1\.044for Gaussian GLM \(Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)c\)\. Models and decoders are fit using ten\-fold cross\-validation within each session; these values summarize session means, with SEM across sessions\.

We then test whether information not supplied as an observed label is preserved after the speed edit\. Temporal frequency is withheld during Neural\-DISCO training, and a logistic\-regression decoder trained on real activity predicts the source temporal frequency from edited activity\. Across the same 32 sessions, temporal\-frequency balanced accuracy is\.541±\.013\.541\\pm\.013for Neural\-DISCO and\.597±\.012\.597\\pm\.012for Gaussian GLM, both above chance \(Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)e\)\. Neural\-DISCO achieves lower speed\-target RMSE, whereas Gaussian GLM preserves more temporal\-frequency information \(Appendix Table[2](https://arxiv.org/html/2610.33145#A1.T2)\)\.

Orientation editing and stimulus preservation\.Across the same 32 sessions, we train a model with orientation and TF observed while running speed is withheld\. We edit by rotating orientation by180∘180^\{\\circ\}\. Target\-orientation and source\-TF balanced accuracies are\.949±\.006\.949\\pm\.006and\.823±\.012\.823\\pm\.012, versus\.714±\.016\.714\\pm\.016and\.532±\.015\.532\\pm\.015for Gaussian GLM \(Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)i; Appendix[A\.10](https://arxiv.org/html/2610.33145#A1.SS10); mean±\\pmSEM across session means\)\. The unregularized ablation has lower orientation fidelity and TF preservation than Neural\-DISCO\. Figures[4](https://arxiv.org/html/2610.33145#S4.F4)j,k show source\-speed preservation; Appendix[A\.10](https://arxiv.org/html/2610.33145#A1.SS10)reports all seven methods and additional stimulus edits\.

Attributions align with dataset\-provided unit metrics\.The Allen dataset provides unit metrics that can be used as independent proxies for attribution quality\. We classify units with significant running modulation using the running\-modulationpp\-value, thresholding atp<0\.05p<0\.05, and compute the AUC obtained from running\-attribution scores\. For stimulus tuning, we compare orientation\-attribution scores with the global orientation selectivity index \(gOSI\)\[[2](https://arxiv.org/html/2610.33145#bib.bib46)\]\. As the residual bottleneck is strengthened, running\-attribution AUC and orientation\-attribution correlation both increase \(Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)f,l\)\.

## 5Discussion

Neural\-DISCO provides a source\-conditioned interface for counterfactual editing of neural population activity\. The key design choice is a supervised factorization with a bottlenecked residual: observed labels remain directly addressable, while source\-specific variability that is not supplied through those labels can still be preserved\. This design targets two common failure modes\. A purely label\-conditioned generator cannot preserve residual structure from source spiking, whereas an unconstrained residual pathway can leak the edited label and prevent the edit from taking effect\[[7](https://arxiv.org/html/2610.33145#bib.bib30)\]\.

The experiments support this edit–preserve trade\-off\. In RatInABox, where edited rates and functional cell classes are known, Neural\-DISCO edits position or head direction while preserving withheld speed, and its attributions recover simulated cell classes\. In macaque reaching recordings, generated activity decodes as the target reach, whereas the unregularized ablation remains source\-like\. In the multi\-session Allen Visual Coding dataset, Neural\-DISCO is able to edit orientation and speed while preserving unobserved labels\. Further, its label attributions align with dataset\-provided unit metrics\.

Neural\-DISCO for Neuroscience\.Neuronal responses to stimuli are not determined solely by the presented stimulus, but are modulated by latent aspects of brain state such as arousal, locomotion, and attention\[[39](https://arxiv.org/html/2610.33145#bib.bib27),[41](https://arxiv.org/html/2610.33145#bib.bib45),[34](https://arxiv.org/html/2610.33145#bib.bib44)\]\. Neural\-DISCO addresses this by providing a flexible editing framework in which the relationship between observed labels and neural activity is conditioned on the preserved source\-specific residual latent\. Two source responses with identical labels can therefore yield different counterfactual edits if they differ in their residual neural state\. We believe this capability can make Neural\-DISCO a useful tool for neuroscientists seeking to understand how both observed and unobserved aspects of brain state modulate neural representations\. We explore this on the Allen Visual Coding dataset by reversing the editing operation: we keep the label inputs fixed, replace a source response’s residual with that of another response with the same orientation and similar running speed, and measure how each neuron’s predicted orientation tuning changes \(Appendix[A\.14](https://arxiv.org/html/2610.33145#A1.SS14)\)\. Across neurons, we find that the variability in Neural\-DISCO’s predicted firing rates over different source residuals recovers the extent to which temporal frequency, which is withheld from the model, modulates each neuron’s firing\. This proof of concept demonstrates that Neural\-DISCO counterfactuals can expose testable hypotheses about how unobserved factors of brain state modulate neural responses\.

Limitations\.Neural\-DISCO counterfactuals are model\-based predictions under a learned factorization, not causal effects from observational data alone\[[16](https://arxiv.org/html/2610.33145#bib.bib56)\]\. As in prior counterfactual modeling\[[47](https://arxiv.org/html/2610.33145#bib.bib14)\], interpretations are most reliable within the training support and when edited labels can vary independently; edits violating dependencies among behavior, stimuli, and internal state should be treated as extrapolations\. Attribution scores likewise measure model sensitivity rather than proving causal or exclusive encoding by individual neurons\. The residual bottleneck can introduce a trade\-off: a weak residual channel may harm reconstruction and preservation of unobserved labels, whereas an overly strong residual channel can leak observed\-label information and make edits ineffective\. Our capacity bound and sensitivity analyses characterize this trade\-off but do not guarantee disentanglement, especially under correlated labels, missing variables, or dataset shift, consistent with known limitations of disentanglement without sufficient inductive bias or supervision\[[29](https://arxiv.org/html/2610.33145#bib.bib23)\]\. The model is static, without temporal or transition dynamics\. Finally, in real neural data, validation remains indirect because decoder\-based assays cannot replace ground\-truth edited responses from perturbation experiments, matched repeated trials, or more realistic simulators\.

Broader Impact\.Neural\-DISCO is intended forin\-silicohypothesis generation in computational neuroscience, using available neural recordings\. By exposing editable behavioral and stimulus labels while preserving source\-specific residual structure, it provides a model\-based interface for prioritizing hypotheses about neural population coding before pursuing more costly targeted analyses or perturbation experiments\.

## Acknowledgments and Disclosure of Funding

This work was supported by start\-up funding from the University of British Columbia and a Discovery Grant from the Natural Sciences and Engineering Research Council of Canada \[RGPIN\-2025\-04460\] to X\.T\. This research was undertaken, in part, thanks to funding from the Canada Research Chairs Program \[CRC\-2024\-00164\]\. Research infrastructure used in this work was supported by the Canada Foundation for Innovation through the John R\. Evans Leaders Fund \[project 45547\] and by the British Columbia Knowledge Development Fund\. This research was also enabled in part by support provided by the Digital Research Alliance of Canada\.

Competing interests\.The authors declare no competing interests\.

## References

## References

- \[1\]A\. A\. Alemi, I\. Fischer, J\. V\. Dillon, and K\. Murphy\(2016\)Deep Variational Information Bottleneck\.External Links:[Document](https://dx.doi.org/10.48550/ARXIV.1612.00410)Cited by:[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p1.1)\.
- \[2\]Allen Institute for Brain Science\(2019\)Allen SDK: visual coding – neuropixels\.External Links:[Link](https://allensdk.readthedocs.io/en/latest/visual_coding_neuropixels.html)Cited by:[Appendix D](https://arxiv.org/html/2610.33145#A4.p4.1),[§4\.3](https://arxiv.org/html/2610.33145#S4.SS3.p1.1),[§4\.3](https://arxiv.org/html/2610.33145#S4.SS3.p5.1)\.
- \[3\]D\. Angelaki, B\. Benson, J\. Benson, D\. Birman, N\. Bonacchi, K\. Bougrova, S\. A\. Bruijns, M\. Carandini, J\. A\. Catarino,et al\.\(2025\)A brain\-wide map of neural activity during complex behaviour\.Nature645\(8079\),pp\. 177–191\.Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1),[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[4\]M\. C\. Aoi, V\. Mante, and J\. W\. Pillow\(2020\)Prefrontal cortex exhibits multidimensional dynamic encoding during decision\-making\.Nature Neuroscience23\(11\),pp\. 1410–1420\.External Links:[Document](https://dx.doi.org/10.1038/s41593-020-0696-5)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[5\]J\. Brehmer, P\. de Haan, P\. Lippe, and T\. S\. Cohen\(2022\)Weakly supervised causal representation learning\.InAdvances in Neural Information Processing Systems,Vol\.35,pp\. 38319–38331\.External Links:[Document](https://dx.doi.org/10.52202/068431-2776),[Link](https://proceedings.neurips.cc/paper_files/paper/2022/hash/fa567e2b2c870f8f09a87b6e73370869-Abstract-Conference.html)Cited by:[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p2.1)\.
- \[6\]X\. Chen, Y\. Duan, R\. Houthooft, J\. Schulman, I\. Sutskever, and P\. Abbeel\(2016\)InfoGAN: interpretable representation learning by information maximizing generative adversarial nets\.InAdvances in Neural Information Processing Systems 29,pp\. 2172–2180\.External Links:[Link](https://proceedings.neurips.cc/paper/2016/hash/7c9d0b1f96aebd7b5eca8c3edaa19ebb-Abstract.html)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[7\]A\. Gabbay and Y\. Hoshen\(2020\)Demystifying inter\-class disentanglement\.arXiv\.External Links:[Link](http://arxiv.org/abs/1906.11796),[Document](https://dx.doi.org/10.48550/arXiv.1906.11796),1906\.11796 \[cs\]Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1),[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2610.33145#S5.p1.1)\.
- \[8\]T\. M\. George, M\. Rastogi, W\. de Cothi, C\. Clopath, K\. Stachenfeld, and C\. Barry\(2024\)RatInABox, a toolkit for modelling locomotion and neuronal activity in continuous environments\.Elife13,pp\. e85274\.Cited by:[§A\.16](https://arxiv.org/html/2610.33145#A1.SS16.SSS0.Px1.p1.1),[Appendix D](https://arxiv.org/html/2610.33145#A4.p2.1),[§4\.1](https://arxiv.org/html/2610.33145#S4.SS1.p1.1)\.
- \[9\]A\. P\. Georgopoulos, J\. F\. Kalaska, R\. Caminiti, and J\. T\. Massey\(1982\)On the relations between the direction of two\-dimensional arm movements and cell discharge in primate motor cortex\.The Journal of Neuroscience2\(11\),pp\. 1527–1537\.External Links:[Document](https://dx.doi.org/10.1523/JNEUROSCI.02-11-01527.1982)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1),[§4\.2](https://arxiv.org/html/2610.33145#S4.SS2.p4.1)\.
- \[10\]A\. P\. Georgopoulos, A\. B\. Schwartz, and R\. E\. Kettner\(1986\)Neuronal population coding of movement direction\.Science233\(4771\),pp\. 1416–1419\.External Links:[Document](https://dx.doi.org/10.1126/science.3749885)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1),[§4\.2](https://arxiv.org/html/2610.33145#S4.SS2.p4.1)\.
- \[11\]J\. I\. Glaser, A\. S\. Benjamin, R\. H\. Chowdhury, M\. G\. Perich, L\. E\. Miller, and K\. P\. Kording\(2020\)Machine learning for neural decoding\.eNeuro7\(4\),pp\. ENEURO\.0506–19\.2020\.External Links:[Document](https://dx.doi.org/10.1523/ENEURO.0506-19.2020)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[12\]R\. Gondur, U\. B\. Sikandar, E\. S\. Schaffer, M\. Aoi, and S\. Keeley\(2024\)Multi\-modal Gaussian Process Variational Autoencoders for neural and behavioral data\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=aGH43rjoe4)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[13\]K\. Hardcastle, N\. Maheswaranathan, S\. Ganguli, and L\. M\. Giocomo\(2017\)A multiplexed, heterogeneous, and adaptive code for navigation in medial entorhinal cortex\.Neuron94\(2\),pp\. 375–387\.e7\.External Links:[Document](https://dx.doi.org/10.1016/j.neuron.2017.03.025)Cited by:[§A\.12](https://arxiv.org/html/2610.33145#A1.SS12.p1.1)\.
- \[14\]M\. M\. Henderson, A\. F\. Luo, S\. Park, M\. J\. Tarr, and L\. Wehbe\(2026\)Diffusion\-based stimulus optimization reveals functional organization across higher visual cortex\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.64898/2026.05.12.724119),[Link](https://www.biorxiv.org/content/10.64898/2026.05.12.724119v1)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1)\.
- \[15\]I\. Higgins, L\. Matthey, A\. Pal, C\. P\. Burgess, X\. Glorot, M\. M\. Botvinick, S\. Mohamed, and A\. Lerchner\(2017\)Beta\-VAE: learning basic visual concepts with a constrained variational framework\.In5th International Conference on Learning Representations \(ICLR\),External Links:[Link](https://openreview.net/forum?id=Sy2fzU9gl)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[16\]P\. W\. Holland\(1986\)Statistics and causal inference\.Journal of the American Statistical Association81\(396\),pp\. 945–960\.External Links:[Document](https://dx.doi.org/10.1080/01621459.1986.10478354),[Link](https://doi.org/10.1080/01621459.1986.10478354)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p2.1),[§5](https://arxiv.org/html/2610.33145#S5.p4.1)\.
- \[17\]X\. Huang, M\. Liu, S\. Belongie, and J\. Kautz\(2018\)Multimodal unsupervised image\-to\-image translation\.InEuropean Conference on Computer Vision \(ECCV\),External Links:[Link](https://www.ecva.net/papers/eccv_2018/papers_ECCV/html/Xun_Huang_Multimodal_Unsupervised_Image-to-image_ECCV_2018_paper.php)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[18\]D\. H\. Hubel and T\. N\. Wiesel\(1962\)Receptive fields, binocular interaction and functional architecture in the cat’s visual cortex\.The Journal of Physiology160\(1\),pp\. 106–154\.External Links:[Document](https://dx.doi.org/10.1113/jphysiol.1962.sp006837)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1)\.
- \[19\]C\. L\. Hurwitz, A\. Srivastava, K\. Xu, J\. Jude, M\. Perich, L\. E\. Miller, and M\. H\. Hennig\(2021\)Targeted neural dynamical modeling\.External Links:[Link](https://openreview.net/forum?id=dnDkuSzNh8)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[20\]J\. J\. Jun, N\. A\. Steinmetz, J\. H\. Siegle, D\. J\. Denman,et al\.\(2017\)Fully integrated silicon probes for high\-density recording of neural activity\.Nature551,pp\. 232–236\.External Links:[Document](https://dx.doi.org/10.1038/nature24636)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[21\]J\. Kapoor, A\. Schulz, J\. Vetter, F\. Pei, R\. Gao, and J\. H\. Macke\(2024\)Latent diffusion for neural spiking data\.InAdvances in Neural Information Processing Systems 37,pp\. 118119–118154\.External Links:[Document](https://dx.doi.org/10.52202/079017-3751)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p1.1)\.
- \[22\]I\. Khemakhem, D\. Kingma, R\. Monti, and A\. Hyvarinen\(2020\)Variational autoencoders and nonlinear ICA: a unifying framework\.InProceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics,Proceedings of Machine Learning Research, Vol\.108,pp\. 2207–2217\.External Links:[Link](https://proceedings.mlr.press/v108/khemakhem20a.html)Cited by:[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p2.1)\.
- \[23\]H\. Kim and A\. Mnih\(2018\)Disentangling by factorising\.InProceedings of the 35th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.80,pp\. 2649–2658\.External Links:[Link](https://proceedings.mlr.press/v80/kim18b.html)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[24\]D\. Kobak, W\. Brendel, C\. Constantinidis, C\. E\. Feierstein, A\. Kepecs, Z\. F\. Mainen, X\. Qi, R\. Romo, N\. Uchida, and C\. K\. Machens\(2016\)Demixed principal component analysis of neural population data\.eLife5,pp\. e10989\.External Links:[Document](https://dx.doi.org/10.7554/eLife.10989)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p1.1)\.
- \[25\]J\. W\. Krakauer, A\. A\. Ghazanfar, A\. Gomez\-Marin, M\. A\. MacIver, and D\. Poeppel\(2017\)Neuroscience needs behavior: correcting a reductionist bias\.Neuron93\(3\),pp\. 480–490\.Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[26\]E\. Kropff, J\. E\. Carmichael, M\. Moser, and E\. I\. Moser\(2015\)Speed cells in the medial entorhinal cortex\.Nature523\(7561\),pp\. 419–424\.Cited by:[§4\.1](https://arxiv.org/html/2610.33145#S4.SS1.p1.1)\.
- \[27\]G\. Lample, N\. Zeghidour, N\. Usunier, A\. Bordes, L\. Denoyer, and M\. Ranzato\(2017\)Fader networks: manipulating images by sliding attributes\.InAdvances in Neural Information Processing Systems 30,pp\. 5967–5976\.External Links:[Link](https://proceedings.neurips.cc/paper/2017/hash/3fd60983292458bf7dee75f12d5e9e05-Abstract.html)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[28\]H\. Lee, H\. Tseng, J\. Huang, M\. K\. Singh, and M\. Yang\(2018\)Diverse image\-to\-image translation via disentangled representations\.InEuropean Conference on Computer Vision \(ECCV\),External Links:[Link](https://hsinyinglee.com/Project_page/DRIT/)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[29\]F\. Locatello, S\. Bauer, M\. Lucic, G\. Raetsch, S\. Gelly, B\. Schölkopf, and O\. Bachem\(2019\)Challenging common assumptions in the unsupervised learning of disentangled representations\.InProceedings of the 36th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.97,pp\. 4114–4124\.Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1),[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.p2.2),[§5](https://arxiv.org/html/2610.33145#S5.p4.1)\.
- \[30\]A\. Luo, M\. Henderson, L\. Wehbe, and M\. Tarr\(2023\)Brain diffusion for visual exploration: cortical discovery using large scale generative models\.InAdvances in Neural Information Processing Systems,Vol\.36\.External Links:[Document](https://dx.doi.org/10.52202/075280-3310),[Link](https://proceedings.neurips.cc/paper_files/paper/2023/hash/ef0c0a23a1a8219c4fc381614664df3e-Abstract-Conference.html)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1)\.
- \[31\]M\. W\. Mathis and A\. Mathis\(2026\)Joint modelling of brain and behaviour dynamics with artificial intelligence\.Nature Reviews Neuroscience27,pp\. 87–100\.External Links:[Document](https://dx.doi.org/10.1038/s41583-025-00996-1)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[32\]M\. W\. Mathis, A\. Perez Rotondo, E\. F\. Chang, A\. S\. Tolias, and A\. Mathis\(2024\)Decoding the brain: from neural representations to mechanistic models\.Cell187\(21\),pp\. 5814–5832\.External Links:[Document](https://dx.doi.org/10.1016/j.cell.2024.08.051)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[33\]J\. D\. McCart, A\. R\. Sedler, C\. Versteeg,et al\.\(2024\)Diffusion\-based generation of neural activity from disentangled latent codes\.External Links:2407\.21195,[Document](https://dx.doi.org/10.48550/arXiv.2407.21195)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p1.1)\.
- \[34\]M\. J\. McGinley, M\. Vinck, J\. Reimer, R\. Batista\-Brito, E\. Zagha, C\. R\. Cadwell, A\. S\. Tolias, J\. A\. Cardin, and D\. A\. McCormick\(2015\)Waking state: rapid variations modulate neural and behavioral responses\.Neuron87\(6\),pp\. 1143–1161\.Cited by:[§5](https://arxiv.org/html/2610.33145#S5.p3.1)\.
- \[35\]K\. Meng, D\. Bau, A\. Andonian, and Y\. Belinkov\(2022\)Locating and editing factual associations in GPT\.InAdvances in Neural Information Processing Systems,Vol\.35,pp\. 17359–17372\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/hash/6f1d43d5a82a37e89b0665b33bf3a182-Abstract-Conference.html)Cited by:[§A\.8](https://arxiv.org/html/2610.33145#A1.SS8.p1.1)\.
- \[36\]K\. Meng, A\. Sen Sharma, A\. Andonian, Y\. Belinkov, and D\. Bau\(2023\)Mass editing memory in a transformer\.InThe Eleventh International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=MkbcAHIYgyS)Cited by:[§A\.8](https://arxiv.org/html/2610.33145#A1.SS8.p1.1)\.
- \[37\]M\. Mirza and S\. Osindero\(2014\)Conditional generative adversarial nets\.External Links:1411\.1784,[Link](https://arxiv.org/abs/1411.1784)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p2.1)\.
- \[38\]S\. Musall, M\. T\. Kaufman, A\. L\. Juavinett, S\. Gluf, and A\. K\. Churchland\(2019\)Single\-trial neural dynamics are dominated by richly varied movements\.Nature Neuroscience22,pp\. 1677–1686\.External Links:[Document](https://dx.doi.org/10.1038/s41593-019-0502-4)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p2.1),[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[39\]C\. M\. Niell and M\. P\. Stryker\(2010\)Modulation of visual responses by behavioral state in mouse visual cortex\.Neuron65\(4\),pp\. 472–479\.External Links:[Document](https://dx.doi.org/10.1016/j.neuron.2010.01.033)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1),[§4\.3](https://arxiv.org/html/2610.33145#S4.SS3.p1.1),[§5](https://arxiv.org/html/2610.33145#S5.p3.1)\.
- \[40\]J\. O’Keefe and J\. Dostrovsky\(1971\)The hippocampus as a spatial map: preliminary evidence from unit activity in the freely\-moving rat\.\.Brain research\.Cited by:[§4\.1](https://arxiv.org/html/2610.33145#S4.SS1.p1.1)\.
- \[41\]J\. M\. Pakan, V\. Francioni, and N\. L\. Rochefort\(2018\)Action and learning shape the activity of neuronal circuits in the visual cortex\.Current opinion in neurobiology52,pp\. 88–97\.Cited by:[§5](https://arxiv.org/html/2610.33145#S5.p3.1)\.
- \[42\]C\. Pandarinath, D\. J\. O’Shea, J\. Collins, R\. Jozefowicz, S\. D\. Stavisky, J\. C\. Kao, E\. M\. Trautmann, M\. T\. Kaufman, S\. I\. Ryu, L\. R\. Hochberg, J\. M\. Henderson, K\. V\. Shenoy, L\. F\. Abbott, and D\. Sussillo\(2018\)Inferring single\-trial neural population dynamics using sequential auto\-encoders\.Nature Methods15\(10\),pp\. 805–815\.External Links:ISSN 1548\-7091, 1548\-7105,[Document](https://dx.doi.org/10.1038/s41592-018-0109-9)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p1.1),[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.p4.2)\.
- \[43\]F\. Pei, J\. Ye, D\. Zoltowski, A\. Wu, R\. H\. Chowdhury, H\. Sohn, J\. E\. O’Doherty, K\. V\. Shenoy, M\. T\. Kaufman, M\. Churchland, M\. Jazayeri, L\. E\. Miller, J\. Pillow, I\. M\. Park, E\. L\. Dyer, and C\. Pandarinath\(2022\)Neural latents benchmark ’21: evaluating latent variable models of neural population activity\.External Links:2109\.04463,[Link](https://arxiv.org/abs/2109.04463)Cited by:[Appendix D](https://arxiv.org/html/2610.33145#A4.p3.1),[§4\.2](https://arxiv.org/html/2610.33145#S4.SS2.p1.1)\.
- \[44\]Z\. Piran, N\. Cohen, Y\. Hoshen, and M\. Nitzan\(2024\)Disentanglement of single\-cell data with biolord\.42\(11\),pp\. 1678–1683\.External Links:ISSN 1546\-1696,[Link](https://www.nature.com/articles/s41587-023-02079-x),[Document](https://dx.doi.org/10.1038/s41587-023-02079-x)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p3.1)\.
- \[45\]O\. G\. Sani, H\. Abbaspourazad, Y\. T\. Wong, B\. Pesaran, and M\. M\. Shanechi\(2021\)Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification\.24\(1\),pp\. 140–149\.External Links:ISSN 1546\-1726,[Document](https://dx.doi.org/10.1038/s41593-020-00733-0)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[46\]O\. G\. Sani, B\. Pesaran, and M\. M\. Shanechi\(2024\)Dissociative and prioritized modeling of behaviorally relevant neural dynamics using recurrent neural networks\.27\(10\),pp\. 2033–2045\.External Links:ISSN 1546\-1726,[Document](https://dx.doi.org/10.1038/s41593-024-01731-2)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[47\]A\. Sauer and A\. Geiger\(2021\)Counterfactual generative networks\.External Links:2101\.06046,[Link](https://arxiv.org/abs/2101.06046)Cited by:[§5](https://arxiv.org/html/2610.33145#S5.p4.1)\.
- \[48\]S\. Schneider, R\. G\. Laiz, A\. Filippova, M\. Frey, and M\. W\. Mathis\(2025\)Time\-series attribution maps with regularized contrastive learning\.External Links:2502\.12977,[Link](https://arxiv.org/abs/2502.12977)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1)\.
- \[49\]S\. Schneider, J\. H\. Lee, and M\. W\. Mathis\(2023\)Learnable latent embeddings for joint behavioural and neural analysis\.617\(7960\),pp\. 360–368\.External Links:ISSN 1476\-4687,[Link](https://www.nature.com/articles/s41586-023-06031-6),[Document](https://dx.doi.org/10.1038/s41586-023-06031-6)Cited by:[§4\.2](https://arxiv.org/html/2610.33145#S4.SS2.p2.1)\.
- \[50\]J\. H\. Siegle, X\. Jia, S\. Durand, S\. Gale, C\. Bennett, N\. Graddis, G\. Heller, T\. K\. Ramirez, H\. Choi, J\. A\. Luviano, P\. A\. Groblewski, R\. Ahmed, A\. Arkhipov, A\. Bernard, Y\. N\. Billeh, D\. Brown, M\. A\. Buice, N\. Cain, S\. Caldejon, L\. Casal, A\. Cho, M\. Chvilicek, T\. C\. Cox, K\. Dai, D\. J\. Denman, S\. E\. J\. De Vries, R\. Dietzman, L\. Esposito, C\. Farrell, D\. Feng, J\. Galbraith, M\. Garrett, E\. C\. Gelfand, N\. Hancock, J\. A\. Harris, R\. Howard, B\. Hu, R\. Hytnen, R\. Iyer, E\. Jessett, K\. Johnson, I\. Kato, J\. Kiggins, S\. Lambert, J\. Lecoq, P\. Ledochowitsch, J\. H\. Lee, A\. Leon, Y\. Li, E\. Liang, F\. Long, K\. Mace, J\. Melchior, D\. Millman, T\. Mollenkopf, C\. Nayan, L\. Ng, K\. Ngo, T\. Nguyen, P\. R\. Nicovich, K\. North, G\. K\. Ocker, D\. Ollerenshaw, M\. Oliver, M\. Pachitariu, J\. Perkins, M\. Reding, D\. Reid, M\. Robertson, K\. Ronellenfitch, S\. Seid, C\. Slaughterbeck, M\. Stoecklin, D\. Sullivan, B\. Sutton, J\. Swapp, C\. Thompson, K\. Turner, W\. Wakeman, J\. D\. Whitesell, D\. Williams, A\. Williford, R\. Young, H\. Zeng, S\. Naylor, J\. W\. Phillips, R\. C\. Reid, S\. Mihalas, S\. R\. Olsen, and C\. Koch\(2021\)Survey of spiking in the mouse visual system reveals functional hierarchy\.Nature592\(7852\),pp\. 86–92\.External Links:ISSN 0028\-0836, 1476\-4687,[Document](https://dx.doi.org/10.1038/s41586-020-03171-x)Cited by:[§4\.3](https://arxiv.org/html/2610.33145#S4.SS3.p1.1)\.
- \[51\]K\. Sohn, H\. Lee, and X\. Yan\(2015\)Learning structured output representation using deep conditional generative models\.InAdvances in Neural Information Processing Systems,Vol\.28\.Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p2.1)\.
- \[52\]N\. A\. Steinmetz, P\. Zatka\-Haas, M\. Carandini, and K\. D\. Harris\(2019\)Distributed coding of choice, action and engagement across the mouse brain\.Nature576,pp\. 266–273\.External Links:[Document](https://dx.doi.org/10.1038/s41586-019-1787-x)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1),[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[53\]C\. Stringer, M\. Pachitariu, N\. Steinmetz, C\. B\. Reddy, M\. Carandini, and K\. D\. Harris\(2019\)Spontaneous behaviors drive multidimensional, brainwide activity\.Science364\(6437\),pp\. eaav7893\.External Links:[Document](https://dx.doi.org/10.1126/science.aav7893)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p2.1),[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[54\]M\. Sundararajan, A\. Taly, and Q\. Yan\(2017\)Axiomatic attribution for deep networks\.External Links:1703\.01365,[Link](https://arxiv.org/abs/1703.01365)Cited by:[§A\.4](https://arxiv.org/html/2610.33145#A1.SS4.p1.1),[§1](https://arxiv.org/html/2610.33145#S1.p4.1),[§3\.3](https://arxiv.org/html/2610.33145#S3.SS3.p1.1)\.
- \[55\]X\. Tang, H\. Shen, S\. Zhao, N\. Li, and J\. Liu\(2023\)Flexible brain–computer interfaces\.Nature Electronics6\(2\),pp\. 109–118\.External Links:[Document](https://dx.doi.org/10.1038/s41928-022-00913-9)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1)\.
- \[56\]J\. S\. Taube, R\. U\. Muller, and J\. B\. Ranck\(1990\)Head\-direction cells recorded from the postsubiculum in freely moving rats\. i\. description and quantitative analysis\.Journal of Neuroscience10\(2\),pp\. 420–435\.Cited by:[§4\.1](https://arxiv.org/html/2610.33145#S4.SS1.p1.1)\.
- \[57\]N\. Tishby, F\. C\. Pereira, and W\. Bialek\(2000\)The information bottleneck method\.arXiv\.External Links:[Document](https://dx.doi.org/10.48550/ARXIV.PHYSICS/0004057)Cited by:[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.SSS0.Px1.p1.1)\.
- \[58\]W\. Truccolo, U\. T\. Eden, M\. R\. Fellows, J\. P\. Donoghue, and E\. N\. Brown\(2005\)A Point Process Framework for Relating Neural Spiking Activity to Spiking History, Neural Ensemble, and Extrinsic Covariate Effects\.Journal of Neurophysiology93\(2\),pp\. 1074–1089\.External Links:ISSN 0022\-3077, 1522\-1598,[Document](https://dx.doi.org/10.1152/jn.00697.2004)Cited by:[§3\.1](https://arxiv.org/html/2610.33145#S3.SS1.p4.2)\.
- \[59\]S\. Tseng, S\. N\. Chettih, C\. Arlt, R\. Barroso\-Luque, and C\. D\. Harvey\(2022\)Shared and specialized coding across posterior cortical areas for dynamic navigation decisions\.Neuron110\(15\),pp\. 2484–2502\.e16\.External Links:[Document](https://dx.doi.org/10.1016/j.neuron.2022.05.012)Cited by:[§A\.14](https://arxiv.org/html/2610.33145#A1.SS14.p3.1)\.
- \[60\]A\. E\. Urai, B\. Doiron, A\. M\. Leifer, and A\. K\. Churchland\(2022\)Large\-scale neural recordings call for new insights to link brain and behavior\.Nature neuroscience25\(1\),pp\. 11–19\.External Links:[Document](https://dx.doi.org/10.1038/s41593-021-00980-9)Cited by:[§1](https://arxiv.org/html/2610.33145#S1.p1.1),[§1](https://arxiv.org/html/2610.33145#S1.p2.1),[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[61\]P\. Vahidi, O\. G\. Sani, and M\. M\. Shanechi\(2024\)Modeling and dissociation of intrinsic and input\-driven neural population dynamics underlying behavior\.121\(7\),pp\. e2212887121\.External Links:[Link](https://www.pnas.org/doi/10.1073/pnas.2212887121),[Document](https://dx.doi.org/10.1073/pnas.2212887121)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[62\]Y\. Wang, J\. Yu, C\. Li, W\. Li, and A\. Wu\(2026\)Uncovering semantic selectivity of latent groups in higher visual cortex with mutual information\-guided diffusion\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=pWX9PUbqPj)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p4.1)\.
- \[63\]D\. Yi, H\. Dong, M\. J\. Higley, A\. Churchland, and S\. Saxena\(2025\)Shared\-AE: automatic identification of shared subspaces in high\-dimensional neural and behavioral activity\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=zXCnIyX9MG)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p2.1)\.
- \[64\]B\. M\. Yu, J\. P\. Cunningham, G\. Santhanam, S\. I\. Ryu, K\. V\. Shenoy, and M\. Sahani\(2009\)Gaussian\-process factor analysis for low\-dimensional single\-trial analysis of neural population activity\.Journal of Neurophysiology102\(1\),pp\. 614–635\.External Links:[Document](https://dx.doi.org/10.1152/jn.90941.2008)Cited by:[§2](https://arxiv.org/html/2610.33145#S2.p1.1)\.
- \[65\]Y\. Zhang, Y\. Wang, M\. Azabou, A\. Andre, Z\. Wang, H\. Lyu, T\. I\. B\. Laboratory, E\. Dyer, L\. Paninski, and C\. Hurwitz\(2025\)Neural encoding and decoding at scale\.External Links:2504\.08201,[Link](https://arxiv.org/abs/2504.08201)Cited by:[Appendix D](https://arxiv.org/html/2610.33145#A4.p6.1),[§2](https://arxiv.org/html/2610.33145#S2.p1.1),[§3\.4](https://arxiv.org/html/2610.33145#S3.SS4.p1.1)\.
- \[66\]D\. Zhou and X\. Wei\(2020\)Learning identifiable and interpretable latent models of high\-dimensional neural activity using pi\-VAE\.arXiv\.External Links:[Link](http://arxiv.org/abs/2011.04798),[Document](https://dx.doi.org/10.48550/arXiv.2011.04798),2011\.04798 \[stat\]Cited by:[Appendix D](https://arxiv.org/html/2610.33145#A4.p5.1),[§2](https://arxiv.org/html/2610.33145#S2.p1.1),[§3\.4](https://arxiv.org/html/2610.33145#S3.SS4.p1.1)\.

## Appendix AAdditional experimental and implementation details

### A\.1Additional RatInABox Figures

Withheld Behavior Flows Through ResidualWe perform a diagnostic test to verify that held\-out behavior is routed through the residual latent\. We train three models on our RatInABox dataset, each receiving two observed labels while the third behavioral variable is held out\. The three settings correspond to withholding position, head direction, or speed\. In Fig\.[5](https://arxiv.org/html/2610.33145#A3.F5), we visualize the residual latent across all models and find that the residual is consistently organized by the withheld behavior, while the observed behaviors are not strongly separated\. In Fig\.[6](https://arxiv.org/html/2610.33145#A3.F6), we quantify the decodability of each label from the residual latent, again showing that the withheld behavior is more decodable from the residual than the observed behaviors\. Together, these support our findings that the residual latent carries information primarily about unobserved source\-specific factors\.

RatInABox Attribution Heatmaps\.In Fig\.[9](https://arxiv.org/html/2610.33145#A3.F9), we plot attribution heatmaps learned by our model for all RatInABox neurons\. We include both the standard model, and the unregularized model without the noisy residual bottleneck\. We observe the expected block structure: in the standard model, position, head\-direction, and residual attributions are concentrated on place, head\-direction, and speed cells, respectively\. In contrast, the unregularized residual control shows broader residual attribution and weaker factor specificity\.

Example Counterfactual Rates across Cell Types\. In Fig\.[10](https://arxiv.org/html/2610.33145#A3.F10)we plot the firing rates of three example cells, one from each cell type, before and after both counterfactual edits\. In the position\-shift edit, the model predicts the target place\-cell rate while maintaining the source head\-direction and speed\-cell rates\. In the head\-direction edit, generated rates track the target head\-direction response while retaining place\- and speed\-cell structure\. All three cell types are compared with their simulator target rates\.

Ablated Model Position Shift DecodingIn Fig\.[11](https://arxiv.org/html/2610.33145#A3.F11), we plot the source and target position shift movement trajectories, alongside our model’s counterfactual prediction and our ablated model’s counterfactual prediction\. To extract trajectories from our model’s predictions, we train a position decoder on real neural activity, and then decode from both model’s predictions\. While our model successfully encodes the target trajectory, the ablated model decoded trajectory looks much less similar to the target\. Quantitative decoding results for this evaluation are in Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)e\.

### A\.2Additional Macaque Reaching Figures

In Fig\.[12](https://arxiv.org/html/2610.33145#A3.F12)we plot the decodingR2R^\{2\}of our ablated model against both source and target trajectories for our macaque reaching edit task\. Across behaviors, we observe that the ablated models’ encoded neural activity remains source like \- decoding each behavior reveals highR2R^\{2\}against the original behaviors, while low, negativeR2R^\{2\}against the edit targets\. This is further visible in Fig\.[13](https://arxiv.org/html/2610.33145#A3.F13), where we show four example trajectories including the original, target, and ablated model prediction through decoding\. Again, we observe that the ablated model encodes the source trial’s trajectory, failing to perform the counterfactual edit\. This further demonstrates that without the noisy information bottleneck from regularization, editing can fail \- likely due to behavior specific information flowing through the residual\.

### A\.3Additional Allen Visual Coding Figures

Orientation Attribution HeatmapsIn Fig\.[14](https://arxiv.org/html/2610.33145#A3.F14), we plot the orientation attribution for all units in session 756029989, sorted by general orientation selectivity index \(gOSI\) for both regularized and un\-regularized models\. We see a trend with higher attributed units having greater gOSI values\. Colour corresponds to normalized attribution, whose calculation is detailed in Appendix[A\.4](https://arxiv.org/html/2610.33145#A1.SS4)\.

### A\.4Integrated Gradients

For a trained model and clean representationz=\(z1,…,zK,u\)z=\(z\_\{1\},\\ldots,z\_\{K\},u\), letgℓ​\(z\)g\_\{\\ell\}\(z\)be neuronℓ\\ell’s generator output before the softplus rate transformation\. We compute Integrated Gradients\[[54](https://arxiv.org/html/2610.33145#bib.bib31)\]using a joint baseline over labels and the residual latent,

b=\(by1,…,byK,bres\)b=\(b\_\{y\_\{1\}\},\\ldots,b\_\{y\_\{K\}\},b\_\{\\mathrm\{res\}\}\)whereby=\(by1,…,byK\)b\_\{y\}=\(b\_\{y\_\{1\}\},\\ldots,b\_\{y\_\{K\}\}\)andbresb\_\{\\mathrm\{res\}\}are the label and residual baselines respectively \(we use zero baselines unless otherwise stated\)\.

We define a straight\-line interpolation path in the joint space:

\(y⁡\(α\),u⁡\(α\)\)=\(by,bres\)\+α⁡\(\(y,u\)−\(by,bres\)\),α∈\[0,1\],\(y\(\\alpha\),u\(\\alpha\)\)=\(b\_\{y\},b\_\{\\mathrm\{res\}\}\)\+\\alpha\\big\(\(y,u\)\-\(b\_\{y\},b\_\{\\mathrm\{res\}\}\)\\big\),\\quad\\alpha\\in\[0,1\],and the corresponding latent path

z⁡\(α\)=\(ϕ1​\(y1​\(α\)\),…,ϕK​\(yK​\(α\)\),u⁡\(α\)\)\.z\(\\alpha\)=\(\\phi\_\{1\}\(y\_\{1\}\(\\alpha\)\),\\ldots,\\phi\_\{K\}\(y\_\{K\}\(\\alpha\)\),u\(\\alpha\)\)\.
The label\-level attribution for factoryay\_\{a\}is defined via the chain rule through the embeddingza=ϕa​\(ya\)z\_\{a\}=\\phi\_\{a\}\(y\_\{a\}\)as

IGℓ,ya​\(z,b\)=\(ya−bya\)⊙∫01∇yagℓ​\(z⁡\(α\)\)​𝑑α\.\\mathrm\{IG\}\_\{\\ell,y\_\{a\}\}\(z;b\)=\(y\_\{a\}\-b\_\{y\_\{a\}\}\)\\odot\\int\_\{0\}^\{1\}\\nabla\_\{y\_\{a\}\}g\_\{\\ell\}\\big\(z\(\\alpha\)\\big\)\\,d\\alpha\.\(12\)
The residual attribution is computed with respect touualong the same path:

IGℓ,res​\(z,b\)=\(u−bres\)⊙∫01∇ugℓ​\(z⁡\(α\)\)​𝑑α\.\\mathrm\{IG\}\_\{\\ell,\\mathrm\{res\}\}\(z;b\)=\(u\-b\_\{\\mathrm\{res\}\}\)\\odot\\int\_\{0\}^\{1\}\\nabla\_\{u\}g\_\{\\ell\}\\big\(z\(\\alpha\)\\big\)\\,d\\alpha\.\(13\)
ForS=32S=32randomly selected source samples per fold, held fixed across theλ\\lambdasweep within that fold, we average signed attributions before taking magnitudes:

Aℓ,ya=1qa​‖1S​∑s=1SIGℓ,ya​\(z\(s\),b\)‖1,Aℓ,res=1m​‖1S​∑s=1SIGℓ,res​\(z\(s\),b\)‖1,A\_\{\\ell,y\_\{a\}\}=\\frac\{1\}\{q\_\{a\}\}\\left\\\|\\frac\{1\}\{S\}\\sum\_\{s=1\}^\{S\}\\mathrm\{IG\}\_\{\\ell,y\_\{a\}\}\(z^\{\(s\)\};b\)\\right\\\|\_\{1\},\\qquad A\_\{\\ell,\\mathrm\{res\}\}=\\frac\{1\}\{m\}\\left\\\|\\frac\{1\}\{S\}\\sum\_\{s=1\}^\{S\}\\mathrm\{IG\}\_\{\\ell,\\mathrm\{res\}\}\(z^\{\(s\)\};b\)\\right\\\|\_\{1\},\(14\)whereqaq\_\{a\}is the number of input dimensions of factoryay\_\{a\}andmmis the residual dimension\. We use Captum’s Integrated Gradients with 50 Gauss–Legendre integration steps and zero baselines\.

For attribution visualizations in Fig\.[9](https://arxiv.org/html/2610.33145#A3.F9)and Fig\.[14](https://arxiv.org/html/2610.33145#A3.F14), and for the attribution summaries in Fig\.[3](https://arxiv.org/html/2610.33145#S4.F3)c,d, we min–max normalize the dataset\-level attribution scores for each factor across neurons to the range\[0,1\]\[0,1\]\.

Aℓ,yanorm​\(z\)=Aℓ,ya​\(z\)−mini⁡Ai,ya​\(z\)maxi⁡Ai,ya​\(z\)−mini⁡Ai,ya​\(z\)A^\{\\mathrm\{norm\}\}\_\{\\ell,y\_\{a\}\}\(z\)=\\frac\{A\_\{\\ell,y\_\{a\}\}\(z\)\-\\min\_\{i\}A\_\{i,y\_\{a\}\}\(z\)\}\{\\max\_\{i\}A\_\{i,y\_\{a\}\}\(z\)\-\\min\_\{i\}A\_\{i,y\_\{a\}\}\(z\)\}

### A\.5Counterfactual Semantics

The edited sample has a causal interpretation only under additional assumptions\. Suppose the data are generated from observed factorsY=\(Y1,…,YK\)Y=\(Y\_\{1\},\\ldots,Y\_\{K\}\)and a source\-specific unobserved factorHH, and suppose an intervention onYaY\_\{a\}leaves\(Y−a,H\)\(Y\_\{\-a\},H\)fixed\. If the true conditional spike distribution factorizes as

p⋆​\(x∣y,h\)=∏ℓ=1nPoisson⁡\(xℓ∣Gℓ⋆​\(ϕ1⋆​\(y1\),…,ϕK⋆​\(yK\),ψ⋆​\(h\)\)\),p^\{\\star\}\(x\\mid y,h\)=\\prod\_\{\\ell=1\}^\{n\}\\mathrm\{Poisson\}\\left\(x\_\{\\ell\}\\mid G\_\{\\ell\}^\{\\star\}\(\\phi\_\{1\}^\{\\star\}\(y\_\{1\}\),\\ldots,\\phi\_\{K\}^\{\\star\}\(y\_\{K\}\),\\psi^\{\\star\}\(h\)\)\\right\),\(15\)if Neural\-DISCO recovers these maps, and if the residual encoder recoversψ⋆​\(h\)\\psi^\{\\star\}\(h\)for the source sample, then replacingϕa​\(ya\)\\phi\_\{a\}\(y\_\{a\}\)byϕa​\(y^a\)\\phi\_\{a\}\(\\hat\{y\}\_\{a\}\)while holding\(z−a,u\)\(z\_\{\-a\},u\)fixed produces samples from

p⋆​\(x∣do⁡\(Ya=y^a\),Y−a=y−a,H=h\)\.p^\{\\star\}\\big\(x\\mid\\mathrm\{do\}\(Y\_\{a\}=\\hat\{y\}\_\{a\}\),Y\_\{\-a\}=y\_\{\-a\},H=h\\big\)\.\(16\)The learned residualu=ϕres​\(x\)u=\\phi\_\{\\mathrm\{res\}\}\(x\)is intended to encode source\-specific information associated withhh, but it is not the same object as the latent generative factorHH\. When these assumptions are not justified, Neural\-DISCO counterfactuals are model\-based predictions under the learned factorization\. Edits far outside the training support, or edits that violate natural dependencies among labels and internal state, should be treated as extrapolations\.

### A\.6Hyperparameter selection and reporting

We selected dataset\-specific hyperparameters using validation diagnostics and preliminary pilot runs, then fixed one configuration for each dataset–model pair before running the cross\-validation experiments reported in the main text\. The final configurations are reported in Table[12](https://arxiv.org/html/2610.33145#A1.T12)\. Theλ\\lambda\-selection criterion favors high validation edit fidelity and low supplied\-label decodability from the clean residual, averaged across ten validation folds\. Reconstruction loss provides an additional diagnostic; checkpoint selection uses validation loss\.

The unregularized condition is an architecture\-matched ablation in which the residual bottleneck is disabled by setting both the residual activation penalty and residual noise to zero\. This ablation tests the necessity of the bottleneck mechanism and was not used to choose the reported Neural\-DISCO hyperparameters\.

We also report sweeps over the residual regularization coefficientλ\\lambdato evaluate how bottleneck strength affects attribution and residual leakage, as in Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4), Fig\.[3](https://arxiv.org/html/2610.33145#S4.F3), and Fig\.[2](https://arxiv.org/html/2610.33145#S4.F2)\. These sweeps are sensitivity analyses rather than hyperparameter\-selection curves for the reported cross\-validation results\. Initialization and training use seed314159\+fold index314159\+\\text\{fold index\}, held fixed acrossλ\\lambda\. Fold indices run from 0 to 9\. The sweeps retain residual noiseσ=\.02\\sigma=\.02for RatInABox and Allen and\.04\.04for MC Maze, including atλ=0\\lambda=0; the zero\-penalty point therefore differs from theλ=σ=0\\lambda=\\sigma=0ablation\.

### A\.7Evaluation metrics

Unless otherwise stated, reported values are calculated on the test set, and averaged across the 10 CV\-fold models\. Uncertainty is reported as the standard error of the mean across folds\. If a metricmkm\_\{k\}is computed on fold\-modelsk∈\{1,…,10\}k\\in\\\{1,\\dots,10\\\}, we report withK=10K=10:

m¯=1K​∑k=1Kmk,SEM⁡\(m\)=smK,\\bar\{m\}=\\frac\{1\}\{K\}\\sum\_\{k=1\}^\{K\}m\_\{k\},\\qquad\\mathrm\{SEM\}\(m\)=\\frac\{s\_\{m\}\}\{\\sqrt\{K\}\},\(17\)wheresms\_\{m\}is the sample standard deviation of the fold\-wise values\.

##### RatInABox\.

For attribution, we treat the neuron\-level attribution score for a given factor as a classifier score for the corresponding synthetic cell type and report area under the ROC curve \(AUC\)\. Writingsis\_\{i\}for the attribution score of neuroniiandyi∈\{0,1\}y\_\{i\}\\in\\\{0,1\\\}for whether that neuron belongs to the relevant synthetic cell class, the AUC is

AUC=Pr⁡\(s\+\>s−\)\+12​Pr⁡\(s\+=s−\),\\mathrm\{AUC\}=\\Pr\(s^\{\+\}\>s^\{\-\}\)\+\\tfrac\{1\}\{2\}\\Pr\(s^\{\+\}=s^\{\-\}\),\(18\)wheres\+s^\{\+\}ands−s^\{\-\}denote scores sampled from the positive and negative classes, respectively\. Position attribution is evaluated against place\-cell labels, head\-direction attribution against head\-direction\-cell labels, and residual attribution against speed\-cell labels because speed is withheld from the observed\-factor embeddings\. We use scikit\-learn’s sklearn\.metrics\.roc\_auc\_score to compute AUC values empirically\. For counterfactual editing, we compare generated ratesr^ℓ​t\\hat\{r\}\_\{\\ell t\}with simulator ground\-truth edited target ratesrℓ​t⋆r^\{\\star\}\_\{\\ell t\}\. For a neuron groupℐ\\mathcal\{I\}, we report

Rℐ2=1\|ℐ\|​∑ℓ∈ℐ\[1−∑t=1T\(r^ℓ​t−rℓ​t⋆\)2∑t=1T\(rℓ​t⋆−r¯ℓ⋆\)2\],R^\{2\}\_\{\\mathcal\{I\}\}=\\frac\{1\}\{\|\\mathcal\{I\}\|\}\\sum\_\{\\ell\\in\\mathcal\{I\}\}\\left\[1\-\\frac\{\\sum\_\{t=1\}^\{T\}\(\\hat\{r\}\_\{\\ell t\}\-r^\{\\star\}\_\{\\ell t\}\)^\{2\}\}\{\\sum\_\{t=1\}^\{T\}\(r^\{\\star\}\_\{\\ell t\}\-\\bar\{r\}^\{\\star\}\_\{\\ell\}\)^\{2\}\}\\right\],\(19\)where

r¯ℓ⋆=1T​∑t=1Trℓ​t⋆\.\\bar\{r\}^\{\\star\}\_\{\\ell\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}r^\{\\star\}\_\{\\ell t\}\.We also check qualitatively that edited\-factor responses move toward the target while non\-target and withheld components are preserved\.

##### MC Maze\.

Because there is no ground\-truth counterfactual neural response in the real recordings, we evaluate generated activity through downstream behavioral decoding\. A CEBRA decoder trained only on real training\-set spikes predicts the target hand position, hand velocity, and eye position from source spikes, target spikes, and generated activity\. For a decoded trajectoryy^t\\hat\{y\}\_\{t\}and target trajectoryyt⋆y\_\{t\}^\{\\star\}, we report the coefficient of determination

R2=1d​∑j=1d\[1−∑t=1T\(y^t​j−yt​j⋆\)2∑t=1T\(yt​j⋆−y¯j⋆\)2\],R^\{2\}=\\frac\{1\}\{d\}\\sum\_\{j=1\}^\{d\}\\left\[1\-\\frac\{\\sum\_\{t=1\}^\{T\}\(\\hat\{y\}\_\{tj\}\-y\_\{tj\}^\{\\star\}\)^\{2\}\}\{\\sum\_\{t=1\}^\{T\}\(y\_\{tj\}^\{\\star\}\-\\bar\{y\}\_\{j\}^\{\\star\}\)^\{2\}\}\\right\],\(20\)whered=2d=2is the number of coordinates andy¯j⋆=1T​∑t=1Tyt​j⋆\\bar\{y\}\_\{j\}^\{\\star\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}y\_\{tj\}^\{\\star\}\. Thus,R2R^\{2\}is averaged over coordinates\. Source\-spike decoding serves as a negative control and target\-spike decoding serves as an empirical reference for the decoder on real data\. For attribution, we again summarize performance by the AUC for distinguishing recorded units from synthetic constant\-rate control neurons using the corresponding factor\-attribution scores\.

##### Allen Visual Coding\.

For running\-speed edits, a decoder trained on real training\-set activity predicts running speed from generated responses, and we report root mean squared error \(RMSE\) relative to the edited target speedvt⋆v\_\{t\}^\{\\star\}:

RMSE=1T​∑t=1T\(v^t−vt⋆\)2\.\\mathrm\{RMSE\}=\\sqrt\{\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\(\\hat\{v\}\_\{t\}\-v\_\{t\}^\{\\star\}\)^\{2\}\}\.\(21\)To test preservation of withheld stimulus information, we train a logistic\-regression decoder on real activity to predict temporal frequency and report balanced accuracy on edited responses,

BA=1C​∑c=1CTPcTPc\+FNc,\\mathrm\{BA\}=\\frac\{1\}\{C\}\\sum\_\{c=1\}^\{C\}\\frac\{\\mathrm\{TP\}\_\{c\}\}\{\\mathrm\{TP\}\_\{c\}\+\\mathrm\{FN\}\_\{c\}\},\(22\)whereCCis the number of temporal\-frequency classes\. For attribution, we use two complementary summaries: AUC for classifying units with significant running modulation from running\-attribution scores, and the Pearson correlation between orientation\-attribution scoresaia\_\{i\}and the dataset\-provided global orientation selectivity indicesgig\_\{i\},

ρ⁡\(a,g\)=∑i\(ai−a¯\)​\(gi−g¯\)∑i\(ai−a¯\)2​∑i\(gi−g¯\)2\.\\rho\(a,g\)=\\frac\{\\sum\_\{i\}\(a\_\{i\}\-\\bar\{a\}\)\(g\_\{i\}\-\\bar\{g\}\)\}\{\\sqrt\{\\sum\_\{i\}\(a\_\{i\}\-\\bar\{a\}\)^\{2\}\}\\sqrt\{\\sum\_\{i\}\(g\_\{i\}\-\\bar\{g\}\)^\{2\}\}\}\.\(23\)

### A\.8Edit–preservation score

Adapting harmonic\-mean summaries used in model\-editing evaluation\[[35](https://arxiv.org/html/2610.33145#bib.bib63),[36](https://arxiv.org/html/2610.33145#bib.bib64)\], we summarize edit fidelity and preservation as

SEP=2​E​PE\+P,\[x\]01=min⁡\{1,max⁡\{0,x\}\},S\_\{\\mathrm\{EP\}\}=\\frac\{2EP\}\{E\+P\},\\qquad\[x\]\_\{0\}^\{1\}=\\min\\\{1,\\max\\\{0,x\\\}\\\},\(24\)where\[x\]01\[x\]\_\{0\}^\{1\}clipsxxto\[0,1\]\[0,1\], andSEP=0S\_\{\\mathrm\{EP\}\}=0whenE\+P=0E\+P=0\. Both components lie in\[0,1\]\[0,1\]; a high score requires both a successful edit and preservation of the remaining information\. We retain the underlying measurements in Tables[1](https://arxiv.org/html/2610.33145#A1.T1)and[2](https://arxiv.org/html/2610.33145#A1.T2)so that failures of either component remain visible\.

For the RatInABox edits, speed is withheld from the observed\-factor embeddings\.E=\[Redited2\]01E=\[R^\{2\}\_\{\\mathrm\{edited\}\}\]\_\{0\}^\{1\}evaluates the edited cell population against the exact simulator target\. Preservation is the mean of the clippedR2R^\{2\}scores for the two other populations,

P=12​\(\[Rother2\]01\+\[Rspeed2\]01\)\.P=\\tfrac\{1\}\{2\}\\left\(\[R^\{2\}\_\{\\mathrm\{other\}\}\]\_\{0\}^\{1\}\+\[R^\{2\}\_\{\\mathrm\{speed\}\}\]\_\{0\}^\{1\}\\right\)\.\(25\)Thus, the head\-direction and speed populations enter preservation for a position edit\. In the head\-direction edit, head\-direction cells defineEE, while place and speed cells definePP; all components use the edit’s simulator target rates\. For the Allen speed\-to\-mean edit,

E=\[1−RMSEeditedRMSEsource\]01,P=\[BAedited−0\.2BAsource−0\.2\]01,E=\\left\[1\-\\frac\{\\mathrm\{RMSE\}\_\{\\mathrm\{edited\}\}\}\{\\mathrm\{RMSE\}\_\{\\mathrm\{source\}\}\}\\right\]\_\{0\}^\{1\},\\qquad P=\\left\[\\frac\{\\mathrm\{BA\}\_\{\\mathrm\{edited\}\}\-0\.2\}\{\\mathrm\{BA\}\_\{\\mathrm\{source\}\}\-0\.2\}\\right\]\_\{0\}^\{1\},\(26\)where both RMSE values are relative to the target speed andBA\\mathrm\{BA\}is temporal\-frequency balanced accuracy\. Chance accuracy is0\.20\.2for the five\-way decoder; we setP=0P=0if source accuracy is at or below chance\. We do not assignSEPS\_\{\\mathrm\{EP\}\}to the original MC Maze reach edit because all three measured factors are edited together, leaving no independently measured non\-edited factor to define preservation\.

The Gaussian GLM editor adds its fitted source\-to\-target label effect to the recorded response,xedit=xsource\+W⁡\(yedit−ysource\)x\_\{\\mathrm\{edit\}\}=x\_\{\\mathrm\{source\}\}\+W\(y\_\{\\mathrm\{edit\}\}\-y\_\{\\mathrm\{source\}\}\), with negative outputs clipped to zero\. The Poisson GLM instead predicts rates from the target labels,r^​\(y\)=exp⁡\(W​y\+b\)\\hat\{r\}\(y\)=\\exp\(Wy\+b\)\. In RatInABox, both GLMs use raw position and heading components\(x,y,hx,hy\)\(x,y,h\_\{x\},h\_\{y\}\)as covariates\. The zero residual control retains the trained Neural\-DISCO model and sets its source residual to zero at edit time\. Because the Gaussian GLM edits recorded spike counts, its output keeps spike\-count noise, whereas the other methods predict mean rates; all are scored against simulator mean rates\.

Table 1:RatInABox editing against exact simulator target trajectories\. PopulationR2R^\{2\}averages neuron\-wise scores\. All three populations are evaluated against simulator target rates in both edits\. Values and joint scores are computed per fold and reported as mean±\\pmSEM across ten folds\. Each component is clipped before computing the edit–preservation score\.We analyze 32 Allen Brain Observatory sessions, retaining units withSNR\>1\\mathrm\{SNR\}\>1and an ISI\-violation score<0\.5<0\.5\. Within each session, we use ten folds and edit running speed toward its mean in the fold’s training data\. All methods share speed and temporal\-frequency decoders trained on real activity in the same training fold\. Both GLMs are trained for 100 epochs with learning rate10−310^\{\-3\}, weight decay10−510^\{\-5\}, and seed314159\+fold index314159\+\\text\{fold index\}\.

Table 2:Allen speed editing and temporal\-frequency preservation across 32 sessions, using a common fold\-training mean target\. Metrics andSEPS\_\{\\mathrm\{EP\}\}are computed per fold and averaged over ten folds within each session; values are mean±\\pmSEM across the 32 session means\. Lower speed RMSE and higher TF balanced accuracy andSEPS\_\{\\mathrm\{EP\}\}are better\.Neural\-DISCO has the highest joint score for both RatInABox edits\. Zeroing its residual maintains edit fidelity but loses withheld speed information\. Across the 32 Allen sessions, Neural\-DISCO has a higher mean joint score than Gaussian GLM and exceeds it in 30 of 32 sessions\. The Gaussian editor has higher mean temporal\-frequency balanced accuracy, whereas Neural\-DISCO has lower mean speed RMSE; zeroing the residual reduces temporal\-frequency decoding to approximately chance\. These results distinguish fitting a requested target from preserving source information\.

### A\.9Editing toward inconsistent behavioral targets

To test whether decoding of generated activity tracks requested labels regardless of their consistency, we retain the source residual and the requested target hand\-position and eye\-position trajectories, but rotate the target two\-dimensional hand velocity\. Rotations preserve every velocity magnitude while making velocity increasingly inconsistent with the unchanged hand\-position trajectory\. The model is evaluated without training on the rotated targets\. Within each fold, we apply a fixed CEBRA decoder trained only on real activity to Poisson spike samples from the generated rates\. We average clockwise and counterclockwise results within each fold at intermediate angles, then report mean±\\pmSEM across ten folds\.

Table 3:MC Maze decoding under velocity rotations\. Position and eye targets remain unchanged\.Requested\-velocity decoding declines progressively as the supplied velocity becomes less compatible with the fixed position trajectory\. Hand and eye decoding also decline, although less strongly\. The control therefore shows degradation under inconsistent behavioral requests with collateral effects on the unchanged targets, rather than selective rejection of velocity alone\. This is evidence against indiscriminate decoder readout of supplied labels; it does not establish biological validity of the generated responses or causal effects of the edits\.

### A\.10Orientation and temporal\-frequency editing

For orientation edits, the observed factors are grating orientation and temporal frequency \(TF\), and running speed is withheld\. We rotate orientation by180∘180^\{\\circ\}across the same 32 sessions used for speed editing\. Models and evaluation decoders are trained on the training presentations in each fold, with model selection on separate validation presentations\. Edits are evaluated on the final drifting\-grating epoch\. Neural\-DISCO retains its source residual; the zero residual control sets this residual to zero at edit time\. The unregularized ablation, pi\-VAE and NEDS use the same orientation and TF inputs\. Training settings are given in Table[12](https://arxiv.org/html/2610.33145#A1.T12)\.

We compute edit fidelityEEas the mean chance\-corrected balanced accuracy over edited stimulus factors\. Chance levels are1/81/8for orientation and1/51/5for TF\. PreservationPPaverages the chance\-corrected accuracy retained relative to source decoding for each unedited stimulus factor and the running\-speed component\(RMSE0−RMSEedited\)/\(RMSE0−RMSEsource\)\(\\mathrm\{RMSE\}\_\{0\}\-\\mathrm\{RMSE\}\_\{\\mathrm\{edited\}\}\)/\(\\mathrm\{RMSE\}\_\{0\}\-\\mathrm\{RMSE\}\_\{\\mathrm\{source\}\}\), where the RMSEs are speed\-decoding errors on edited and source activity andRMSE0\\mathrm\{RMSE\}\_\{0\}is the error of always predicting the mean held\-out speed\. The speed component is zero when source decoding is no more accurate than this constant prediction\. Each component is clipped to\[0,1\]\[0,1\]before averaging, andSEP=2​E​P/\(E\+P\)S\_\{\\mathrm\{EP\}\}=2EP/\(E\+P\)\. When both stimulus factors are edited, preservation measures running speed alone\.

Table 4:Orientation editing and source preservation across 32 Allen sessions\. Metrics andSEPS\_\{\\mathrm\{EP\}\}are computed per fold and averaged over ten folds within each session; values are mean±\\pmSEM across the 32 session means\. All methods share decoders trained on real fold\-training activity and are evaluated on the same held\-out presentations and targets\. Speed RMSE measures error relative to source speed\. Lower speed RMSE and higher balanced accuracy andSEPS\_\{\\mathrm\{EP\}\}are better\.Neural\-DISCO has higher orientation fidelity and TF preservation than Gaussian GLM, pi\-VAE, NEDS and the unregularized ablation\. Poisson GLM and the zero residual control predict the stimulus labels accurately but do not preserve source\-speed information\.

##### Additional single\-session stimulus edits\.

In session 756029989, we also advance TF cyclically through1,2,4,8,151,2,4,8,15Hz or jointly edit orientation and TF\. Table[5](https://arxiv.org/html/2610.33145#A1.T5)reports the component and joint scores for these edits\.

Table 5:Additional Allen stimulus edits in session 756029989\. Values are mean±\\pmSEM across ten folds, withEE,PPandSEPS\_\{\\mathrm\{EP\}\}computed within each fold before aggregation\. All methods share decoders trained on real fold\-training activity and are evaluated on the same held\-out presentations and targets\. Higher values are better\.The component scores distinguish successful stimulus editing from preservation of source running speed\. Poisson GLM attains high edit fidelity but has a zero preservation score when orientation and TF are edited together\. It generates rates from the target stimulus labels without a source residual\. pi\-VAE and NEDS also have zero preservation scores for this joint edit\. Gaussian GLM has higher preservation than Neural\-DISCO for the joint edit, while Neural\-DISCO has higher edit fidelity and joint score\. Zeroing the residual yields a zero joint\-edit preservation score\.

### A\.11Decoding omitted and supplied factors from the clean residual

We trained three Allen models, each supplied with two factors while withholding the third, usingλ=\.01\\lambda=\.01\. We extracted the clean residual without injecting noise and trained probes only on residuals from the model’s training fold\. Categorical probes used a training\-fitted standardizer followed by balanced logistic regression; the speed probe used a training\-fitted standardizer and ridge regression\. Probes were evaluated on the corresponding held\-out validation fold\.

Table 6:Decoding from the clean Allen residual\. Each row represents a separately trained model; columns identify the factor decoded from its residual\. Bold diagonal entries measure retention of the omitted factor, while off\-diagonal entries measure decodability of supplied factors\. Values are mean±\\pmSEM across ten folds\. Balanced accuracies are reported without chance correction; speed uses rawR2R^\{2\}\.The residual preferentially retains omitted\-factor information, although some supplied\-factor information remains\. Each factor is more decodable when omitted from the observed inputs; supplied\-factor decoding includes TF balanced accuracies of\.252\.252and\.265\.265and speedR2=\.201±\.084R^\{2\}=\.201\\pm\.084\.

##### MC Maze residual decoding\.

We decode supplied hand position, hand velocity, and eye position from the clean residual using MLP decoders on held\-out data\. Table[7](https://arxiv.org/html/2610.33145#A1.T7)also includes source\-spike decoding as a reference, using Ridge probes \(α=1\\alpha=1\)\. Feature standardization and training of these source\-spike probes use only the training trials of each cross\-validation fold; evaluation uses held\-out validation trials\. The clean residual is the residual encoder output without injected noise\. Table[7](https://arxiv.org/html/2610.33145#A1.T7)shows low residual decoding relative to source spikes\. These decoding scores do not establish statistical independence\.

Table 7:Held\-out decoding of supplied MC Maze covariates from the clean residual, with source\-spike decoding as a reference\. Values are mean±\\pmSEM across ten folds\. Source\-spike decoding uses Ridge probes and uniform\-averageR2R^\{2\}across the two coordinates\. NegativeR2R^\{2\}is retained\.

### A\.12Robustness to mixed selectivity

Motivated by mixed selectivity in navigational coding\[[13](https://arxiv.org/html/2610.33145#bib.bib65)\], we tested dependence on cleanly separated tuning by introducing multiplicative interactions into the 300\-neuron RatInABox population\. The three 100\-neuron groups combined place×\\timesspeed, head direction×\\timesposition, and speed×\\timeshead direction tuning\. For a neuron with original rateria​\(t\)r\_\{i\}^\{a\}\(t\)and a fixed paired neuron from another population with raterπ⁡\(i\)b​\(t\)r\_\{\\pi\(i\)\}^\{b\}\(t\), we used

r~i​\(t\)=ci​ria​\(t\)​\[\(1−α\)\+α​rπ⁡\(i\)b​\(t\)rπ⁡\(i\)b¯\],\\widetilde\{r\}\_\{i\}\(t\)=c\_\{i\}r\_\{i\}^\{a\}\(t\)\\left\[\(1\-\\alpha\)\+\\alpha\\frac\{r\_\{\\pi\(i\)\}^\{b\}\(t\)\}\{\\overline\{r\_\{\\pi\(i\)\}^\{b\}\}\}\\right\],\(27\)where the overbar denotes the mean on the original training trajectory andcic\_\{i\}preserves the neuron’s mean training rate\. We retrained Neural\-DISCO at each interaction strength using the RatInABox architecture and training settings in Table[12](https://arxiv.org/html/2610.33145#A1.T12)\. The simulator supplied exact counterfactual rates by combining the requested target factor with the preserved source factors\.

Table[8](https://arxiv.org/html/2610.33145#A1.T8)reports the mean of neuron\-wiseR2R^\{2\}across all 300 neurons, followed by the edit–preservation score defined in Appendix[A\.8](https://arxiv.org/html/2610.33145#A1.SS8)\. Theα=0\\alpha=0row comes from separately trained models scored against target rates constructed for this study, so values differ slightly from Table[1](https://arxiv.org/html/2610.33145#A1.T1)\. For this experiment, edit fidelity uses the primary edited population, while preservation averages the other two populations’R2R^\{2\}scores after clipping each to\[0,1\]\[0,1\]\. All components are scored against the exact edited simulator rates: neurons recruited through an interaction may change even when their primary factor is unchanged\. Atα=1\\alpha=1, paired all\-populationR2R^\{2\}retained99\.17±0\.33%99\.17\\pm 0\.33\\%of itsα=0\\alpha=0value for position edits and99\.29±0\.24%99\.29\\pm 0\.24\\%for direction edits\. These results support robustness to the simulated interactions tested here\.

Table 8:Neural\-DISCO counterfactual performance under increasing multiplicative mixed selectivity\. Values are mean±\\pmSEM across ten independently trained cross\-validation folds\. TheR2R^\{2\}columns average neuron\-wise scores across the full population;SEPS\_\{\\mathrm\{EP\}\}combines primary\-population edit fidelity with performance on the other two populations\.
### A\.13Data and sampling requirements

We retrained Neural\-DISCO on nested, fold\-specific subsets of the training data, holding validation and test sets, counterfactual edits, and training settings fixed within each dataset\. Reduced RatInABox subsets were contiguous trajectory intervals that did not cross the held\-out validation interval\. MC Maze subsets contained complete trials balanced across 108 task conditions; Allen subsets contained complete presentations balanced across 40 orientation–temporal\-frequency conditions\. Evaluation decoders were held fixed across fractions: MC Maze used the full real training recording, and Allen used the complete real training fold\. The Allen scaling experiment edited speed to the mean of the full training fold, held fixed across generator\-training fractions\.

We computed retention relative to each fold’s full\-data retraining and then averaged across ten folds\. Retention is the reduced\-data/full\-data ratio forR2R^\{2\}, the ratio after subtracting the0\.200\.20chance level for Allen temporal\-frequency balanced accuracy, and the full\-data/reduced\-data ratio for Allen speed RMSE\. Table[9](https://arxiv.org/html/2610.33145#A1.T9)gives the smallest tested subset whose mean paired retention reached90%90\\%for every evaluated outcome\. This criterion applies to the fold mean and does not require each individual fold to exceed90%90\\%\.

Table 9:Tested training\-subset requirements for retaining at least90%90\\%of full\-data performance on every outcome\. Retention is mean±\\pmSEM across ten folds for the limiting outcome; the full\-fold reference equals100%100\\%by definition\. Median counts are per condition\. These training\-subset sizes assume the same held\-out data and full\-data evaluation decoders; the Allen speed target uses the full training fold\.For Allen, 170 presentations \(50%\) met both retention criteria\. At the next smaller tested subset \(37\.5%; 128 presentations\), TF retention was88\.4±3\.0%88\.4\\pm 3\.0\\%\. Retention is not clipped at100%100\\%\. At the next smaller tested subsets, the limiting retention fell to78\.9±3\.0%78\.9\\pm 3\.0\\%for RatInABox at 1\.8 minutes and89\.4±\.8%89\.4\\pm\.8\\%for MC Maze at 97 trials\. RatInABox tested five training fractions; MC Maze and Allen each tested nine\. These results apply to the tested training subsets and do not establish universal data requirements\.

Training used seed314159\+fold index314159\+\\text\{fold index\}and subset selection used seed271828\+fold index271828\+\\text\{fold index\}, with 300 epochs for RatInABox and 100 epochs for each real dataset\.

### A\.14Hypothesis generation from source\-dependent tuning

We use Allen Visual Coding recordings to test whether a score constructed without temporal\-frequency \(TF\) labels can prioritize neurons whose activity carries TF information in a later epoch\. The models are trained with running speed and orientation; TF is excluded from Neural\-DISCO training, donor matching, and score construction\. Scoring uses out\-of\-fold responses from the first two drifting\-grating epochs; its association with TF contribution is evaluated in the final epoch\. Changing only the residual shows how much a neuron’s predicted tuning varies with source context\. The score provides a ranking for follow\-up experiments, which we evaluate against TF contribution in a later epoch\.

For each source bin, we choose a donor from a different presentation in the same validation fold, matching presented orientation and within\-presentation time bin, with the nearest running speed\. We generate two predicted tuning curves over the same eight target orientations using the source running speed in both cases: one uses the source residual and the other the donor residual\. Thus, the observed inputs to the two predicted curves are identical, and only the residual differs\. Two of 378 early presentations lack a valid donor and are excluded, leaving 376 presentations\. For each presentation, predicted rates are averaged across its bins before calculating the squared log\-rate difference between curves\. The neuron score averages this difference over orientations and presentations:

cℓ=1376​∑p=137618​∑o=18\[log⁡\(r¯ℓ​p​osource\+ϵ\)−log⁡\(r¯ℓ​p​odonor\+ϵ\)\]2,ϵ=10−4\.c\_\{\\ell\}=\\frac\{1\}\{376\}\\sum\_\{p=1\}^\{376\}\\frac\{1\}\{8\}\\sum\_\{o=1\}^\{8\}\\left\[\\log\(\\bar\{r\}^\{\\mathrm\{source\}\}\_\{\\ell po\}\+\\epsilon\)\-\\log\(\\bar\{r\}^\{\\mathrm\{donor\}\}\_\{\\ell po\}\+\\epsilon\)\\right\]^\{2\},\\qquad\\epsilon=10^\{\-4\}\.\(28\)
Independent TF contribution is evaluated on 222 final\-epoch presentations using presentation\-grouped two\-fold cross\-validation, repeated five times\. The reduced Poisson model contains categorical orientation and linear running speed; the full model additionally contains categorical TF, and the null model contains only an intercept\. Each neuron’s contribution is\(Dreduced−Dfull\)/Dnull\(D\_\{\\mathrm\{reduced\}\}\-D\_\{\\mathrm\{full\}\}\)/D\_\{\\mathrm\{null\}\}\[[59](https://arxiv.org/html/2610.33145#bib.bib66)\], where the deviances are evaluated on held\-out presentations\. Running\-speed scaling and the null rate are fitted only within each training half and then applied to its held\-out counterpart\.

We compare Neural\-DISCO with an out\-of\-fold, identity\-link Gaussian GLM with categorical orientation and linear running speed\. Its edit adds the GLM’s predicted change to the source rates:redit=rsource−f⁡\(ysource\)\+f⁡\(ytarget\)r\_\{\\mathrm\{edit\}\}=r\_\{\\mathrm\{source\}\}\-f\(y\_\{\\mathrm\{source\}\}\)\+f\(y\_\{\\mathrm\{target\}\}\)\. Using identical source–donor pairs, its score is the squared difference between presentation\-averaged raw\-rate curves, averaged over orientations and presentations\. The shared target prediction cancels in this difference, leaving the paired additive residual difference\. We also evaluate squared log\-rate differences after flooring edited rates atϵ\\epsilonbefore presentation averaging\.

Table 10:Hypothesis prioritization across all 672 units and the anatomical subset of 311 visual\-cortex units\. Spearman correlation between the early\-epoch score and final\-epoch TF contribution\. Values are correlation±\\pmbootstrap standard error from 2,000 presentation resamples, stratified by fold and orientation in the early epochs and by orientation and TF in the final epoch\.The visual\-cortex comparison in Table[10](https://arxiv.org/html/2610.33145#A1.T10)supports using source\-dependent predicted tuning to prioritize hypotheses about an omitted experimental variable\. It does not identify a unique hidden cause or establish a causal effect of residual state\.

### A\.15Compute resources

All experiments were run on a single H200 GPU with 141 GB RAM\. Training and evaluation of Neural\-DISCO and all baselines across all datasets required approximately 6 GPU hours in total\. RatInABox sweeps overλ\\lambdarequired approximately 2 GPU hours, MC Mazeλ\\lambdasweeps required approximately 1 GPU hour, Allenλ\\lambdasweeps required approximately 1 GPU hour, and baseline training and evaluation required approximately 2 GPU hours\. During Neural\-DISCO training, peak VRAM usage stayed below 1 GB\.

### A\.16RatInABox synthetic dataset experiment details

##### Dataset generation\.

We use the RatInABox generator\[[8](https://arxiv.org/html/2610.33145#bib.bib39)\]to generate synthetic neural spiking activity, simulator firing rates, and associated behavioral variables\. The synthetic environment has aspect ratio 2:1 and scale 1 m\. Dataset generation uses random seed 42\. The simulation timestep is 50 ms, the duration is 20 min, the initial position is\(0\.5,0\.5\)\(0\.5,0\.5\), and the mean speed is 0\.2 m/s\. We generate four trajectories:

- •Train:standard random\-walk trajectory\.
- •Inverted:the same trajectory astrain, but temporally reversed\.
- •Shifted:a trajectory spatially confined to the left side of the environment,x≤1x\\leq 1\.
- •Cfact\_shifted:the same asshifted, but with a\+1\+1offset along thexx\-axis\.

##### Synthetic cell types\.

The synthetic population contains 300 neurons: 100 place cells, 100 head\-direction cells, and 100 speed cells\.

##### Loaded arrays and observed variables\.

Each generated RatInABox file contains 24,000 bins\. Spike counts and simulator rates have shape\(24000,300\)\(24000,300\), position has shape\(24000,2\)\(24000,2\), head direction has shape\(24000,2\)\(24000,2\), and velocity has shape\(24000,2\)\(24000,2\)\. Neural\-DISCO receives position and head direction as observed editable factors\. Speed is withheld from the observed\-factor embeddings, so speed\-related information can only be preserved through the residual latent\.

##### Splits and counterfactual edits\.

RatInABox uses contiguous sample splits rather than trial\-grouped splits\. Each fold contains 17,280 training samples and 1,920 validation samples\. The held\-out test loader uses the first 20% of theinvertedtrajectory, corresponding to 4,800 samples\. For counterfactual generation, the model keeps the source residual and swaps the configured target behavior blocks\. In the position\-shift edit, the target position comes from the shifted counterfactual trajectory while head direction and speed are held fixed\. In the head\-direction edit, head direction is reversed while position and speed are held fixed\. Since the simulator provides ground\-truth rates under the edited variables, generated rates are compared directly against the edit’s simulator target rates by cell group: place cells, head\-direction cells, speed cells, and all cells\.

#### A\.16\.1Generalization to unseen factor combinations

During training, we withhold all samples whose position satisfiesx∈\[1\.5,2\.0\]x\\in\[1\.5,2\.0\],y∈\[0\.5,1\.0\]y\\in\[0\.5,1\.0\], and whose heading has a non\-negativeyy\-component \(hy≥0h\_\{y\}\\geq 0\)\. Herehyh\_\{y\}is theyy\-component of the unit head\-direction vector\. The joint position–heading combinations are therefore absent from training\. Table[11](https://arxiv.org/html/2610.33145#A1.T11)reports position\-edit fidelity and preservation of the withheld speed factor: pi\-VAE achieves slightly higher position fidelity, whereas Neural\-DISCO preserves speed\-cell activity\.

Table 11:Position editing with joint position–heading combinations withheld from training\. Position\-edit and speed\-preservationR2R^\{2\}average neuron\-wise scores within the place\-cell and speed\-cell populations, respectively, against simulator target rates\. Values are mean±\\pmSEM across ten cross\-validation folds\.

### A\.17MC Maze Experiment Details

##### Dataset\.

We use the MC Maze reaching dataset from the Neural Latents Benchmark / DANDI dataset000128/0\.220113\.0400, subjectsub\-Jenkins\. Spikes are loaded fromtrain\_spikes\_heldin; no additional unit filtering is applied beyond the NLB held\-in unit selection\. The dataset contains 137 recorded units together with paired behavioral trajectories\. Spike counts are binned at 50 ms, yielding 14 bins per trial in the benchmark representation\. We use the NLB train split for 10\-fold cross\-validation and reserve the public NLB validation split as the held\-out test set for decoding analyses\.

##### Observed factors and preprocessing\.

The observed editable factors are hand position, hand velocity, and eye position, each represented as a two\-dimensional trajectory\. We do not apply additional spike smoothing beyond the benchmark binning\. Behavioral targets are normalized per coordinate before training and decoding\.

##### Splits\.

The NLB train split contains 1,721 trials and the public validation split contains 574 trials\. After flattening 14 time bins per trial, these correspond to 24,094 train\-split samples and 8,036 held\-out samples, respectively\. Cross\-validation is grouped by full trials rather than individual bins so that all time bins from a trial remain in the same fold\. In fold 0, the model split contains 21,672 training samples and 2,422 validation samples\.

##### Training details\.

The MC Maze Neural\-DISCO model is trained for 100 epochs with train/validation/test batch sizes of 512/256/256, learning rate10−310^\{\-3\}, weight decay10−510^\{\-5\}, and scheduler decay factorγ=1\\gamma=1\. The behavior latent dimension isdy=96d\_\{y\}=96for each observed factor and the residual latent dimension isdres=128d\_\{\\mathrm\{res\}\}=128\. The factor encoder, residual encoder, and generator have depths of 3/4/4, with residual penaltyλ=0\.48\\lambda=0\.48and residual noise standard deviationσ=0\.04\\sigma=0\.04\. The unregularized control uses the same architecture and optimizer settings withλ=0\\lambda=0andσ=0\\sigma=0\.

##### Source\-target pairing and counterfactual construction\.

For each held\-out source trial, we choose a target trial with matched task context and a different reach trajectory\. Pairing uses thecontext\_matchedstrategy: candidate targets must match the source trial ontrial\_version,num\_targets,num\_barriers, andactive\_target, and must differ in target direction by at least30∘30^\{\\circ\}\. Among the eligible candidates, the target trial is chosen by nearest standardized delay and reaction time\. Pairing uses seed 0\. On the public validation split, this procedure retains 557 source\-target pairs from 574 trials\. In the main counterfactual experiment, the hand\-position, hand\-velocity, and eye\-position blocks are jointly replaced with the target\-trial values while the source residual is held fixed\.

##### Decoder training and attribution controls\.

Downstream MC Maze decoders are trained only on real training\-split spikes and behavior\. We use a KNN decoder on CEBRA embedding with theoffset10\-model, batch size 1000, learning rate10−410^\{\-4\}, temperature 1\.0, output dimension 8, 15,000 training iterations, cosine distance,time\_deltaconditioning, time offset 10, KNN withK=36K=36, cosine metric, and seed 0\. For paper figures, the same trained decoder is applied to source spikes, target spikes, and generated activity\. Generated activity is obtained by Poisson sampling from the counterfactual Neural\-DISCO rates before downstream decoding\. For attribution sanity checks, recorded neurons are augmented with synthetic constant\-rate control neurons, and we report AUC for distinguishing recorded units from these controls using the total attribution score to all behaviors for each neuron\.

### A\.18Allen Visual Coding Neuropixels experiment details

##### Spike binning and preprocessing\.

For the 32\-session speed\- and orientation\-editing comparisons, we apply the preprocessing and splitting protocol below within each session\. Examples are drawn from the 32\-session cohort\. We use only spiking and behavior during drifting\-grating stimulus presentations, filtering out gray periods in between drifting stimuli\. For speed editing, running speed and stimulus orientation are used as observed factors, while drifting\-grating temporal frequency is withheld from the observed\-factor embeddings for preservation analyses\. We remove presentations with blank or invalid orientation labels\. In session 756029989, the filtered dataset contains 378 training/CV presentations and 222 held\-out test presentations\. Unit filtering uses strict inequalitiesSNR\>1\\mathrm\{SNR\}\>1andisi\_violations<0\.5<0\.5, leaving 672 units in session 756029989\. Spike trains and running speed are binned at 50 ms resolution over each 2 s drifting\-grating presentation, yielding 40 bins per presentation\.

##### Data splitting\.

Train/test splitting is by drifting\-grating epoch: the first two drifting\-grating epochs are used for model training and cross\-validation, and the final epoch is reserved as held\-out test data\. Cross\-validation splits are grouped bypresentation\_idusing deterministic contiguousnp\.array\_splitpartitions of sorted presentation ids\. In fold 0 of session 756029989, the training split contains 13,600 bins from 340 presentations, and the validation split contains 1,520 bins from 38 presentations\. The held\-out test split contains 8,880 bins from 222 presentations\.

##### Decoder training\.

For speed\-edit evaluation, we train an MLP decoder only on real fold\-training activity\. Real train/test spikes are Gaussian\-smoothed withσ=2\.0\\sigma=2\.0bins before decoder fitting and scoring\. The speed decoder has hidden dimension 256, 3 layers, GELU nonlinearities, dropout 0\.0, learning rate10−310^\{\-3\}, weight decay10−410^\{\-4\}, maximum training duration 300 epochs, patience 20, and seed161803\+fold index161803\+\\text\{fold index\}\. Its internal validation split uses seed42\+fold index42\+\\text\{fold index\}\. For temporal\-frequency preservation, we train a multinomial logistic\-regression decoder only on real training activity and evaluate it on edited activity

LogisticRegression\(max\_iter=2000,\\displaystyle\\texttt\{LogisticRegression\(max\\\_iter=2000,\}solver="lbfgs", class\_weight="balanced"\)\\displaystyle\\texttt\{solver="lbfgs", class\\\_weight="balanced"\)\}

##### Unit metrics for attribution experiments\.

For the experiment shown in Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)f, ground\-truth labels for running modulation are derived from the Allen unit metricrun\_pval\_dg, with units classified as significantly modulated whenrun\_pval\_dg<0\.05<0\.05\. In Fig\.[4](https://arxiv.org/html/2610.33145#S4.F4)l, we compare orientation\-attribution scores with the Allen\-provided global orientation selectivity index,g\_osi\_dg\. Stimulus examples report the most frequent decoded label within each presentation\.

##### Training time\.

Training the Allen 10\-fold models across the reportedλ\\lambdavalues takes approximately 1 GPU hour on an H200 GPU\.

Table 12:Hyperparameters and training settings for the main editing experiments\. For Neural\-DISCO, architecture entries give factor encoder / residual encoder / generator depths\. Depths count input, hidden and output layers, so depthddhasd−1d\-1weight matrices\. Batch sizes are train/validation/test\. The attribution experiments sweepλ\\lambda\. The unregularized control uses the same architecture and training settings as Neural\-DISCO within each dataset and task, withλ=0\\lambda=0andσ=0\\sigma=0\.Architecture and residual bottleneck

Training settings

Neural\-DISCO models were trained with Adam and checkpoints selected by minimum validation loss\. pi\-VAE uses validation stopping with patience 40 and reloads the minimum\-loss checkpoint\. Allen NEDS retains checkpoints when validation loss improves by at least10−410^\{\-4\}, with early\-stopping patience 50, and uses training\-fold input min–max normalization\. The Allen stimulus\-edit models use the same latent widths and optimizer settings, with factor/residual/generator depths 3/4/4 andλ=\.01\\lambda=\.01; their ablation setsλ=σ=0\\lambda=\\sigma=0\. The same pi\-VAE and NEDS settings apply to speed and stimulus edits\.

Table 13:Trainable parameter counts for the main editing models\. Allen counts are shown for session 756029989 \(672 units\)\. The unregularized control \(w/o reg\.\) has the same architecture as Neural\-DISCO within each dataset and task and therefore has the same number of trainable parameters\.

### A\.19Baselines

For baseline models, we used the authors’ released implementations where available, adapting input and covariate dimensions and sequence length to our preprocessing\. Hyperparameters and training settings for pi\-VAE and NEDS are given in Table[12](https://arxiv.org/html/2610.33145#A1.T12)\. Training and evaluation of the baselines across all experiments takes approximately 2 GPU hours on an H200 GPU\.

## Appendix BProof of the residual capacity bound

###### Proof\.

LetS∼Unif​\{1,…,N\}S\\sim\\mathrm\{Unif\}\\\{1,\\ldots,N\\\}, letXS=x\(S\)X\_\{S\}=x^\{\(S\)\}, and letUS=ϕres​\(XS\)U\_\{S\}=\\phi\_\{\\mathrm\{res\}\}\(X\_\{S\}\)\. ForU~=US\+ϵ\\tilde\{U\}=U\_\{S\}\+\\epsilonwith independentϵ∼𝒩⁡\(0,σ2​Im\)\\epsilon\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}I\_\{m\}\),

I⁡\(S,U~\)=h⁡\(U~\)−h⁡\(U~∣S\)=h⁡\(U~\)−h⁡\(ϵ\)\.I\(S;\\tilde\{U\}\)=h\(\\tilde\{U\}\)\-h\(\\tilde\{U\}\\mid S\)=h\(\\tilde\{U\}\)\-h\(\\epsilon\)\.\(29\)LetΣU\\Sigma\_\{U\}be the covariance ofUSU\_\{S\}\. Thentr⁡\(ΣU\)≤𝔼​‖US‖22=P^\\mathrm\{tr\}\(\\Sigma\_\{U\}\)\\leq\\mathbb\{E\}\\\|U\_\{S\}\\\|\_\{2\}^\{2\}=\\hat\{P\}\. Since Gaussian random variables maximize differential entropy under a covariance constraint,

h⁡\(U~\)≤12​log⁡\(\(2​π​e\)m​det\(ΣU\+σ2​Im\)\)\.h\(\\tilde\{U\}\)\\leq\\frac\{1\}\{2\}\\log\\left\(\(2\\pi e\)^\{m\}\\det\(\\Sigma\_\{U\}\+\\sigma^\{2\}I\_\{m\}\)\\right\)\.\(30\)By the arithmetic–geometric mean inequality,

det\(ΣU\+σ2​Im\)≤\(tr⁡\(ΣU\)\+m​σ2m\)m≤\(σ2\+P^m\)m\.\\det\(\\Sigma\_\{U\}\+\\sigma^\{2\}I\_\{m\}\)\\leq\\left\(\\frac\{\\mathrm\{tr\}\(\\Sigma\_\{U\}\)\+m\\sigma^\{2\}\}\{m\}\\right\)^\{m\}\\leq\\left\(\\sigma^\{2\}\+\\frac\{\\hat\{P\}\}\{m\}\\right\)^\{m\}\.\(31\)Alsoh⁡\(ϵ\)=m2​log⁡\(2​π​e​σ2\)h\(\\epsilon\)=\\frac\{m\}\{2\}\\log\(2\\pi e\\sigma^\{2\}\)\. Subtracting gives

I⁡\(S,U~\)≤m2​log⁡\(1\+P^m​σ2\)\.I\(S;\\tilde\{U\}\)\\leq\\frac\{m\}\{2\}\\log\\left\(1\+\\frac\{\\hat\{P\}\}\{m\\sigma^\{2\}\}\\right\)\.\(32\)BecauseYa,S=ya\(S\)Y\_\{a,S\}=y\_\{a\}^\{\(S\)\}is a deterministic function ofSS, data processing givesI⁡\(Ya,S,U~\)≤I⁡\(S,U~\)I\(Y\_\{a,S\};\\tilde\{U\}\)\\leq I\(S;\\tilde\{U\}\)\. BecauseZa,S=ϕa​\(Ya,S\)Z\_\{a,S\}=\\phi\_\{a\}\(Y\_\{a,S\}\)is a deterministic function ofYa,SY\_\{a,S\}, another application of data processing givesI⁡\(Za,S,U~\)≤I⁡\(Ya,S,U~\)I\(Z\_\{a,S\};\\tilde\{U\}\)\\leq I\(Y\_\{a,S\};\\tilde\{U\}\)\. All logarithms are natural, so the mutual informations are measured in nats\. ∎

## Appendix CSupplemental Figures

![Refer to caption](https://arxiv.org/html/2610.33145v1/rat_unknown_latent_grid.png)Fig\. 5:Withheld\-label diagnostics visualization\.Across the three held\-out\-label settings, the residual latent is organized by the omitted variable and not by the observed variables\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/rat_unknown_latent_confusion_matrix.png)Fig\. 6:Withheld\-label diagnostics decoding for Rat in a Box dataset\.Residual\-latent decoders recover the held\-out variable much more strongly than the observed variables\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/allen_residual_decoding_heatmap_raw_2.png)Fig\. 7:Residual decoding with withheld factors on Allen\.Each row withholds one factor from the model’s supplied variables; columns show how well each factor can be decoded from the residual latent\. Chance is 0\.125 for orientation, and 0\.2 for temporal frequency\. Each factor is more decodable when it is withheld, as reflected by the diagonal entries\. The residual preferentially retains omitted\-factor information, although some supplied\-factor information remains \(Table[6](https://arxiv.org/html/2610.33145#A1.T6)\)\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/mc_maze_residual_decoding_heatmap_2.png)Fig\. 8:Residual leakage on MC Maze\.Hand position, hand velocity, and eye position are strongly correlated, as such no factor is held out; all three are supplied to the model\. Leakage is therefore measured by decoding each supplied factor from the residual latent\. All three factors have low held\-out decodingR2R^\{2\}\.![Refer to caption](https://arxiv.org/html/2610.33145v1/rib_heatmap.png)Fig\. 9:RatInABox Attribution Heatmap\.Rows correspond to neurons ordered by synthetic cell type: place cells, head\-direction cells, and speed cells\. Columns show normalized attribution to head\-direction, position, and residual blocks\. \(a\) Attribution results of Neural\-DISCO\. \(b\) Attribution results of architecture\-matched unregularized residual control\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplement_rat_rates.png)Fig\. 10:RatInABox Counterfactual Rates\.Example rates for one cell of each type in \(a\) the position\-shift edit and \(b\) the head\-direction edit\. Generated rates are compared with simulator target rates for each population\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/rat_position_shift_ablation_trajectory.png)Fig\. 11:RatInABox position\-shift trajectory ablation\.Decoded trajectories for the position\-shift counterfactual are shown for the source activity, target trajectory, the regularized Neural\-DISCO counterfactual, and the architecture\-matched unregularized control\. The regularized model tracks the shifted target manifold substantially better, whereas the unregularized control exhibits diffuse trajectories and poorer preservation of the intended edit\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/mc_maze_no_reg_decode_r2.png)Fig\. 12:MC Maze decoding for the unregularized control\.Downstream decodingR2R^\{2\}for generated activity from the architecture\-matched unregularized control, evaluated against both the source and target trajectories\. For all three observed factors, the unregularized control remains closer to the source than to the target, indicating that it fails to perform the intended counterfactual edit\.![Refer to caption](https://arxiv.org/html/2610.33145v1/supplemental_figs/mc_maze_no_reg_reach_trajectories.png)Fig\. 13:MC Maze qualitative reach examples for the unregularized control\.Example hand\-position trajectories for four source\-target trial pairs\. The generated counterfactual trajectories remain between the source and target reaches rather than matching the target movement, consistent with the poor target\-decoding performance of the unregularized control\.![Refer to caption](https://arxiv.org/html/2610.33145v1/allen_heatmap.png)Fig\. 14:Allen Attribution Heatmap sorted by Orientation Selectivity Index\.X\-axis corresponds to individual units, with x\-ticks indicating gOSI cutoffs, colored by orientation attribution value\. Attribution scores are averaged over ten folds before min–max normalization within each method\. \(a\) Attribution results of Neural\-DISCO\. \(b\) Attribution results of architecture\-matched unregularized residual control\.
## Appendix DDataset and Software Licenses

We use the following datasets and software tools in this work\. We comply with the terms of use and licensing conditions of each resource\.

RatInABox SimulatorRatInABox\[[8](https://arxiv.org/html/2610.33145#bib.bib39)\]is an open\-source Python toolkit for modelling rodent locomotion and generating synthetic neuronal data\. The software is available at[https://github\.com/RatInABox\-Lab/RatInABox](https://github.com/RatInABox-Lab/RatInABox)and is released under theMIT License\.NLB\-MC\_Maze Dataset

The MC\_Maze dataset is part of the Neural Latents Benchmark ’21\[[43](https://arxiv.org/html/2610.33145#bib.bib17)\]and is publicly available via the DANDI Archive \([https://dandiarchive\.org/dandiset/000128](https://dandiarchive.org/dandiset/000128)\)\. The dataset was originally collected by Kaufman, Churchland, and Shenoy at Stanford University and is distributed under theCreative Commons Attribution 4\.0 International \(CC BY 4\.0\)license\.Allen Brain Observatory – Visual Coding Neuropixels

NEDS\. We used the official implementation of NEDS\[[65](https://arxiv.org/html/2610.33145#bib.bib38)\], available at[https://github\.com/yzhang511/NEDS](https://github.com/yzhang511/NEDS), which is released under the MIT License\.

## NeurIPS Paper Checklist

1. 1\.Claims
2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
3. Answer:\[Yes\]\.
4. Justification: The abstract and introduction frame the paper around source\-conditioned counterfactual editing and explicitly state that the counterfactuals are model\-based predictions rather than identified causal effects\.
5. 2\.Limitations
6. Question: Does the paper discuss the limitations of the work performed by the authors?
7. Answer:\[Yes\]\.
8. Justification: Section[5](https://arxiv.org/html/2610.33145#S5)discusses causal interpretation, residual\-bottleneck sensitivity, and indirect validation in real neural data\.
9. 3\.Theory assumptions and proofs
10. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete and correct proof?
11. Answer:\[Yes\]\.
12. Justification: Proposition 1 states the finite\-sample residual capacity bound and Appendix[B](https://arxiv.org/html/2610.33145#A2)gives the proof and assumptions\.
13. 4\.Experimental result reproducibility
14. Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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?
15. Answer:\[Yes\]\.
16. Justification: A complete description of model structure and hyperparameters is provided in Appendix[A](https://arxiv.org/html/2610.33145#A1), along with preprocessing instructions, random seeds, fold definitions, and session/unit filtering details needed for full reproduction\.
17. 5\.Open access to data and code
18. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
19. Answer:\[Yes\]\.
20. Justification: Our experiments use public datasets and simulators that are publicly available and we provide exact instructions for acquiring/setting them up\. As well, the manuscript package includes a[code release](https://github.com/TangLab-UBC/Neural-DISCO-neurips)with end\-to\-end reproduction commands/scripts for exact reproducibility\.
21. 6\.Experimental setting/details
22. Question: Does the paper specify all the training and test details necessary to understand the results?
23. Answer:\[Yes\]\.
24. Justification: The paper includes core training settings and hyperparameters in Appendix[A](https://arxiv.org/html/2610.33145#A1), as well as all configuration parameters and information, including optimizer settings, preprocessing commands, split definitions, session identifiers, and filtering criteria\.
25. 7\.Experiment statistical significance
26. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
27. Answer:\[Yes\]\.
28. Justification: Section[4](https://arxiv.org/html/2610.33145#S4)states that reported±\\pmvalues and error bars denote standard error of the mean across CV models\.
29. 8\.Experiments compute resources
30. Question: For each experiment, does the paper provide sufficient information on the computer resources needed to reproduce the experiments?
31. Answer:\[Yes\]\.
32. Justification: Appendix[A](https://arxiv.org/html/2610.33145#A1)reports the GPU type, memory, approximate wall\-clock GPU hours, and per\-experiment compute breakdown\.
33. 9\.Code of ethics
34. Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics?
35. Answer:\[Yes\]\.
36. Justification: The work uses public animal neural datasets and a simulator, and the manuscript does not include human\-subject data collection or high\-risk deployment\.
37. 10\.Broader impacts
38. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
39. Answer:\[Yes\]\.
40. Justification: Section[5](https://arxiv.org/html/2610.33145#S5)discusses the expected broader impact of the work as a research tool forin\-silicohypothesis generation in computational neuroscience\. The work uses public animal neural recordings and simulations, introduces no new human\-subject dataset, and is not intended for clinical, diagnostic, or closed\-loop interventional deployment\. Its broader societal risk is therefore limited, while its primary impact is to help prioritize neural–behavioral hypotheses and guide follow\-up analyses\.
41. 11\.Safeguards
42. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse?
43. Answer:\[N/A\]\.
44. Justification: The work does not release a high\-risk generative model such as a language, image, or scraped\-data model; it targets animal neural activity modeling\.
45. 12\.Licenses for existing assets
46. Question: Are the creators or original owners of assets used in the paper properly credited and are the license and terms of use explicitly mentioned and properly respected?
47. Answer:\[Yes\]\.
48. Justification: The paper cites the public datasets and baseline methods\. Appendix[D](https://arxiv.org/html/2610.33145#A4)lists asset sources and licenses; the code package provides pinned software dependencies and environment installation instructions\.
49. 13\.New assets
50. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
51. Answer:\[Yes\]\.
52. Justification: The paper does not introduce a new dataset\. The code package includes reproduction instructions, experiment configurations, and pinned dependency locks\.
53. 14\.Crowdsourcing and research with human subjects
54. 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 as details about compensation?
55. Answer:\[N/A\]\.
56. Justification: The reported experiments do not involve crowdsourcing or new human\-subject data collection\.
57. 15\.Institutional review board approvals or equivalent for research with human subjects
58. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board approvals or an equivalent review were obtained?
59. Answer:\[N/A\]\.
60. Justification: The reported experiments do not involve new human\-subject data collection\.
61. 16\.Declaration of LLM usage
62. Question: 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?
63. Answer:\[N/A\]\.
64. Justification: The core methodology does not use LLMs as a model component\. Use of LLMs only for writing or formatting does not need to be declared under this question\.