Drop, Swap, and Generate: A Self-Supervised Approach for Generating Neural Activity

1 Introduction

• Without labels or additional inputs, distinguishing different sources of variability in neural activity remains a major challenge.Unsupervised representations can reveal neural response stability, individual differences, and remapping through learning.Source: §1 · p. 1
• The method aligns content representations of augmented views, created by dropping neurons and jittering samples in time, while reconstructing the original inputs with content and style variables.These transformations are intended to encourage temporal consistency and invariance to the neurons used to represent a brain state.Source: §1 · p. 2
• On synthetic data and non-human-primate reaching datasets from two individuals, Swap-VAE disentangled behavior and movement dynamics without labels.The authors propose general-purpose measures of representation quality to assess this separation.Source: §1 · p. 2
• Swap-VAE is a generative method designed to learn interpretable latent factors of neural activity and generate realistic activity.Its content and style framing draws on image decomposition and maps movement target to content and movement dynamics to style.Source: §1 · p. 2
• BlockSwap swaps content variables between views and predicts the original view from content taken from the other view.The paper also proposes metrics for disentanglement and applies them to neural datasets from different non-human primates.Source: §1 · p. 2

2.1 Variational autoencoders and their application in neural data analysis

• VAEs optimize a tractable evidence lower bound; β > 1 strengthens prior-alignment regularization, which empirically induces more disentanglement and interpretability.The standard VAE uses β = 1, while β > 1 defines the beta-VAE.Source: §2.1 · pp. 2–3
• Figure 1 depicts augmented brain-state views being encoded into content and style, aligned in content, and reconstructed through a decoder using both spaces.It illustrates movement content as reach direction and style as movement dynamics.Source: Figure 1 · p. 3
• Neural-data VAEs include sequential models for population dynamics and single-trial reconstruction, while pi-VAE uses an MLP encoder on individual neural-state vectors.This work likewise uses an MLP encoder and does not explicitly incorporate temporal structure into its architecture.Source: §2.1 · p. 3

2.2 Instance-specific alignment and self-supervision

• Self-supervised alignment trains representations to be similar across augmented views of the same sample, rather than using labels to regularize them.Prior work applied alignment with dropout, temporal shifts, and noise to neural recordings; this work couples alignment with a generative model without a dual network.Source: §2.2 · p. 3

3.1 Model for neural datasets

• Each neural state vector records estimated firing rates for d neurons at one of N time bins.Source: §3.1 · p. 4

3.2 Unifying instance-specific alignment and generative modeling

• The encoder splits each view's latent code into content and style, aligning content across views while reconstructing inputs and regularizing style toward an isotropic Gaussian prior.For neural datasets, the reconstruction loss is Poisson.Source: §3.2 · p. 4

3.3 BlockSwap: A novel latent space augmentation for disentanglement

• BlockSwap exchanges content codes between views while retaining each view's style, then adds reconstruction of swapped codes to the original-view reconstruction objective.The alignment term can be removed when using this loss with style-space KL regularization.Source: §3.3 · p. 4

3.4 Representational quality and disentanglement metrics

• The multi-task disentanglement score compares each latent variable's variance under changes in reach direction versus temporal structure, then aggregates normalized scores across variables.A larger score indicates greater specificity to one factor than the other.Source: §3.4 · p. 5
• Linear readouts decode reach direction or temporal structure from frozen full, content, or style representations, with direction quality measured by acc and delta-acc.The linear weights are retrained for each target label while the representation stays fixed.Source: §3.4 · p. 5

4.1 Synthetic experiments

• On the synthetic datasets, Swap-VAE achieved a multi-task disentanglement score of 0.93, versus 0.46 for beta-VAE and 0.12 for the supervised model, and recovered discrete classes and sequential structure.Source: §4.1 · p. 6, Figure 2 · p. 6

4.2 Experiments on neural datasets

• Content separates reach direction, while Style organizes time with reach direction collapsed; across four datasets, Chewie's multi-task covariance scores were lower than Mihi's on average.Content provides good reach decoding but weaker temporal decoding, while the reverse pattern holds for Style.Source: §4.2 · pp. 7–8, Figure 3 · p. 7
• Swap-VAE matches or exceeds benchmarks on reach decoding and is comparable to beta-VAE and Sup-Time on temporal decoding, unlike self-supervised models that miss temporal structure.Temporal jitter builds invariance that can make temporal structure difficult to decode; Swap-VAE's Style space carries additional information for both tasks.Source: §4.2 · p. 8
• Swap-VAE reconstructs firing rates with lower RMSE than VAE, and generated samples improved classifier accuracy by roughly 5% on Chewie and 3.5% on Mihi without surpassing Swap-VAE.The classifier gains were observed as generated samples were added to training data at 50%, 100%, and 200%.Source: §4.2 · p. 9, Figure 4 · p. 9

4.3 Model ablations: Testing our BlockSwap augmentation

• BlockSwap yields the highest decoding accuracies among tested variants; combining spatial and temporal augmentations outperforms either alone, while removing L2 with content swapping changes performance minimally.The authors report that including L2 offers additional flexibility in controlling performance.Source: §4.3 · pp. 9–10, Table 3 · p. 10

5 Conclusion

• Swap-VAE uses self-supervised alignment to decompose neural activity into latent subspaces, aiming to reveal meaningful and interpretable factors and generate realistic neural activity.The authors frame these representations as a way to gain insight into the relationship between neural activity and animal behavior.Source: §5 · p. 10
• Performance relies on the choice of augmentations, which may depend on the downstream task of interest.The authors identify examining how augmentations affect different downstream tasks as future work.Source: §5 · p. 10
• Disentanglement was more pronounced in Mihi than Chewie, but the authors leave open whether preparation or method weakness explains the difference.Mihi was instructed to wait for an auditory cue before moving; the authors say the preparation hypothesis warrants study across more individuals.Source: §5 · p. 10

A.1 Model implementation details

• The encoder has three linear layers with output sizes [d, k, k], and the symmetric decoder has three with sizes [k, k, d].Intermediate layers use batch normalization and ReLU; the decoder's last layer uses SoftPlus with beta = 1 and threshold = 20.Source: §A.1 · p. 15
• Synthetic models train for 100,000 iterations with Adam at learning rate 0.0001, and both synthetic and monkey experiments use batch size 256.The synthetic experiments use an Nvidia Titan RTX GPU.Source: §A.1 · p. 15
• Reach and temporal decoding use a one-layer linear decoder trained with Adam at learning rate 0.005 and weight decay 0.00001.The hyperparameters were selected by grid search on the supervised pre-trained feature space for reach directions.Source: §A.1 · p. 15

A.2 BlockSwap implementation details

• BlockSwap swaps content between two views while aligning the first latent-space part and reconstructing from both parts.In the illustration, pink and purple indicate shared content, while blue and green indicate styles.Source: Figure S1 · p. 15, §A.2 · p. 15

A.3 Dataset details

• The datasets contain M1 spiking activity from two rhesus macaques performing cursor reaches, with Mihi waiting for an auditory “Go” cue before moving.After removing zero-variance neurons, the datasets contain 163 and 148 neurons for Chewie days 1 and 2, and 163 and 152 for Mihi days 1 and 2.Source: §A.3 · p. 15

B Comparing performance across multiple initializations

• The section reports standard deviations for benchmark models and examines Swap-VAE variance across initializations.The supplied passage does not state the resulting values or comparisons.Source: B Comparing performance across… · p. 16

B.1 Stability of the benchmark models

• Benchmark-model accuracy tables report standard deviations over five random seeds, with MYOW results from a reproduction using BYOL-selected settings.Model weights, training, and linear evaluation layers were randomized by seed.Source: §B.1 · p. 16, Table S1 · p. 16, Table S2 · p. 16

B.2 Stability of the SwapVAE

• Swap-VAE maintains a gap over other methods in reach-direction accuracy across random initializations.The reported analysis separates variance from training and evaluating whole networks from variance introduced by retraining the linear evaluation layer.Source: Table S3 · p. 16, §B.2 · p. 16

C.1 Latent space visualizations

• Figure S2 uses tSNE to visualize content and style spaces across four datasets, with rows colored by reach direction, dynamics, and reach-direction trajectories.Content appears on the left and style on the right for each dataset.Source: Figure S2 · p. 17, §C.1 · p. 18

C.2 Time decoding details

• Temporal decoding predicts how far into a reach each sample falls, using the first nine bins as nine classification classes.The first bin marks the beginning of each reach.Source: §C.2 · p. 18

• Each Figure S3 plot shows reach and time decoding for a different neural recording.Source: Figure S3 · p. 18
• The plots compare decoding from the model’s Full, Content, and Style spaces with other baseline methods.Source: Figure S3 · p. 18
• For the supervised decoder, separate decoders are trained for the reaching and temporal decoding tasks.Source: Figure S3 · p. 18

D Additional details on the generative quality experiment

• To create a new datapoint, the experiment fixes the existing datapoint’s content latent and varies its style latent with noise sampled from 0.2 × N(0, I).Source: D Additional details on the generative… · p. 18
• The new content-and-style latent vector is passed through the decoder to generate instance-specific neural activity.Source: D Additional details on the generative… · p. 18
• Generated activities improve supervised models for Chewie-1 and Mihi-2, where the generative models surpass the supervised models, but not for Chewie-2 and Mihi-1.The authors surmise that generative-model performance may upper-bound supervised-model performance trained on generated activities; generated activities may still help other downstream tasks.Source: D Additional details on the generative… · p. 18

E.1 Different latent dimensions

• With the Full latent space fixed at 128 dimensions, models perform comparably across different Content–Style splits.At s = 128, the model converges to beta-VAE.Source: §E.1 · p. 19, Table S4 · p. 19
• With a 128-dimensional Content space, adding Style dimensions can yield better linear-evaluation reach-direction accuracy on Content features.The comparison is based on the Content space because only it is regularized by the alignment loss; the Full-space features are generally similar or slightly better.Source: §E.1 · p. 19, Table S5 · p. 19

E.2 Disentanglement score for different ablation models

• Removing alignment while retaining BlockSwap gives better disentanglement than removing BlockSwap while retaining alignment and the original reconstruction term.Source: §E.2 · p. 19
• Swap-VAE has the best disentanglement performance among the tested ablations.Table S6 reports multi-task disentanglement scores; No-Swap does not outperform beta-VAE, while Swap-only is closest to Swap-VAE.Source: §E.2 · p. 19, Table S6 · p. 20

E.3 More details about data augmentation operations

• Adding Pepper or Noise to Dropout and Temporal shift did not further improve performance, and the authors suggest generative models may not need as intense augmentations as recent self-supervised methods.Source: §E.3 · p. 20