Train/Test splits for the embedding itself

[wulfdewolf]

Are these not important for evaluating the learnt embedding?
In the documentation and the paper, I understand that all shown embeddings are always for the training data.

From my own experiments, even though the embedding for my training data looks very nice, when I project some hold-out data at the end of the session into the same space (without adapt=True), the embedding practically collapses.
Does that mean my original embedding was overfitted?

I might be missing something simple here. I understand that we are interested in the structure in the embedding space and if we can find one that has structure for the training data, that is very nice.
But, if it doesn't transfer to some unseen data from a not-so different distribution, should I trust it?
This is one example, but I've tried it for multiple sessions and other variables.

Train:

Test:

[MMathisLab]

Hi @wulfdewolf you'd have to tell us a bit more about your scenario for me to understand the images; test and train embeddings should look similar, if the data is similarly IID. You test here isn't collapsed, but I don't know the coloring. You'd also first have to establish the consistency of train runs - it should be very high and not overfitting.

[wulfdewolf]

Hi, apologies, I could've given a lot more information.
This is data where mice are seeing sequences of images.
The colouring indicates image position in the sequence, sequences are 25 images long.
I passed them as np.int32 such that the discrete solver is used.
The test data is very simply a small chunk at the start of the session + the same size chunk at the end of the session.
The train data is everything in between.

I used the following call:

model = CEBRA(
                model_architecture="offset10-model",
                batch_size=2048,
                learning_rate=3e-4,
                temperature_mode="auto",
                output_dimension=4,
                max_iterations=20000,
                num_hidden_units=64,
                distance="cosine",
                conditional="time_delta",
                device="cuda_if_available",
                verbose=True,
                time_offsets=10,
                optimizer_kwargs=(
                    ("betas", (0.9, 0.999)),
                    ("eps", 1e-08),
                    ("weight_decay", 0.0001),
                    ("amsgrad", False),
                ),
            )
model.fit(spike_counts[train_mask], behaviour[vars][train_mask].astype(np.int32))

I added some weight decay because I thought it might help improving the projection of the unseen data, but to no avail.

[wulfdewolf]

You're saying that I should try different seeds for the train embedding to verify if it is consistent, before looking at the projection of the test data?
I agree that my test projection here isn't entirely collapsed, but compared to the train data, it looks relatively collapsed.
The data comes from the same session, which is why I am so shocked it doesn't transfer...

[MMathisLab]

Correct; first you should be sure you have a consistent fit to pick the best parameters. Additionally, it sounds like your train vs test might not cover the same images; in that case, it doesn't seems like you are testing IID. If you take one full session and test in a held out session, how does that look?

[wulfdewolf]

Re the same images; I checked. Each image is shown very briefly, so the test set contains every image, multiple times.

Re consistency; what parameters would I play with if consistency isn't high?
Should I compute consistency across different seeds for the training data?
I could optimise hyperparameters by maximising consistency between train and test embeddings, but that feels like cheating?

All that aside, the train embedding I am seeing at the moment is that of my dreams, if I find it to be inconsistent, what does that mean?
Is the model too powerful and using small noise correlations to still get something like this?

[MMathisLab]

Hi @wulfdewolf sorry for my delay. To be clear, you should pick parameters that maximize consistency without overfitting, of course; in our user guide Step 3 is our recommendation for doing train/validation splits. The nature of your data (trial based) makes this easy to do. If you are over-parameterizing your model on small data (as you suggest might be the case/you train longer than typical, ie 20K is a lot), then that will become apparent in the train/val set up. Generally we recommend finding parameters that provide high consistency across runs for train/val and only then going on to use those parameters on the full data (much like in the way one uses tSNE or UMAP on the full dataset). I hope that clarifies.

You could consider checkpointing the model at regular intervals (e.g., every 1k steps) and then plot the consistency (on both train and test) along with the loss curves over time.

As a side comment, you could set the temperature mode to ‘constant’ and set the temperature to 0.1 (for a smoother embedding) or 1.0 (for a more “clustered” embedding).

[wulfdewolf]

Hi again, I have been working on different things and haven't had the time to look at this further...
I will at some point though.
From your comments, I'm assuming my first pass at this was way too naive and there's a bunch of things I wasn't doing that needed doing.
I wasn't looking at consistency at all...

Thanks for taking the time!!
I appreciate the work you all put into this.
Happy to keep this closed.

[MMathisLab]

Hi @wulfdewolf no problem; thanks for circling back around. Also, if you pick this back up, we made a new notebook that hopefully helps better explain some of the workflow suggestions we make! https://github.com/AdaptiveMotorControlLab/CEBRA-demos/blob/main/CEBRA_best_practices.ipynb