`How to test for`: overfitting in CEBRA Supervised Embeddings

[jgraving]

Hi @stes and @MMathisLab et al.,

I've been testing CEBRA supervised embeddings and observed that even when the neural data are shuffled, the resulting embeddings remain highly consistent with those for the original (unshuffled) data—with high decoding accuracy and low InfoNCE loss. This suggests that the model is overfitting to the supervised labels, and high consistency and decoding performance is the null expectation for any sufficiently flexible encoder network. Naturally, this calls into question the soundness of the method as well as the validity of any scientific inferences drawn from the use of such a method.

Edit: This initially suggested that the model might be overfitting to the supervised labels. However, after further analysis and clarification, it is clear this is not a widespread issue, and only occurs in specific circumstances with small datasets, large models and extended training time.

• Method:
  Experiments were conducted using example rat and monkey data with the provided code for CEBRA supervised embeddings. Embeddings for the original (unshuffled) and shuffled (randomly permuted) neural data were compared, with the supervised labels left intact (unshuffled).

• Code & Models:
  Code, figures, and saved models for the experiments are available at cebra_shuffle_overfit.

• Observations:

  • When the neural data are shuffled, the embeddings remain highly consistent with those for the original (unshuffled) data.
  • Decoding accuracy remains high.
  • InfoNCE loss remains low.

Concerns

- Overfitting:
Experiments strongly indicate that the encoder network for the neural data is readily overfitting to the supervised labels.

Edit: Upon further analysis and discussion, this overfitting behavior is observed only with larger models trained for an extended period. When proper validation protocols—such as using a train/validation split and monitoring loss curves—are employed, this failure mode is detectable and manageable.

- Lack of Proper Validation:
The software package doesn't provide any obvious functionality for proper validation and testing of generalization for supervised embeddings, and it doesn’t appear that this issue is discussed or addressed in the paper.

Edit: The package does provide validation functionality, which has been made more prominent in the demos. The published results are based on heldout test sets.

• Documentation Issues:
  The documentation encourages users to increase the network size ad libitum for a "better embedding" without explicitly discussing potential pitfalls. This effectively suggests that users should overfit embeddings to their data. Documentation reference:

  num_hidden_units (int): The number of dimensions to use within the neural network model.
  Higher numbers slow down training, but make the model more expressive and can result in a better embedding.
  Especially if you find that the embeddings are not consistent across runs, increase :py:attr:~.num_hidden_units and :py:attr:~.output_dimension 
  to increase the model size and output dimensionality.
  |Default:| 32.

- Impact:
These limitations raise serious concerns about the central claims and soundness of the work, especially given the high-impact venue where the paper was published and the seemingly widespread adoption of the method.

Proposed Next Steps

1. Paper Correction and Software Update:
The paper should be corrected and the software package updated to address these limitations—or retracted if the issues cannot be adequately resolved.

2. User Awareness:
You, as the maintainers/authors, should use GitHub and your social media channels to alert users to this issue, ensuring they can make informed decisions about using the method.

Edit: After further review and clarification (see conversation below), it’s clear that the overfitting observed under specific conditions (large models and extended training) is expected behavior that can be managed with proper validation. These statements overestimated the scope of the issue.

I invite you as maintainers, as well as others contributors and users, to review these findings and investigate this problem yourselves. This issue aims to ensure that you and other users are aware of this critical limitation and that necessary actions are taken to address it.
Edit: Demos have been updated to make validation more prominent and documentation on best practices has been added.

Best regards,
Jake Graving

[MMathisLab]

Hi @jgraving, thanks for your post. The issue you observe can occur when on limited data, too large of models get heavily overtrained, and no train/validation splits or sanity checks on the training curves are performed to uncover such an issue. This is true for any deep learning-based approach.

We therefore want to immediately push back that our paper should be retracted or corrected; we never make any claims that it's impossible to overfit, and below we provide both intuition about why your experiments are not what we show in our paper, and that your experiments do not invalidate our claims (i.e., plainly speaking, we show results with reasonable and justifiable model parameters and training).

Shuffle controls:

In fact, it is entirely expected that the distribution of the labels shape the embedding (see Proposition 7, Supplementary Note 2). In Figure 2c we therefore show shuffle controls demonstrating that if labels are shuffled, it is not possible to fit them.

In your example, you shuffle the neural data. If the model has sufficient capacity and the data is limited, it can fit a “lookup table” from input data to the output embedding to match the label distribution (because this is intact and still has structure), which is what you observe in your example.

​​To give more context for the two shuffling schemes:

(1) In our paper, we shuffle the labels across time. When the time_delta strategy for positive sampling is used, this changes the distribution of the positive samples. We do this kind of control to show that fitting a model on nonsensical label structure (in a real experiment, this would be a behavior time series without connection to the neural data) is not possible (again, assuming a sufficiently large dataset).

(2) In your example, you shuffle the neural data across time, keeping the labels intact. Now the question is whether the label structure can be forced on a latent representation of the shuffled data. This will be possible as long as the model has enough capacity to fit a lookup table, where for each timepoint the embedding is arranged to fit the behavior.

Overfitting:

For hypothesis testing applications of CEBRA, it is crucial to closely monitor the loss curves. A typical signature of overfitting is a high loss in the beginning (around log N, typically a bit smaller due to numerical effects), and then after many iterations a drop when the model enters the overfitting regime. This is exactly what is observable in your dataset.

Below we plot your provided achilles embeddings (neural shuffled and non-shuffled) and the expected behavior can be clearly seen: (a) smaller models have the overfitting point much later in training, larger models overfit earlier (b) overfitting leads to a sudden drop very late in training:

For ease of reference, in our publication we used hidden units = 32. If we use your neural shuffling scheme, at a reasonable point in training, we see a clear difference between shuffle and non-shuffled, as expected:

Consistency

In your example, visualizing the consistency across multiple steps of training and/or tracking consistency on both train and validation set would be important, if in doubt.

For example, here is the four rats from the hippocampus dataset with your neural-shuffle, time-only, and non-shuffled; at 5K steps this is the result:

Hyperparameter tuning:

Contrary to your claim, we do provide built-in tooling for easily plotting loss curves and performing train/validation splits of data before settling on parameters.

In our user guide we suggest performing a train/validation split to be sure that, in addition to shuffle controls, one’s parameters don’t drive overfitting (the essence of a train/val). To note, just like in UMAP or tSNE where there is not a train vs. test split typically, one doesn’t necessarily need to only use validation for comparing models. Users should just choose parameters appropriately before training on the full data.

For decoding applications, it is crucial to perform a train/validation split.

Updates: let’s be constructive

Updated demo

While our main hippocampus demo notebook already included how to do the train/validation, we added a new section at the top to highlight this more PR here. We see this as a step to validate parameters as we say in our user guide. As stated above, just like in UMAP or tSNE where there is not a train vs. test split typically, one doesn’t necessarily need to only use validation for comparing models. Users should just choose parameters appropriately before training on the full data.

Note in User Guide:

In our 9 Steps to work with CEBRA, the Step 3, already in our docs clearly highlights to users they can use a train/validation set:

 # 3. Split data and labels
     (
         train_data,
         valid_data,
         train_discrete_label,
         valid_discrete_label,
         train_continuous_label,
         valid_continuous_label,
     ) = train_test_split(neural_data,
                         discrete_label,
                         continuous_label,
                         test_size=0.3)

But, we added a PR to update the docs to educate users about spotting this issue when applying on their datasets earlier on:

#. To note, you should still be mindful of performing train/validation splits and shuffle controls to avoid overfitting <https://developers.google.com/machine-learning/crash-course/overfitting/overfitting>_.

Best Regards,
Mackenzie @MMathisLab and Steffen @stes

[jgraving]

Hi @stes and @MMathisLab,

Thanks for the explanation and updates. I realize now my earlier suggestion to retract or correct the paper was an overstatement. My intent was simply to highlight a potentially important issue, which you have addressed, so I will close the issue.

Best,
Jake

[jgraving]

Hi @stes and @MMathisLab,

I'm unable to recover a model that doesn't overfit using your example code and "reasonable" default parameters. Your example uses the InfoNCE loss to detect overfitting which is dominated by the ln(batch size) term from the logsumexp and maybe gives the illusion of good performance on the test set by adding a large constant value to both losses, thereby minimizing the apparent scale of the difference. If instead you use the goodness of fit score metric, which (as far as I understand) calculates the denominator as logmeanexp, and is interpreted as the difference between expected similarity for positive and negative pairs, the model never really generalizes to the test set and cannot accurately classify positive and negative pairs. Decoding performance for the test set is also poor. Consistency between train and test is relatively high though, so they are still consistent embeddings. Just not very useful embeddings if they can't generalize. But perhaps I'm doing something wrong? Do let me know if so.

The shuffling example above was meant as an illustrative worst case scenario akin to Simpson's paradox. And this overfitting scenario is exactly what I was concerned about. The encoders/embeddings don't appear to learn anything meaningful about the neural data.

Best,
Jake

Code is here: https://github.com/jgraving/cebra_shuffle_overfit/blob/main/cebra_overfit_hippocampus_pos.ipynb

Figures below:

[MMathisLab]

Hi @jgraving,

before digging deeper into this, the numbers you are sharing here are strictly worse than the ones in our demo notebook for decoding (https://cebra.ai/docs/demo_notebooks/Demo_decoding.html). Could you check this out, and let us know if that already contributes to solving your concern?

Please also note that consistency alone is not the metric to rely on. The loss, as we already pointed out, clearly matters here: 20K iterations overfits rat 1 in your own prior demo, and we do not use 20K in our paper. In Supplementary Fig. 7 (https://cebra.ai/docs/cebra-figures/figures/ExtendedDataFigure7.html), for instance, where we compare single and multi session embedding, we compare consistency at a fixed decoding performance:

The results in this notebook (on the held out test set) clearly demonstrate that decoding is possible on a held-out portion of the dataset (all of our reported decoding is on held out data). A comparison between InfoNCE and the decoding performance is also provided in the last cell of the notebook:

I would also note, infoNCE is the proper metric to use given that is our loss, and that is the only metric we used in our paper. We recently included GoF in the codebase to help ease interpretation of model fit, but again, unrelated to our paper (and GoF on the small validation data is noisy and not interpretable and we never suggest looking at this metric on validation unless with very large datasets).

[jgraving]

Hi @MMathisLab,

Yes I explored a range of training steps as you can clearly see in the plots. Results change very little after a couple thousand training steps. Your example uses 5000, which is included.

It seems your results are using median error not mean, which might explain the difference. Median effecitvely ignores the tail of the distribution, so higher magnitude errors are weighted less heavily.

And you argue goodness of fit is noisy for small datasets? But it is a bijection of infonce simply shifted by an additive constant

CEBRA/cebra/integrations/sklearn/metrics.py

Line 251 in 823c9ca

return (chance_level - infonce) * nats_to_bits

if goodness of fit is noisy then so is infonce, and so your argument is then that neither should be used for detecting overfitting?

[jgraving]

Hi @MMathisLab ,

The median error is indeed closer to your published/example results:

And, yes it works well only on part of the test set, as you mentioned:

It doesn't seem that infonce or goodness of fit are all that useful for detecting overfitting, especially when the task is decoding. But then brute forcing the hyperparameter search to find good decoding generalization is inefficient and requires a lot of effort (and compute) on the part of the user. Perhaps something to consider for future updates.

Either way, my results now match yours sufficiently well. Not quite the performance I expected, but it is what it is. So I will go ahead and close the issue.

Thanks for your time and efforts.

Best,
Jake

[MMathisLab]

Hi @jgraving,

We now had a look at the latest notebook you provided. There were multiple issues that deviated from the demo and methodology we discussed in the paper. Please have a look yourself at our modified version of your notebook: https://colab.research.google.com/drive/14olkCXowefz3cNRhI54vmFhCV2JP_f2B?usp=sharing .

Results are now on par with the decoding demo provided in the cebra documentation and we hope the comments in the notebook help. In short, in your notebook the train/valid was incorrect (a quick spot was the loss was high for both in your notebook); the decoding was also different.

Regarding the GoF/InfoNCE question: Yes, the metrics are essentially the same, just scaled and offset for ease of interpretation (see PR #190). GoF on validation should only be done if a proper split with sufficient number of samples was used (this was not the case in the demo you shared). We have a new notebook in development with best practices for future readers to look into.

Hope this clarifies any lingering concerns you had.

As a general request: As the discussion here does not concern issues around the CEBRA codebase, but rather questions around the correct usage, we also would like you to consider opening a reply in the discussion forum (https://github.com/AdaptiveMotorControlLab/CEBRA/discussions). In any case thanks for raising this, as a result we made some hopefully helpful comments in our usage guides and hope that future readers find this useful.

Regards,
@MMathisLab & @stes

[jgraving]

Thanks, that makes more sense. I was using the train test split you posted in your previous reply, but I guess this was just meant as a generic example.

Note in User Guide:

In our 9 Steps to work with CEBRA, the Step 3, already in our docs clearly highlights to users they can use a train/validation set:

3. Split data and labels

 (
     train_data,
     valid_data,
     train_discrete_label,
     valid_discrete_label,
     train_continuous_label,
     valid_continuous_label,
 ) = train_test_split(neural_data,
                     discrete_label,
                     continuous_label,
                     test_size=0.3)

Also in your quick start guide you import train test split from sklearn but then don't use it and split the last contiguous indices instead.

Seems my other error was decoding position only from a position direction embedding. Also makes sense to decode the same labels used for training and then calculate the error for position only.

[jgraving]

Just to confirm, the example you sent previously was using sklearn train test split so that's what I used. Seems you've now fixed this but didn't mention it...

https://github.com/AdaptiveMotorControlLab/CEBRA-demos/blob/cfd35addabc3be68a49252a2c62557b54abc1ee8/Demo_hippocampus.ipynb

[MMathisLab]

@jgraving actually no; the decoding demo notebook always has had the correct split (as I edited your notebook to have), and is what I linked above and asked you to use/investigate: https://github.com/AdaptiveMotorControlLab/CEBRA-demos

The usage doc indeed had a general sklearn split demo (which can be fine for some data/models), and on Saturday as part of some minor refactoring to take in your feedback, which I pointed out in my "lets be constructive" section, for the WIP notebook it had it ... the current merged and displayed demo does not have this.

Can I ask, what is the goal here? We have civilly engaged with your questions, spent a lot of time on this to address your concerns, took your feedback to update any clarifications, and spent hours going though and finding errors in your demo that is attempting to "retract our paper"... at this point, I respectfully ask we end here.