Skip to content

Add dynamics contrastive learning#301

Open
stes wants to merge 6 commits into
mainfrom
dcl-release
Open

Add dynamics contrastive learning#301
stes wants to merge 6 commits into
mainfrom
dcl-release

Conversation

@stes

@stes stes commented Jun 22, 2026

Copy link
Copy Markdown
Member

This PR add a basic version of dynamics contrastive learning to CEBRA, with support for linear dynamics. Additional dynamics models can be easily added once the API is fixed.

The minimal required changes are:

  • Add a new negative batch size. This is already tackled in Decouple batch size and number of negatives #263. The right order is to merge Decouple batch size and number of negatives #263, rebase this branch, and then merge.
  • Add a cebra.dynamics registry. Linear dynamics are added there as a first example, additional dynamics models can be ported as well. Some of these would need deeper changes, starting with linear is a pragmatic choice to get a first version merged soon
  • A change to the CEBRA API, to allow addition of a dynamics model. **An open Q is how we handle "good defaults": For CEBRA without a dynamics model, we would set some parameters different than with a dynamics model; right now I added a DCL subclass to CEBRA which has these defaults for initialization; tbd.

Out of scope for this PR

  • support for multi trials (needs API change to batching, can be follow up)
  • support for an explicit control signal (needs API change to batching, can be follow up)

On top of dynamics learning in "time contrastive mode", this integration of now also allows to combine fitting linear models in other sampling modes (e.g. using continuous or even discrete labels to fit "linear actions").

Two example scripts are added to produce an output. We can discuss adding a demo notebook to reproduce a plot like this as an example on the rat-hippocampus dataset:

image

The two-minute example run to test the API should produce a result like this (colored by position only like in the demo), visualized is a 2D ICA projection of the embedding:

image

Code: https://gist.github.com/stes/b1be59683b8e0aa37150ab528c4033d0

Co-authored-by: Tobias Schmidt <tsbaunatal@gmail.com>
Co-authored-by: Rodrigo Gonzalez Laiz <gonlairo@gmail.com>
@stes stes self-assigned this Jun 22, 2026
@cla-bot cla-bot Bot added the CLA signed label Jun 22, 2026
@stes

stes commented Jun 22, 2026

Copy link
Copy Markdown
Member Author

Unrelated issue, looks like some upstream software update again:

image

Should be tackled in a separate PR first

cause identified: #303

@stes stes requested a review from MMathisLab June 22, 2026 21:43
@stes

stes commented Jun 22, 2026

Copy link
Copy Markdown
Member Author

Issue related to the PR, most likely model solving/loading and backward compat

image

@stes

stes commented Jun 22, 2026

Copy link
Copy Markdown
Member Author

pickle error should be fixed; was due to PR not being compatible with #287 yet. hopefully some progress on tests now

@stes

stes commented Jun 22, 2026

Copy link
Copy Markdown
Member Author

pin numpy to below 2.5 temporarily until all changes there propagated

@stes

stes commented Jun 22, 2026

Copy link
Copy Markdown
Member Author

docs build green finally!

@MMathisLab MMathisLab requested a review from myscience June 23, 2026 03:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant