Co-smoothing: A cross-validation framework

Let be spiking neural activity of N channels recorded over a finite window of time, i.e., a trial, and subsequently quantised into T time-bins. X_t,n represents the number of spikes in channel n during time-bin t. The dataset , partitioned as and , consists of S trials of the experiment. The latent-variable model (LVM) approach posits that each time-point in the data is a noisy measurement of a latent state .

To infer the latent trajectory is to learn a mapping . On what basis do we validate the inferred ? We cannot access the ground truth , so instead we test the ability of to predict unseen or held-out data. Data may be held-out in time, e.g., predicting future data points from the past, or in space, e.g., predicting neural activities of one set of neurons (or channels) based on those of another set. The latter is called co-smoothing [6].

The set of N available channels is partitioned into two: held-in channels and held-out channels. The S trials are partitioned into train and test. During training, both channel partitions are available to the model and during test, only the held-in partition is available. During evaluation, the model must generate the rate-predictions for the held-out partition. This framework is visualised in Fig 1A.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. Prediction framework and its relation to ground truth.

A. To evaluate a neural LVM with co-smoothing, the dataset is partitioned along the neurons and trials axes. B. The held-in neurons are used to infer latents , while the held-out serve as targets for evaluation. The encoder f and decoder g are trained jointly to maximise co-smoothing . After training, the composite mapping is evaluated on the test set. C. We hypothesise that models with high co-smoothing may have an asymmetric relationship to the true system, ensuring that model representation contains the ground truth, but not vice-versa. We reveal this in a synthetic student(S)-teacher(T) setting by the unequal performance of regression on the states in the two directions. denote decoding error of model v latents from model u latents .

https://doi.org/10.1371/journal.pcbi.1013789.g001

Importantly, the encoding-step or inference of the latents is done using a full time-window, i.e., analogous to smoothing in control-theoretic literature, whereas the decoding step, mapping the latents to predictions of the data is done on individual time-steps:

(1)

(2)

where the subscripts ‘’ and ‘’ denote partitions of the neurons (Fig 1B). During evaluation, the held-out data from test trials is compared to the rate-predictions from the model using the co-smoothing metric defined as the normalised log-likelihood, given by:

(3)

(4)

where is poisson log-likelihood, is a the mean rate for channel n, and is the total number of spikes, following [6].

Thus, the inference of LVM parameters is performed through the optimisation:

(5)

using , without access to the test trials from . For clarity, apart from (5), we report only , omitting the superscript.

Good co-smoothing does not guarantee correct latents

It is common to assume that being able to predict held-out parts of will guarantee that the inferred latent aligns with the true one [6,14,16–28]. To test this assumption, we use a student-teacher scenario where we know the ground truth. To compare how two models (u,v) align, we infer the latents of both from , then do a regression from latents of u to v. The regression error is denoted (i.e. for teacher to student decoding). Contrary to the above assumption, we hypothesise that good prediction guarantees that the true latents are contained within the inferred ones (low ), but not vice versa (Fig 1C). It is possible that the inferred latents possess additional features, unexplained by the true latents (high ).

We demonstrate this phenomenon in three different student-teacher scenarios: task-trained RNNs, Hidden Markov Models (HMMs) and linear gaussian state space models. We start with RNNs, as they are a standard tool to investigate computation through dynamics in neuroscience [29], and expand upon the other models in the appendix. A 128-unit RNN teacher (Methods) is trained on a 2-bit flip-flop task, inspired by working memory experiments. The network receives input pulses and has to maintain the identity of the last pulse (see Methods). The student is a sequential autoencoder, where the encoder f is composed of a neural network that converts observations into an initial latent state, and another recurrent neural network that advances the latent state dynamics [29] (see Methods).

We generated a dataset of observations from this teacher, and then trained 30 students with latent-dimensionality 3–64 on the same teacher data using gradient-based methods (see Methods). Co-smoothing scores of students increased with the size of the latents, but are high for models in the range of 5-15 dimensional latents (S1 Fig). Consistent with our hypothesis, the ability to decode the teacher from the student was highly correlated to the co-smoothing score (Fig 2 top left). In contrast, the ability to decode the student from the teacher has a very different pattern. For students with low co-smoothing, this decoding is good – but meaningless. For students with high co-smoothing, there is a large variability, and little correlation to the co-smoothing score (Fig 2 top right). In this simple example, it would seem that one only needs to increase the dimensionality of the latent until co-smoothing saturates. This minimal value would satisfy both demands. This is not the case for real data, as will be shown below.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Upper panel.

Several students, sequential autoencoders (SAE, see Methods), are trained on a dataset generated by a single teacher, a noisy GRU RNN trained on a 2-bit flip flop (2BFF, see Methods). The StudentTeacher decoding error is low and tightly related to the co-smoothing score. The TeacherStudent decoding error is more varied and uncorrelated to co-smoothing. A score of corresponds to predicting the mean firing-rate for each neuron at all trials and time points. Green and red points are representative “Good” and “Bad” students respectively, whose latents are visualised below along-side the ground truth . The visualisations are projections of the latents along the top three principal components of the data. The ground truth latents are characterised by 4 stable states capturing the 2² memory values. This structure is captured in the “Good” student. The bad student also includes this structure in addition to an extraneous variability along the third component. Lower panel. The same experiment conducted with HMMs. The teacher is a nearly deterministic 4-cycle and students are fit to its noisy emissions. Dynamics in selected models are visualised. Circles represent states, and arrows represent transitions. Circle area and edge thickness reflect fraction of visitations or volume of traffic after sampling the HMM over several trials. The colours also reflect the same quantity – brighter for higher traffic. Edges with values below 0.01 are removed for clarity (S5 Fig). The teacher (M = 4) is a 4-cycle. Note the prominent 4-cycles (orange) present in the good student (M = 10), and the bad student (M = 8). In the good student, the extra states are seldom visited, whereas in the bad student there is significant extraneous dynamics involving these states (dark arrows).

https://doi.org/10.1371/journal.pcbi.1013789.g002

What is it about a student model, that produces good co-smoothing with the wrong latents? It’s easiest to see this in a setting with discrete latents, so we first show the HMM teacher and two exemplar students – named “Good” and “Bad” (marked by green and red arrows in S3 FigAB) – and visualise their states and transitions using graphs in Fig 2. The teacher is a cycle of 4 steps. The good student contains such a cycle (orange), and the initial distribution is restricted to that cycle, rendering the other states irrelevant. In contrast, the bad student also contains this cycle (orange), but the initial distribution is not consistent with the cycle, leading to an extraneous branch converging to the cycle, as well as a departure from the main cycle (both components in dark colour). Note that this does not interfere with co-smoothing, because the emission probabilities of the extra states are consistent with true states, i.e., the emission matrix conceals the extraneous dynamics. In the RNN, we see a qualitatively similar picture, with the bad students having dynamics in task-irrelevant dimensions (Fig 2 “Bad” S).

Few-shot prediction selects better models

Because our objective is to obtain latent models that are close to the ground truth, the co-smoothing prediction scores described above are not satisfactory. Can we devise a new prediction score that will be correlated with ground truth similarity? The advantage of prediction benchmarks is that they can be optimised, and serve as a common language for the community as a whole to produce better algorithms [30].

We suggest few-shot co-smoothing as a complementary prediction score to co-smoothing, to be used on models with good scores on the latter. Similarly to standard co-smoothing, the functions g and f are trained using all trials of the training data (Fig 3A). The key difference is that a separate group of neurons is set aside (Table 1), and only k trials of these neurons are used to estimate a mapping (Fig 3B), similar to g in (2). The neural LVM is then evaluated on both the standard co-smoothing using and the few-shot version using (Fig 3C).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. Co-smoothing and few-shot co-smoothing; a composite evaluation framework for Neural LVMs.

A. The encoder f and decoder g are trained jointly using held-in and held-out neurons. B. A separate decoder is trained to readout k-out neurons using only k trials. Meanwhile, f and g are frozen. C. The neural LVM is evaluated on the test set resulting in two scores: co-smoothing and k-shot co-smoothing .

https://doi.org/10.1371/journal.pcbi.1013789.g003

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Table 1. Dimensions of real and synthetic datasets.

Number of train and test trials , , time-bins per trial for co-smoothing T, and forward-prediction , held-in, held-out and k-out neurons , , . ^†In all the NLB [6] datasets as well the RNN dataset we use the same set of neurons for and .

https://doi.org/10.1371/journal.pcbi.1013789.t001

For small values of k, the scores can be highly variable. To reduce this variability, we repeat the procedure s times on independently resampled sets of k trials, producing s estimates of , each with its own score . For each student , we then report the average score across the s resamples. A theoretical analysis of the choice of k is given in the next section, with practical guidelines provided in S2 Fig. The number of resamples s is chosen empirically to ensure high confidence in the estimated average (Methods).

To demonstrate the utility of the proposed prediction score, we return to the RNN students from Fig 2 and evaluate for each. This score provides complementary information about the models, as it is uncorrelated with standard co-smoothing (Fig 4A), and it is not merely a stricter version of co-smoothing (S6 Fig). Since we are only interested in models with good co-smoothing, we restrict attention to students satisfying . Among these students, despite their nearly identical co-smoothing scores, the k-shot scores are strongly correlated with the ground-truth measure (Fig 4B). Together, these findings suggest that simultaneously maximizing and —both prediction-based objectives—produces models with low and , yielding a more complete measure of model similarity to the ground truth.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4. Few-shot prediction selects better models.

A. Few-shot measures something new. Student models with high co-smoothing have highly variable 2-shot co-smoothing, which is uncorrelated to co-smoothing. Error bars reflect standard error of the mean across several few-shot regressions (see Methods). B. For the set of students with high co-smoothing, i.e., satisfying , 2-shot co-smoothing to held-out neurons is negatively correlated with decoding error from teacher-to-student. Green and red points represent the example “Good” and “Bad” models (Fig 2).

https://doi.org/10.1371/journal.pcbi.1013789.g004

Why does few-shot work?

The example HMM and RNN students of Fig 2 can help us understand why few-shot prediction identifies good models. The students differ in that the bad student has more than one state corresponding to the same teacher state. Because these states provide the same output, this feature does not hurt co-smoothing. In the few-shot setting, however, the output of all states needs to be estimated using a limited amount of data. Thus the information from the same amount of observations has to be distributed across more states. We make this data efficiency argument more precise in three settings: linear regression (LR), HMMs, and binary classification prototype learning (BCPL).

In the case of LR, the teacher latent is a scalar random variable z and the student latent is a random p-vector, whose first coordinate is z and the remaining p–1 coordinates are the extraneous noise:

(6)

where . In other words, a single teacher state is represented by several possible student states.

Next, we model the neural-data – noisy observations of the teacher latent , where . The few-shot learning is captured by minimum-norm k-shot least-squares linear regression:

(7)

where is the 2-norm.

The generalisation error of the few-shot learner is given by:

(8)

where is the true mapping.

We solve for as using the theory of [31], and demonstrate a good fit to numerical simulations at finite p,k (Methods). We do similar analyses for Bernoulli HMM latents with maximum likelihood estimation of the emission parameters (Methods) and BCPL [32] (Methods).

Across the three scenarios, model performance decreases with extraneous variability (Fig 5). Crucially, this difference appears at small k, and vanishes as . With HMMs and BCPL this is a gradual decrease, while in LR, there is a known critical transition at p = k [31,33,34].

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5. Theoretical analysis of k-shot learner performance as a function of k and extraneous noise , in three different settings.

Points show numerical simulations and dashed lines show analytical theory. A. Hidden Markov Models (HMMs) (Methods), Bernoulli observations, MLE estimator. B. Minimum norm least squares linear regression with and p = 50 (main text and Methods). C. binary classification, prototype learning (Methods).

https://doi.org/10.1371/journal.pcbi.1013789.g005

Interestingly, the scenarios differ in the bias-variance decomposition of their performance deficits. In LR, extraneous noise leads to increased bias with identical variance (Methods, Claim 2), whereas in the HMM and BCPL, it leads to increased variance and zero bias (Methods, (28) and (52) respectively).

How does one choose the value of k in practice? The intuition and theoretical results suggest that we want the smallest possible value. In real data, however, we expect many sources of noise that could make small values impractical. For instance, for low firing rates, small k values can mean that some neurons will not have any spikes in k trials and thus there will nothing to regress from. Our suggestion is therefore to use the smallest value of k that allows robust estimation of few-shot co-smoothing. (S2 Fig) shows the effect of this choice for various datasets.

SOTA LVMs on neural data

In previous sections, we showed that models with near perfect co-smoothing may possess latents with extraneous dynamics. We established this in a synthetic student-teacher setting with RNNs, HMMs and LGSSM models.

To show the applicability in more realistic scenarios, we consider four datasets mc_maze_20 [35], mc_rtt_20 [36], dmfc_rsg_20 [37], area2_bump_20 [38] from the Neural Latent Benchmarks suite [6] (see Methods). They consist of neural activity (spikes) recorded from various cortical regions of monkeys as they perform specific tasks. The 20 indicates that spikes were binned into 20ms time bins. We trained several SpatioTemporal Neural Data Transformers (STNDTs) [39–42], that achieve near state-of-the-art (SOTA) co-smoothing on these datasets. We evaluate co-smoothing on a test set of trials and define the set of models with the best co-smoothing (see Methods and Table 1).

A key component of training modern neural network architectures such as STNDT is the random sweep of hyperparameters, a natural step in identifying an optimal model for a specific data set [19]. This process generates several candidate solutions to the optimisation problem (5), yielding models with similar co-smoothing scores but, as we demonstrate in this section, varying amounts of extraneous dynamics.

Two proxies for : cycle consistency and cross-decoding.

To reveal extraneous dynamics in the synthetic examples (RNNs, HMMs), we had access to ground truth that enabled us to directly compare the student latent to that of the teacher. With real neural data, we do not have this privilege. This limitation has been recognised in the past and a proxy was suggested [8,29,43] – cycle consistency. Instead of decoding the student latent from the teacher latent, cycle consistency attempts to decode the student latent from the student’s own rate prediction . In our notation this is (Fig 6A and Methods). If the student has perfect co-smoothing (see S3 Appendix), this should be equivalent to as it would ensure that teacher and student have the same rate-predictions .

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 6. Cycle consistency and cross-decoding as a proxy for distance to the ground truth in the absence of ground-truth.

A Cycle consistency [8,29,43] involves learning a mapping g⁻¹ from the rates back to the latents (see Methods). B The latents of each pair of models are cross-decoded from one another. Minimal models can be fully decoded by all models but extraneous models only by some. C Cross-decoding matrix for SAE NODE models trained on data from the NoisyGRU (Fig 2). D, E For models with high co-smoothing () the proxy metrics – cross-decoding column average , and cycle-consistency ) – are both highly correlated to ground truth .

https://doi.org/10.1371/journal.pcbi.1013789.g006

Because we cannot rely on perfect co-smoothing, we also suggest a novel metric – cross-decoding – where we compare the models to each other. The key idea is that all high co-smoothing models contain the teacher latent. One can then imagine that each student contains a selection of several extraneous features. The best student is the one containing the least such features, which would imply that all other students can decode its latents, while it cannot decode theirs (Fig 6B). Instead of computing and as in Fig 2, we perform decoding from latents of model u to model v () for every pair of models u and v using linear regression and evaluating an R² score for each mapping (see Methods). In Fig 6C the results are visualised by a matrix with entries for all pairs of models u and v. The ideal model v^* would have no extraneous dynamics, therefore, all the other models should be able to decode its latents perfectly, i.e., . Provided a large and diverse population of models only the ‘pure’ ground truth would satisfy this condition. To evaluate how close a model v is to the ideal v^* we propose a simple metric: the column average . This will serve as proxy for the distance to ground truth, analogous to in Fig 4. We validate this procedure using the RNN student-teacher setting in Fig 6D, where we show that is highly correlated to the ground truth measure . We also validate cycle-consistency against using the RNN setting (Fig 6E). In both cases we find a high correlation between the metrics.

Having developed a proxy for the ground truth we can now correlate it with the few-shot co-smoothing to held-out neurons. Following the discussion in the previous section, we choose the smallest value of k that ensures no trials with zero spikes (S2 Fig). Fig 7 shows a negative correlation of with both proxy measures and across the STNDT models in the four data sets. Moreover, regular co-smoothing for the same models is relatively uncorrelated with these measures. As an illustration of the latents of different models, Fig 7(bottom) shows the PCA projection of latents from two STNDT models trained on mc_maze_20. Both have high co-smoothing scores but differ in their few-shot scores . We note smoother trajectories and better clustering of conditions in the model with higher . We also quantify the ability to decode behavior from these two models, and found the top-PCs perform better in the “Good” model (S7 Fig).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 7. Few-shot scores correlate with the proxies of distance to the ground truth, cycle-consistency and the cross-decoding column average .

We train several STNDT models on four neural recordings from monkeys [35–38], curated by [6] and filter for models with high co-smoothing . The few-shot co-smoothing scores negatively correlate with the two proxies and (orange points), while regular co-smoothing (turquoise points) does not (one-tailed p-values shown for p < 0.05 and ^*** for p < 0.001). Green and red arrows indicate the extreme models whose latents are visualised below. values may be compared against an EvalAI leaderboard [6]. Note that we evaluate using an offline train-test split, not the true test set used for the leaderboard scores, for which held-out neuron data is not publicly accessible. (Bottom) Principal component analysis of the latent trajectories of two STNDT models trained on mc_maze_20 with similar co-smoothing scores but contrasting few-shot co-smoothing. The “Good” model scores , and the “Bad” model , . The trajectories are coloured by task conditions and start at a circle and end in a triangle.

https://doi.org/10.1371/journal.pcbi.1013789.g007

Glossary

1. Latent variable model (LVM) (f and g) : A function mapping neural time-series data to an inferred latent space (f). The latents can then be used to predict held-out data (g).
2. Smoothing : mapping a sequence of observations to a sequence of inferred latents . It is often formalised as a conditional probability .
3. Extraneous dynamics : the notion that inferred latent variables may contain features and temporal structure not present in the true system from which the data was observed.
4. Co-smoothing () : A metric evaluating LVMs by their ability to predict the activity of held-out neurons provided held-in neural activity over a window of time. The two sets of neurons are typically random subsets from a single population.
5. Few-shot co-smoothing () : A variant of co-smoothing in which the mapping from latents to held-out neurons () is learned from a small number of trials.
6. State-of-the-art (SOTA) : the best performing method or model current available in the field. This is usually based on a specific benchmark, i.e., a dataset and associated evaluation metric. In active fields the SOTA is constantly improving.
7. Cycle consistency () : a measure of extraneousness of model latents as compared to their rate predictions. Computed by learning and evaluating the inverse mapping from rate predictions to latents.
8. Cross-decoding () : another measure of model extraneousness. It is evaluated on a population of models trained on the same dataset. It involves regressing from one model latents to another model, for all pairs in the population. A scalar measure is the obtained for each model: the cross-decoding column mean . It reflects the average ‘decodability’ of a model, by all the other models.

Student-teacher Recurrent Neural Networks (RNN)

Both teacher and student are based on an adapted version of [29]. In the following, we provide a brief description.

Teacher

We train a noisy 64 Gated Recurrent Unit (NoisyGRU) RNN [50], on a 2-bit flip flop 2BFF task [3], implemented by [29]. The GRU RNN follows standard dynamics, which we repeat here using the typical notation of GRUs. This notation is not consistent with the Results section, and we explain the relation below.

(9)

(10)

(11)

(12)

(13)

where , , , , , , , , , are trainable parameters. The latent used in the Results section () is the hidden unit activity h. After model training, the NoisyGRU units are subsampled, centered, normalised, and rectified to give synthetic neural firing rates - which are of the Results section. These firing rates are used to define a stochastic Poisson process to generate the synthetic neural data.

Students

The student models are sequential autoencoders (SAEs) consisting of a bidirectional GRU that predicts the initial latent state, a Neural ODE (NODE) that evolves the latent dynamics (together these form the encoder, f, under our notation), and a linear readout layer mapping the latent states to the data (the decoder, g). We train several randomly initialised models with a range of latent dimensionalities (3, 5, , 32, 64). Models are trained to minimise a Poisson negative loglikelihood reconstruction loss, using the Adam [51] optimiser.

Student-teacher Hidden Markov Models (HMMs)

We choose both student and teacher to be discrete-space, discrete-time Hidden Markov Models (HMMs). As a teacher model, they simulate two important properties of neural time-series data: its dynamical nature and its stochasticity. As a student model, they are perhaps the simplest LVM for time-series, yet they are expressive enough to capture real neural dynamics ( of 0.29 for HMMs vs. 0.24 for GPFA and 0.35 for LFADS, on mc_maze_20). The HMM has a state space , and produces observations (emissions in HMM notation) along neurons , with a state transition matrix , emission model and initial state distribution . More explicitly:

(14)

The same HMM can serve two roles: a) data-generation by sampling from (14) and b) inference of the latents from data on a trial-by-trial basis:

(15)

i.e., smoothing, computed exactly with the forward-backward algorithm [52]. Note that although z is the latent state of the HMM, we use its posterior probability mass function as the relevant intermediate representation because it reflects a richer representation of the knowledge about the latent state than a single discrete state estimate. To make predictions of the rates of held-out neurons for co-smoothing we compute:

(16)

As a teacher, we constructed a 4-state model of a noisy chain , with , , and sampled once and frozen (Fig 2, left). We generated a dataset of observations from this teacher (see Table 1).

For each student, we evaluate . This involves estimating the bernoulli emission parameters , given the latents using (26) and then generating rate predictions for the k-out neurons using (16).

HMM training

HMMs are traditionally trained with expectation maximisation, but they can also be trained using gradient-based methods. We focus here on the latter as these are used ubiquitously and apply to a wide range of architectures. We use an existing implementation of HMMs with differentiable parameters: dynamax [53] – a library of differentiable state-space models built with jax.

We seek HMM parameters that minimise the negative log-likelihood loss, L of the held-in and held-out neurons in the train trials:

(17)

(18)

To find the minimum we do full-batch gradient descent on L, using dynamax together with the Adam optimiser [51] .

Decoding across HMM latents

Consider two HMMs u and v, of sizes M(u) and M(v), both candidate models of a dataset . Following (15), each HMM can be used to infer latents from the data, defining encoder mappings f^u and f^v. These map a single trial i of the data to and .

Since HMM latents are probability mass functions, we do not do use linear regression to learn the mappings across model latents. Instead we perform a multinomial regression from to .

(19)

(20)

where , and is the softmax. During training we sample states from the target PMFs thus arriving at a more well-known problem scenario: classification of M(v)-classes. We optimise W and to minimise a cross-entropy loss to the target using the fit() method of sklearn.linear_model.LogisticRegression.

We define decoding error, as the average Kullback-Leibler divergence D_KL between target and predicted distributions:

(21)

where D_KL is implemented with scipy.special.rel_entr.

In section and Fig 1, the data X is sampled from a single teacher HMM, , and we evaluate and for each student notated simply as .

Analysis of LVMs without access to ground truth

We denote the set of high co-smoothing models as those satisfying for Fig 4 and in Fig 7, , the the encoders and decoders respectively. Note that STNDT is a deep neural network given by the composition , and the choice of intermediate layer whose activity is deemed the ‘latent’ is arbitrary. Here we consider g the last ‘read-out’ layer and f to represent all the layers up-to g.

Few-shot co-smoothing

To perform few-shot co-smoothing, we learn , which takes the same form as g, a Poisson Generalised Linear Model (GLM) for each held-out neuron. We use sklearn.linear_model.PoissonRegressor, which has a hyperparameter alpha, the amount of l2 regularisation. For the results in the main text, in Fig 7, we select . We partition the training data into several random subsets of k trials and train an independently initialised GLM on each subset. Each GLM is then evaluated on a fixed test set of trials (Fig 3), yielding a score for each subset. We report the mean over such repetitions, , along with the standard error of the mean (error bars in Fig 4, Fig 7). Scores are more variable at small k, so we need more repetitions to better estimate the average score. To implement this in a standarised way, we incorporate this chunking of data into several subsets in the nlb_tools library (Data and code availability). This way we ensure that all models are trained and tested on identitical subsets. We report the compute-time for few-shot co-smoothing in S2 Appendix.

Cross-decoding

We perform a cross-decoding from the latents of model u, , to those of model v, , for every pair of models u and v using a linear mapping implemented with sklearn.linear_model.LinearRegression:

(22)

minimising a mean squared error loss. We then evaluate a R² score (sklearn.metrics.r2_score) of the predictions, , and the target, , for each mapping. We define the decoding error . The results are accumulated into a matrix (see Fig 6).

Cycle consistency

We evaluate cycle-consistency [8,29] for a model u also using a linear mapping from its rate predictions back to its latents implemented with sklearn.linear_model.LinearRegression:

(23)

again minimising a squared error loss. As in cross-decoding we evaluate R² score (sklearn.metrics.r2_score) and the decoding error (Fig 6A).

Summary of Neural Latent Benchmark (NLB) datasets

Here are brief descriptions of the datasets used in this study. All datasets were collected from macaque monkeys performing sensorimotor or cognitive tasks. More comprehensive details can be found in the Neural Latents Benchmark paper [6].

1. mc_maze [35] Motor cortex recordings during a delayed reaching task where monkeys navigated around virtual barriers to reach visually cued targets. The task involved 108 unique maze configurations, with several repeated trials for each one, thus serving as a “neuroscience MNIST”. We choose this dataset to visualise the latents in Fig 7.
2. mc_rtt [36] Motor cortex recordings during naturalistic, continuous reaching toward randomly appearing targets without imposed delays. The task lacks trial structure and includes highly variable movements, emphasizing the need for modeling unpredictable inputs and non-autonomous dynamics.
3. dmfc_rsg [37] Recordings from dorsomedial frontal cortex during a time-interval reproduction task, where monkeys estimated and reproduced time intervals between visual cues using eye or hand movements. The task involves internal timing, variable priors, and mixed sensory-motor demands.

Somatosensory cortex recordings during a visually guided reach task in which unexpected mechanical bumps to the limb occurred in half of the trials. The task probes proprioceptive feedback processing and requires modeling input-driven neural responses.

Dimensions of datasets

We analyse several datasets in this work. Three synthetic datasets generated by an RNN, HMM (Methods, Fig 2) and LGSMM (S4 Fig) and the four datasets from the Neural Latent Benchmarks (NLB) suite [6,35–38]. In Table 1, we summaries the dimensions of all these datasets. To evaluate k-shot on the existing SOTA methods while maintaining the NLB evaluations, we conserved the forward-prediction aspect. During model training, models output rate predictions for future time bins in each trial, i.e., (1) and (2) are evaluated for while input remains as . Although we do not discuss the forward-prediction metric in our work, we note that the SOTA models receive gradients from this portion of the data.

In all the NLB datasets as well as the RNN dataset we reuse held-out neurons as k-out neurons. We do this to preserve NLB evaluation metrics on the SOTA models, as opposed to re-partitioning the dataset resulting in different scores from previous works. This way existing co-smoothing scores are preserved and k-shot co-smoothing scores can be directly compared to the original co-smoothing scores. The downside is that we are not testing the few-shot on ‘novel’ neurons. Our numerical results (Fig 7) show that our concept still applies.

Theoretical analysis of few shot learning in HMMs.

Consider a student-teacher scenario as in section . We let T = 2 and use a stationary teacher . Now consider two examples of inferred students. To ensure a fair comparison, we use two latent states for both students. In the good student, , these two states statistically do not depend on time, and therefore it does not have extraneous dynamics. In contrast, the bad student, , uses one state for the first time step, and the other for the second time step. A particular example of such students is:

(24)

(25)

where each vector corresponds to the two states, and we only consider two time steps.

We can now evaluate the maximum likelihood estimator of the emission matrix from k trials for both students. In the case of bernoulli HMMs the maximum likelihood estimate of given a fixed f and k trials has a closed form:

(26)

We consider a single neuron, and thus omit n, reducing the estimates to:

(27)

where C_t is the number of times x = 1 at time t in k trials. We see that C_t is a sum of k i.i.d. Bernoulli random variables (RVs) with the teacher parameter B^*, for both .

Thus, and are scaled binomial RVs with the following statistics:

(28)

The test loss is given by . For , for both values of t, and for , and . Ultimately,

(29)

(30)

To see how these variations affect the test loglikelihood L of the few-shot regression on average, we do a taylor expansion around B^*, recognising that the function is maximised at B^*, so .

(31)

(32)

(33)

We see that this second order truncation of the loglikelihood is decomposed into a bias and a variance term. We recognise that the bias term goes to zero because we know the estimator is unbiased (28). To compute the variance term, we compute the hessians which differ for the two models:

(34)

where .

Incorporating these hessians into (33), we obtain:

(35)

(36)

Fig 5A shows these analytical results against the left hand side of (35) and (36) evaluated numerically.

Theoretical analysis of ridgeless least squares regression with extraneous noise.

Teacher latents generate observations x_i:

(37)

where is observation noise.

In this setup there is no time index: we consider only a single sample index i.

We consider candidate student latents, , that contain the teacher along with extraneous noise, i.e:

(38)

where is a vector of i.i.d. extraneous noise, and is the identity matrix.

We study the minimum l₂ norm least squares regression estimator on k training samples:

(39)

with the regression weights . More succinctly, , where .

Note that, by construction, the true mapping is:

(40)

Test loss or risk is a mean squared error:

(41)

given a test sample . The error can be decomposed as:

(42)

The scenario described above is a special case of [31]. What follows is a direct application of their theory, which studies the risk R, in the limit such that , to our setting. The alignment of the theory with numerical simulations is demonstrated in Fig 5B.

Claim 1. , i.e., the underparameterised case k > p.

B = 0 and the risk is just variance and is given by:

(43)

with no dependence on .

Proof: This is a direct restatement of Proposition 2 in [31].

Claim 2. , i.e., the overparameterised case k < p.

The following is true as :

(44)

(45)

Proof: For the non-isotropic case [31] define the following distributions based on the eigendecomposition of .

(46)

(47)

In the limit we take . This greatly simplifies calculations and nevertheless provide a good fit for numerical results with finite k and p. We solve for using equation 12 in [31].

(48)

We then compute the limiting values of B and V:

(49)

(50)

Substituting completes the proof.

The extraneous noise, , influences the risk of ridgeless regression only in the regime k < p, and its effect is confined to the bias term, leaving the variance unaffected. In contrast, observation noise contributes exclusively to the variance term. Consequently, the dependence of the risk on persists even in the absence of observation noise, i.e., when .

Fig 5B presents the theoretical predictions alongside the empirical average k-shot performance of minimum-norm least-squares regression, computed numerically using the function numpy.linalg.lstsq.

Theoretical analysis of prototype learning for binary classification with extraneous noise

Teacher latents are distributed as , that is either or with probability , representing two classes a and b respectively.

We consider candidate student latents, , that contain the teacher along with extraneous noise, i.e:

(51)

where is a M-vector of i.i.d. extraneous noise, and is the identity matrix and .

We consider the prototype learner , , where and are the sample means of k latents from class and k latents from class respectively. The classification rule is given by the sign of : classifying the input as if positive and otherwise.

This setting is a special case of [54]. They provide a theoretical prediction for average few-shot classification error rate for class , , given by where is a monotonously decreasing function.

(52)

the difference of the population centroids of the two classes.

In our case this reduces to:

(53)

To obtain this we note radii of manifold is with an average radius and participation ratio .

We substitute .

The bias term is zero since .

The and terms are both zero.

The participation ratio D_a = M. Our construction is symmetric in that .

The classification error, , decreases monotonically with the number of samples k, tending to zero as for all finite values of . In contrast, increases monotonically with extraneous noise , deviating significantly from zero once .

Fig 5C presents the numerically computed error in comparison with the theoretical prediction given in (53).