[myscience]

Auxiliary solver checkpoints omit reference_model. In DCL, that model is nn.Sequential(encoder, dynamics_model), so saving with the sklearn backend and loading a fresh estimator restores the encoder but leaves the dynamics freshly initialized. Embeddings can match after a round trip while predicted latent transitions change.

AuxiliaryVariableSolver now saves and restores reference_model alongside the base solver state. Both minibatch and full-batch DCL inherit this behavior, as do auxiliary solvers with independent reference encoders. The encoder remains model, preserving training and embedding-transform behavior.

Reproducer

import numpy as np
import torch
from cebra.integrations.sklearn.cebra import CEBRA, DCL

torch.manual_seed(4)
X = np.random.default_rng(4).normal(size=(100, 5)).astype("float32")
model = DCL(
    device="cpu", output_dimension=3, max_iterations=2,
    batch_size=8, batch_size_negatives=11,
).fit(X)
z = torch.randn(7, 3)
before = model.dynamics_model_(z).detach().clone()
model.save("dcl.pt", backend="sklearn")
restored = CEBRA.load("dcl.pt", backend="sklearn")

# The encoder round trip already works.
np.testing.assert_allclose(model.transform(X), restored.transform(X))
# Fails before this patch; passes with it.
torch.testing.assert_close(restored.dynamics_model_(z), before)

The failure was independently reproduced on upstream commit b7e4b9e.

Checkpoint compatibility

Historical auxiliary solver checkpoints lack reference_model. Strict loading raises an explanatory error. Explicit low-level solver.load_state_dict(state, strict=False) warns and restores available components; reference-model parameters absent from those components retain their current values. Solvers without auxiliary reference models are unchanged.

Validation

• Full required CPU suite: 3,193 passed, 102 skipped, zero failures; 6,965 dataset-dependent cases deselected.
• Regression coverage includes independent single-session and multisession reference encoders and their buffers, both DCL batch modes, linear and orthogonal-linear dynamics, shared encoder ownership, resumed training, and strict/non-strict loading. All six round-trip cases fail with the original serializer and pass with this revision.
• Full strict Sphinx build (-W --keep-going -n) passes with no warnings.
• Formatting/import ordering, spelling, docstring coverage, citation validation, binary checks, package build, and an installed-wheel checkpoint smoke test pass.
• Local environment: Python 3.10.12, PyTorch 2.5.1, scikit-learn 1.7.2, NumPy 2.2.6, CPU.

I have signed the Contributor License Agreement.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.

[stes]

This feels like a patch to an upstream bug; the dynamics model should be part of the state dict (?). Did we forget an ``nn.Module'' somewhere?

[myscience]

Yeah, I think that's a fair point. From what I understand the issue is that Solver has it's own state_dict method that saves self.model, while the DCL puts the dynamics module in self.reference_model = Sequential(model, dynamics), so the issue is that the custom saving never saves reference_model (DCL still needs model as a separate attribute to compute the embeddings and the pos/neg pairs etc...). I guess one clean option would be to add a custom state_dict and load_state_dict methods to AuxiliaryVariableSolver (which DCL inherits) that save and load reference_model. Otherwise one modifies the current method of the base Solver to handle the optional reference_model if present... what do you think @stes ?