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

Nico Policzer; James Ho; Joseph Soo; Xin Tang

cs.LGOct 23, 2026

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

Authors: Nico Policzer, James Ho, Joseph Soo, Xin Tang

Organizations: Michael Smith Laboratories and Department of Computer Science University of British Columbia 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 introduce Neural-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.

Figures & tables

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.

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.

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-decoding $R^{2}$ comparisons; dotted lines show the empirical target-spike decoding reference.

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 $\pm$ SEM 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.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Position shift
Method	Edited $R^{2}$	Other observed $R^{2}$	Speed $R^{2}$	$S_{\mathrm{EP}}$
Neural-DISCO	$0.940\pm 0.004$	$0.888\pm 0.001$	$0.964\pm 0.001$	$0.933\pm 0.002$
Zero residual	$0.976\pm 0.002$	$0.911\pm 0.0002$	$-0.004\pm 0.001$	$0.621\pm 0.0004$
Gaussian GLM edit	$-208.575\pm 0.391$	$-2.116\pm 0.001$	$-1.258\pm 0.002$	$0.000\pm 0.000$
Poisson GLM	$0.036\pm 0.012$	$0.914\pm 0.00002$	$-0.034\pm 0.002$	$0.065\pm 0.019$
Head-direction reversal

Appendix

Table 1: RatInABox editing against exact simulator target trajectories. Population

R^{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

\pm

SEM across ten folds. Each component is clipped before computing the edit–preservation score.

Method	Speed RMSE (cm/s)	TF balanced accuracy	$S_{\mathrm{EP}}$
Neural-DISCO	$3.629\pm 0.567$	$0.541\pm 0.013$	$0.592\pm 0.026$
w/o reg.	$6.182\pm 1.067$	$0.570\pm 0.013$	$0.391\pm 0.030$
Zero residual	$3.114\pm 0.462$	$0.199\pm 0.0004$	$0.012\pm 0.001$
pi-VAE	$4.004\pm 0.684$	$0.197\pm 0.0008$	$0.010\pm 0.001$
NEDS	$3.979\pm 0.645$	$0.194\pm 0.0008$	$0.007\pm 0.001$
Gaussian GLM edit	$6.017\pm 1.044$	$0.597\pm 0.012$	$0.372\pm 0.040$

Appendix

Table 2: Allen speed editing and temporal-frequency preservation across 32 sessions, using a common fold-training mean target. Metrics and

S_{\mathrm{EP}}

are computed per fold and averaged over ten folds within each session; values are mean

\pm

SEM across the 32 session means. Lower speed RMSE and higher TF balanced accuracy and

S_{\mathrm{EP}}

are better.

Velocity rotation	Velocity $R^{2}$	Hand-position $R^{2}$	Eye-position $R^{2}$
$0^{\circ}$	$0.772\pm 0.003$	$0.895\pm 0.002$	$0.630\pm 0.004$
$45^{\circ}$	$0.618\pm 0.002$	$0.845\pm 0.002$	$0.578\pm 0.003$
$90^{\circ}$	$0.344\pm 0.006$	$0.744\pm 0.004$	$0.484\pm 0.005$
$135^{\circ}$	$0.160\pm 0.011$	$0.647\pm 0.005$	$0.410\pm 0.004$
$180^{\circ}$	$0.102\pm 0.009$	$0.603\pm 0.008$	$0.388\pm 0.008$

Appendix

Table 3: MC Maze decoding under velocity rotations. Position and eye targets remain unchanged.

Method	Orientation BA	TF BA	Speed RMSE (cm/s)	$S_{\mathrm{EP}}$
Neural-DISCO	$0.949\pm 0.006$	$0.823\pm 0.012$	$6.268\pm 1.017$	$0.838\pm 0.022$
w/o reg.	$0.842\pm 0.012$	$0.693\pm 0.013$	$5.613\pm 0.933$	$0.807\pm 0.024$
Zero residual	$0.961\pm 0.005$	$0.855\pm 0.013$	$11.943\pm 1.741$	$0.656\pm 0.002$
Gaussian GLM	$0.714\pm 0.016$	$0.532\pm 0.015$	$6.119\pm 1.015$	$0.705\pm 0.023$
Poisson GLM	$0.976\pm 0.003$	$0.914\pm 0.008$	$11.894\pm 1.751$	$0.660\pm 0.001$
pi-VAE	$0.317\pm 0.015$	$0.620\pm 0.018$	$11.743\pm 1.720$	$0.263\pm 0.013$

Appendix

Table 4: Orientation editing and source preservation across 32 Allen sessions. Metrics and

S_{\mathrm{EP}}

are computed per fold and averaged over ten folds within each session; values are mean

\pm

SEM 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 and

S_{\mathrm{EP}}

are better.

Temporal-frequency shift
Method	Edit fidelity $E$	Preservation $P$	$S_{\mathrm{EP}}$
Neural-DISCO	$0.840\pm 0.005$	$0.895\pm 0.010$	$0.867\pm 0.005$
w/o reg.	$0.479\pm 0.018$	$0.956\pm 0.009$	$0.636\pm 0.016$
Gaussian GLM	$0.466\pm 0.010$	$0.919\pm 0.003$	$0.618\pm 0.009$
Poisson GLM	$0.939\pm 0.002$	$0.500\pm 0.000$	$0.653\pm 0.0004$
pi-VAE	$0.647\pm 0.040$	$0.145\pm 0.040$	$0.226\pm 0.047$

Appendix

Table 5: Additional Allen stimulus edits in session 756029989. Values are mean

\pm

SEM across ten folds, with

E

P

and

S_{\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.

Factor withheld	Orientation BA	Speed $R^{2}$	TF BA
Orientation	$\mathbf{.486\pm.007}$	$.201\pm.084$	$.252\pm.016$
Speed	$.142\pm.010$	$\mathbf{.649\pm.034}$	$.265\pm.013$
Temporal frequency	$.163\pm.008$	$.083\pm.104$	$\mathbf{.550\pm.012}$
Chance / reference	$.125$	$0$	$.20$

Appendix

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

\pm

SEM across ten folds. Balanced accuracies are reported without chance correction; speed uses raw

R^{2}

Covariate	Clean residual $R^{2}$	Source spikes $R^{2}$
Hand position	$0.045\pm 0.003$	$0.448\pm 0.004$
Hand velocity	$0.049\pm 0.005$	$0.387\pm 0.005$
Eye position	$0.053\pm 0.003$	$0.364\pm 0.007$

Appendix

Table 7: Held-out decoding of supplied MC Maze covariates from the clean residual, with source-spike decoding as a reference. Values are mean

\pm

SEM across ten folds. Source-spike decoding uses Ridge probes and uniform-average

R^{2}

across the two coordinates. Negative

R^{2}

is retained.

	Position edit		Direction edit
$\alpha$	All-population $R^{2}$	$S_{\mathrm{EP}}$	All-population $R^{2}$	$S_{\mathrm{EP}}$
0	$0.937\pm.002$	$0.941\pm.003$	$0.945\pm.001$	$0.931\pm.001$
0.25	$0.930\pm.002$	$0.934\pm.003$	$0.939\pm.001$	$0.927\pm.001$
0.50	$0.931\pm.001$	$0.937\pm.001$	$0.940\pm.001$	$0.933\pm.001$
0.75	$0.930\pm.001$	$0.936\pm.002$	$0.941\pm.001$	$0.936\pm.001$
1	$0.929\pm.001$	$0.933\pm.001$	$0.938\pm.001$	$0.937\pm.001$

Appendix

Table 8: Neural-DISCO counterfactual performance under increasing multiplicative mixed selectivity. Values are mean

\pm

SEM across ten independently trained cross-validation folds. The

R^{2}

columns average neuron-wise scores across the full population;

S_{\mathrm{EP}}

combines primary-population edit fidelity with performance on the other two populations.

Dataset	Training subset	Condition coverage	Limiting outcome	Retention
RatInABox	3.6 min; 4,320 bins (25%)	Contiguous interval	Position-edit $R^{2}$	$95.1\pm 1.9\%$
MC Maze	145 trials (9.375%)	108/108; median 1 trial	Velocity $R^{2}$	$93.9\pm.3\%$
Allen	170 presentations (50%)	40/40; median 4 presentations	TF balanced accuracy	$96.3\pm 2.4\%$

Appendix

Table 9: Tested training-subset requirements for retaining at least

90\%

of full-data performance on every outcome. Retention is mean

\pm

SEM across ten folds for the limiting outcome; the full-fold reference equals

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.

Editor	All units	Visual cortex
Neural-DISCO	$0.386\pm 0.030$	$0.448\pm 0.030$
Gaussian GLM (raw-rate score)	$0.192\pm 0.022$	$0.186\pm 0.027$
Gaussian GLM (log-rate score)	$-0.061\pm 0.020$	$0.121\pm 0.024$

Appendix

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

\pm

bootstrap 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.

Method	Position-edit $R^{2}$	Speed-preservation $R^{2}$
Neural-DISCO	$0.940\pm 0.004$	$0.964\pm 0.001$
pi-VAE	$0.966\pm 0.001$	$-0.041\pm 0.003$

Appendix

Table 11: Position editing with joint position–heading combinations withheld from training. Position-edit and speed-preservation

R^{2}

average neuron-wise scores within the place-cell and speed-cell populations, respectively, against simulator target rates. Values are mean

\pm

SEM across ten cross-validation folds.

Dataset	Model	Labels	Latents	Architecture	$\lambda$	$\sigma$
RatInABox	Neural-DISCO	pos, head dir	$d_{y}=16$ , $d_{\mathrm{res}}=32$	3/3/3	0.10	0.02
RatInABox	unregularized	pos, head dir	$d_{y}=16$ , $d_{\mathrm{res}}=32$	3/3/3	0	0
MC Maze	Neural-DISCO	hand pos, hand vel, eye pos	$d_{y}=96$ , $d_{\mathrm{res}}=128$	3/4/4	0.48	0.04
MC Maze	unregularized	hand pos, hand vel, eye pos	$d_{y}=96$ , $d_{\mathrm{res}}=128$	3/4/4	0	0
Allen speed	Neural-DISCO	speed, orientation	$d_{y}=48$ , $d_{\mathrm{res}}=96$	3/4/2	0.0015	0.02
Allen stimulus	Neural-DISCO	orientation, TF	$d_{y}=48$ , $d_{\mathrm{res}}=96$	3/4/4	0.01	0.02

Appendix

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 depth

d

has

d-1

weight 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

\lambda=0

and

\sigma=0

Dataset	Model	Trainable params
RatInABox	Neural-DISCO	122,414
RatInABox	w/o reg.	122,414
RatInABox	pi-VAE	232,183
MC Maze	Neural-DISCO	308,772
MC Maze	w/o reg.	308,772
MC Maze	pi-VAE	64,435

Appendix

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.

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.

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.

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 ).

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 decoding $R^{2}$ .

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.

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.

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.

Fig. 12: MC Maze decoding for the unregularized control. Downstream decoding $R^{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.

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.

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.

Neural Encoding

Neural encoding studies how information is represented in biological neural activity and how models map stimuli or signals into neural codes.

See all 92 papers

Oct 27, 2026

Neuronal Identity as an Organizational Basis for Analyzing Neural Population Dynamics

Individual neurons exhibit diverse molecular, morphological, and electrophysiological properties that shape their distinct roles in neural circuits. While these multimodal features are largely shared across individuals, they remain an underutilized source of information in analyzing neural population dynamics. We posit that these shared biological profiles provide a common reference for relating neural dynamics across subjects and recording sessions. Our framework leverages this structure by routing neurons into latent groups based on their biological profiles and analyzing population dynamics through group-level representations. Because these assignments are anchored in biological properties that are consistent across subjects, the learned organization supports transfer to unseen animals without retraining or neuron-wise correspondence. We evaluate our framework on whole-brain calcium imaging in larval zebrafish and large-scale electrophysiology in the mouse visual cortex. Across both species and recording modalities, the framework delivers state-of-the-art performance in both neural decoding and manifold analysis: it achieves superior decoding of behavior and stimulus labels while better preserving the temporal consistency and local geometric structure of latent trajectories on unseen data. Our experiments demonstrate that not only does providing biological features improve baseline performance, but our proposed framework achieves even stronger results by using these properties to explicitly structure the population rather than treating them as auxiliary inputs. Together, these results suggest that organizing neural populations by their multimodal biological features offers a promising framework for generalizing population dynamics across subjects and recording sessions.

Seungjae Han, Joshua Yedam You, Soi Kim +5

School of Electrical Engineering, KAIST, Daejeon, Republic of Korea · Department of Brain and Cognitive Sciences, KAIST, Daejeon, Republic of Korea · Graduate School of Data Science, KAIST, Daejeon, Republic of Korea +3

Neural Decoding Shared Representation Learning Neural Representation Geometry Dynamical Systems

Ask PDF

Oct 2, 2026

NeuroLens: Learning Latent Embeddings of Neural Semantics from Chronic Recordings

Understanding how neural activity represents higher-order cognition and how these representations evolve over time has long been a central pursuit in neuroscience. However, current analytical tools cannot easily distinguish representational plasticity from recording instability in chronic neural recordings. Here, we introduce NeuroLens (Latent Embeddings of Neural Semantics), a self-supervised model based on the Joint-Embedding Predictive Architecture (JEPA) framework that learns denoised, semantically informative latents from chronic neural recordings. An adaptive encoder maps changing neural populations into a common latent space, while a temporal predictor learns structure that supports prediction of future latent states. By predicting in latent space, NeuroLens captures temporally predictable structure and reduces sensitivity to transient, recording-specific variability. Across chronic intracortical data in mice and humans, the learned representations improve decoding of decision-making and semantic task variables. Multi-day pretraining enables generalization to future sessions, rapid few-shot adaptation to unseen neural populations, and more stable decoding over time than state-of-the-art baselines. Together, these results establish NeuroLens as a new paradigm for studying how neural representations change during learning and over long timescales.

Hanrui Lyu, Baiyuan Chen, Tianshu Tan +8

Northwestern University · University of Cambridge · Johns Hopkins University +3

Self-Supervised Time Series Representation Learning Neural Encoding Neural Decoding

Ask PDF

Aug 5, 2026

NeuroPB: Scaling Neural Decoding with Pretrained Behavioral Representations

Decoding continuous motor trajectories from neural activity is essential for developing practical brain-computer interfaces (BCIs). However, current neural decoders are constrained by the limited scale and heterogeneity of neural recordings. In contrast, behavioral data can be collected more readily and at substantially larger scale from humans, animals, simulations, and robotic systems. Here, we introduce NeuroPB, a framework that scales neural decoding by transferring knowledge from pretrained behavioral representations. NeuroPB first pretrains a motor encoder on large-scale motor behavior data and then aligns neural activity with the resulting behavioral representation space using a limited set of paired neural-behavioral recordings. A neural encoder and lightweight motor decoder are subsequently optimized to reconstruct continuous movement from the aligned neural representations. Across multiple macaque motor datasets, behavioral pretraining improves trajectory decoding, including an 11%

R^2

increase on center-out and 8% on random-target compared with training the motor encoder from scratch. Notably, pretraining on robotic trajectories achieves performance comparable to pretraining on macaque trajectories, demonstrating that transferable kinematic structure is shared across biological and artificial models. Moreover, decoding performance improves as the scale and diversity of robotic pretraining data increase, when the amount of neural data is fixed. Pretraining also enhances generalization across recording sessions, subjects, and motor tasks, with only 10% calibration needed to match training from scratch. Overall, these results establish behavioral pretraining as a scalable source for neural decoding and provide a promising route toward high-performance and calibration-efficient BCIs under limited neural data.

Luyao Jin, Yonghao Song, Huan Zhao +2

Department of Mechanical and Automation Engineering, The Chinese University of Hong Kong · Department of Biomedical Engineering, School of Medicine, Tsinghua University · School of Biomedical Sciences, The Chinese University of Hong Kong

Representation Learning Brain-Computer Interfaces Neural Decoding Cross-Domain Transfer Learning

Ask PDF

Abstract

Neural Encoding

Explore similar work

Neuronal Identity as an Organizational Basis for Analyzing Neural Population Dynamics

NeuroLens: Learning Latent Embeddings of Neural Semantics from Chronic Recordings

NeuroPB: Scaling Neural Decoding with Pretrained Behavioral Representations