cca_zoo.deep¶

Deep CCA variants. Requires pip install cca-zoo[deep].

Base class¶

BaseDeep ¶

BaseDeep(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
)

Bases: LightningModule

Base class for deep multiview models, as Lightning modules.

Train with a :class:lightning.pytorch.Trainer. Calling the model, or trainer.predict, returns each view's encoding; for canonical variates, fit :class:~cca_zoo.linear.CCA or :class:~cca_zoo.linear.MCCA to the training encodings. Subclasses implement :meth:loss.

The covariance-based losses are estimated from each process's batch, so under multi-GPU (DDP) training they see the per-GPU batch, not the global one.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view, each mapping it to `n_components` outputs.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

Source code in cca_zoo/deep/_base.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
) -> None:
    super().__init__()
    # Modules are not hyperparameters: they are saved in the state dict and
    # passed back to load_from_checkpoint.
    self.save_hyperparameters(
        ignore=[
            name
            for cls in type(self).__mro__
            if "__init__" in vars(cls)
            for name, param in inspect.signature(cls.__init__).parameters.items()
            if "Module" in str(param.annotation)
        ]
    )
    self.encoders = nn.ModuleList(encoders)
    self.n_components = n_components
    self.learning_rate = learning_rate

forward ¶

forward(views: list[Tensor]) -> list[torch.Tensor]

Encode each view; one tensor of shape (batch_size, n_components) per view.

Source code in cca_zoo/deep/_base.py

def forward(self, views: list[torch.Tensor]) -> list[torch.Tensor]:
    """Encode each view; one tensor of shape (batch_size, n_components) per view."""
    return [self._checked(enc(v)) for enc, v in zip(self.encoders, views)]

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors, and optionally `"independent_views"` from an independent batch.	required

Returns:

Type	Description
`dict[str, Tensor]`	A dictionary whose `"objective"` entry is minimised.

Source code in cca_zoo/deep/_base.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors, and optionally
            ``"independent_views"`` from an independent batch.

    Returns:
        A dictionary whose ``"objective"`` entry is minimised.
    """
    raise NotImplementedError

predict_step ¶

predict_step(
    batch: Batch, batch_idx: int
) -> list[torch.Tensor]

Encodings of a batch, one tensor per view.

Source code in cca_zoo/deep/_base.py

def predict_step(self, batch: Batch, batch_idx: int) -> list[torch.Tensor]:
    """Encodings of a batch, one tensor per view."""
    return self.forward(batch["views"])

configure_optimizers ¶

configure_optimizers() -> torch.optim.Optimizer

Adam with learning_rate.

Source code in cca_zoo/deep/_base.py

def configure_optimizers(self) -> torch.optim.Optimizer:
    """Adam with ``learning_rate``."""
    return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

Data¶

MultiviewDataset ¶

MultiviewDataset(
    views: list[ArrayLike],
    partials: ArrayLike | None = None,
)

Bases: Dataset[dict[str, Any]]

Dataset of in-memory views, yielding {"views": [...]} samples.

The deep models read batches as dictionaries with a "views" list, and :class:DPCCA also a "partials" tensor; any dataset returning that shape works.

Parameters:

Name	Type	Description	Default
`views`	`list[ArrayLike]`	Arrays of shape (n_samples, n_features_i), one per view.	required
`partials`	`ArrayLike \| None`	Conditioning variable for :class:`DPCCA`, shape (n_samples, n_partials). Default is None.	`None`

Examples:

>>> import numpy as np
>>> from torch.utils.data import DataLoader
>>> from cca_zoo.deep import MultiviewDataset
>>> rng = np.random.default_rng(0)
>>> X1 = rng.standard_normal((100, 10)).astype("float32")
>>> X2 = rng.standard_normal((100, 8)).astype("float32")
>>> batch = next(iter(DataLoader(MultiviewDataset([X1, X2]), batch_size=32)))
>>> [v.shape for v in batch["views"]]
[torch.Size([32, 10]), torch.Size([32, 8])]

Source code in cca_zoo/deep/_data.py

def __init__(
    self, views: list[ArrayLike], partials: ArrayLike | None = None
) -> None:
    self.views: list[torch.Tensor] = [
        torch.as_tensor(np.asarray(v), dtype=torch.float32) for v in views
    ]
    self.partials = (
        None
        if partials is None
        else torch.as_tensor(np.asarray(partials), dtype=torch.float32)
    )

len ¶

__len__() -> int

Number of samples.

Source code in cca_zoo/deep/_data.py

def __len__(self) -> int:
    """Number of samples."""
    return int(self.views[0].shape[0])

getitem ¶

__getitem__(index: int) -> dict[str, Any]

Sample index as {"views": [tensor per view]}, with its partials.

Source code in cca_zoo/deep/_data.py

def __getitem__(self, index: int) -> dict[str, Any]:
    """Sample ``index`` as ``{"views": [tensor per view]}``, with its partials."""
    item: dict[str, Any] = {"views": [v[index] for v in self.views]}
    if self.partials is not None:
        item["partials"] = self.partials[index]
    return item

Models¶

DCCA ¶

DCCA(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: _ObjectiveModel

Deep CCA of two views: encoders trained to maximise canonical correlation.

Minimises, per mini-batch,

\[ \mathcal{L} = -\bigl\| \Sigma_{11}^{-1/2} \Sigma_{12} \Sigma_{22}^{-1/2} \bigr\|_F^2 \]

(:class:~cca_zoo.deep.objectives.CCALoss). As linear CCA generalises to MCCA, GCCA and TCCA, it generalises to more views as :class:DMCCA, :class:DGCCA and :class:DTCCA.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

Raises:

Type	Description
`ValueError`	If there are not two encoders.

References

Andrew, G., Arora, R., Bilmes, J., & Livescu, K. (2013). Deep canonical correlation analysis. ICML.

Examples:

>>> import lightning.pytorch as pl
>>> import numpy as np
>>> import torch
>>> import torch.nn as nn
>>> from torch.utils.data import DataLoader
>>> from cca_zoo.deep import DCCA, MultiviewDataset
>>> rng = np.random.default_rng(0)
>>> X1 = rng.standard_normal((64, 10)).astype("float32")
>>> X2 = rng.standard_normal((64, 8)).astype("float32")
>>> loader = DataLoader(MultiviewDataset([X1, X2]), batch_size=32)
>>> model = DCCA(n_components=2, encoders=[nn.Linear(10, 2), nn.Linear(8, 2)])
>>> trainer = pl.Trainer(max_epochs=2, logger=False, enable_progress_bar=False,
...                      enable_checkpointing=False, enable_model_summary=False)
>>> trainer.fit(model, loader)
>>> batches = trainer.predict(model, loader)
>>> Z1, Z2 = (torch.cat(z).numpy() for z in zip(*batches))
>>> from cca_zoo.linear import CCA
>>> U1, U2 = CCA(n_components=2).fit([Z1, Z2]).transform([Z1, Z2])
>>> U1.shape
(64, 2)

Source code in cca_zoo/deep/_dcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    if len(encoders) != 2:
        raise ValueError(
            f"DCCA is defined for two views, got {len(encoders)}; use DMCCA, "
            "DGCCA or DTCCA for more."
        )
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.objective = CCALoss(reg_covar=reg_covar)

DCCAEY ¶

DCCAEY(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Deep CCA by minimising the Eckart-Young loss.

\[ \mathcal{L}_{EY} = -2 \operatorname{tr}(C) + \operatorname{tr}(V V), \]

(:mod:cca_zoo._utils._ey), which is unconstrained and so suits mini-batches. With independent representations the penalty is \(\operatorname{tr}(V V_{\text{ind}})\), an unbiased estimate.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Chapman, J., Wells, L., & Lawry Aguila, A. (2024). Unconstrained Stochastic CCA: Unifying Multiview and Self-Supervised Learning. arXiv:2310.01012.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DCCAEY
>>> model = DCCAEY(n_components=4, encoders=[nn.Linear(10, 4), nn.Linear(8, 4)])

Source code in cca_zoo/deep/_base.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
) -> None:
    super().__init__()
    # Modules are not hyperparameters: they are saved in the state dict and
    # passed back to load_from_checkpoint.
    self.save_hyperparameters(
        ignore=[
            name
            for cls in type(self).__mro__
            if "__init__" in vars(cls)
            for name, param in inspect.signature(cls.__init__).parameters.items()
            if "Module" in str(param.annotation)
        ]
    )
    self.encoders = nn.ModuleList(encoders)
    self.n_components = n_components
    self.learning_rate = learning_rate

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The EY loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors and, optionally, `"independent_views"` from an independent batch, which give an unbiased estimate of the penalty.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "rewards", "penalties"}`.

Source code in cca_zoo/deep/_dcca_ey.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The EY loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors and, optionally,
            ``"independent_views"`` from an independent batch, which
            give an unbiased estimate of the penalty.

    Returns:
        ``{"objective", "rewards", "penalties"}``.
    """
    c, v = _cca_cv(self(batch["views"]))
    independent = batch.get("independent_views")
    v_ind = v if independent is None else _cca_cv(self(independent))[1]
    rewards = torch.trace(2.0 * c)
    penalties = torch.trace(v @ v_ind)
    return {
        "objective": -rewards + penalties,
        "rewards": rewards,
        "penalties": penalties,
    }

DCCANOI ¶

DCCANOI(
    n_components: int,
    encoders: list[Module],
    rho: float = 0.9,
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: BaseDeep

Deep CCA by nonlinear orthogonal iterations.

Wang et al.'s Algorithm 2: regresses each view's encoding on the others' whitened, centred encodings, held fixed,

\[ \mathcal{L} = \sum_{i \neq j} \bigl\| z_i - \operatorname{sg}\bigl(\Sigma_j^{-1/2} (z_j - \mu_j)\bigr) \bigr\|_2^2, \]

with sg a stop-gradient, and the running moments updated by each training batch as \(\Sigma_j \leftarrow \rho \Sigma_j + (1 - \rho)\, \widehat{\Sigma}_j\), and \(\mu_j\) likewise.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`rho`	`float`	Weight of the previous running moments in `[0, 1)`, the paper's time constant. Default is 0.9.	`0.9`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

Raises:

Type	Description
`ValueError`	If `rho` is outside `[0, 1)`.

References

Wang, W., Arora, R., Livescu, K., & Srebro, N. (2015). Stochastic optimization for deep CCA via nonlinear orthogonal iterations. Allerton.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DCCANOI
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4)]
>>> model = DCCANOI(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dcca_noi.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    rho: float = 0.9,
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    if not 0.0 <= rho < 1.0:
        raise ValueError(f"rho must be in [0, 1), got {rho}.")
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.rho = rho
    self.mse = nn.MSELoss(reduction="sum")
    self.bws = nn.ModuleList(
        [_BatchWhiten(n_components, rho=rho, reg_covar=reg_covar) for _ in encoders]
    )

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The NOI loss of a batch.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective": loss}`.

Source code in cca_zoo/deep/_dcca_noi.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The NOI loss of a batch.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective": loss}``.
    """
    representations = self(batch["views"])
    targets = [bw(z).detach() for z, bw in zip(representations, self.bws)]
    objective = sum(
        self.mse(representations[i], targets[j])
        for i, j in itertools.permutations(range(len(representations)), 2)
    )
    return {"objective": objective}

DCCASDL ¶

DCCASDL(
    n_components: int,
    encoders: list[Module],
    lam: float = 0.5,
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Deep CCA with a stochastic decorrelation loss.

Batch-normalised encodings are aligned by mean squared error, summed over pairs of views, and decorrelated within each view:

\[ \mathcal{L} = \sum_{a < b} \operatorname{MSE}(z_a, z_b) + \lambda \sum_a \operatorname{mean}\lvert \text{offdiag}(\operatorname{Cov}(z_a)) \rvert. \]

With two views this is the original loss, with one difference: Chang et al. estimate the covariance by a running average across batches, while here each batch's covariance is used, so no state is carried between batches, as in :class:DPCCA. Use batches well above n_components.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`lam`	`float`	Weight of the decorrelation term. Default is 0.5.	`0.5`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Chang, X., Xiang, T., & Hospedales, T. M. (2018). Scalable and effective deep CCA via soft decorrelation. CVPR.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DCCASDL
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4)]
>>> model = DCCASDL(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dcca_sdl.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    lam: float = 0.5,
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.lam = lam
    self.bns = nn.ModuleList(
        [nn.BatchNorm1d(n_components, affine=False) for _ in encoders]
    )

forward ¶

forward(views: list[Tensor]) -> list[torch.Tensor]

Batch-normalised encodings of each view.

Source code in cca_zoo/deep/_dcca_sdl.py

def forward(self, views: list[torch.Tensor]) -> list[torch.Tensor]:
    """Batch-normalised encodings of each view."""
    return [bn(enc(v)) for enc, bn, v in zip(self.encoders, self.bns, views)]

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The SDL loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "l2", "sdl"}`.

Source code in cca_zoo/deep/_dcca_sdl.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The SDL loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "l2", "sdl"}``.
    """
    representations = self(batch["views"])
    l2 = torch.stack(
        [F.mse_loss(a, b) for a, b in itertools.combinations(representations, 2)]
    ).sum()
    sdl = torch.stack([_sdl_loss(r) for r in representations]).sum()
    return {
        "objective": l2 + self.lam * sdl,
        "l2": l2,
        "sdl": sdl,
    }

DCCAE ¶

DCCAE(
    n_components: int,
    encoders: list[Module],
    decoders: list[Module],
    lam: float = 0.5,
    objective: Module | None = None,
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: _ObjectiveModel

Deep CCA with per-view autoencoder reconstruction.

\[ \mathcal{L} = (1 - \lambda) \mathcal{L}_{\text{CCA}}(z_1, \dots, z_M) + \lambda \sum_i \operatorname{MSE}(x_i, \text{decoder}_i(z_i)). \]

lam=0 is :class:DMCCA (:class:DCCA with two views) and lam=1 independent autoencoders.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`decoders`	`list[Module]`	One module per view mapping its encoding back.	required
`lam`	`float`	Reconstruction weight in `[0, 1]`. Default is 0.5.	`0.5`
`objective`	`Module \| None`	Correlation loss; None uses `MCCALoss(reg_covar)`, the sum of `CCALoss` over pairs of views. Default is None.	`None`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

Raises:

Type	Description
`ValueError`	If `lam` is outside `[0, 1]`.

References

Wang, W., Arora, R., Livescu, K., & Bilmes, J. (2015). On deep multi-view representation learning. ICML.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DCCAE
>>> model = DCCAE(
...     n_components=4,
...     encoders=[nn.Linear(10, 4), nn.Linear(8, 4)],
...     decoders=[nn.Linear(4, 10), nn.Linear(4, 8)],
... )

Source code in cca_zoo/deep/_dccae.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    decoders: list[nn.Module],
    lam: float = 0.5,
    objective: nn.Module | None = None,
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    if not 0.0 <= lam <= 1.0:
        raise ValueError(f"lam must be in [0, 1], got {lam}.")
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.objective = (
        MCCALoss(reg_covar=reg_covar) if objective is None else objective
    )
    self.lam = lam
    self.decoders = nn.ModuleList(decoders)

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The DCCAE loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "cca", "reconstruction"}`.

Source code in cca_zoo/deep/_dccae.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The DCCAE loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "cca", "reconstruction"}``.
    """
    views = batch["views"]
    representations = self(views)
    cca = self.objective(representations)
    reconstruction = _reconstruction_loss(
        views, [dec(z) for dec, z in zip(self.decoders, representations)]
    )
    return {
        "objective": (1.0 - self.lam) * cca + self.lam * reconstruction,
        "cca": cca,
        "reconstruction": reconstruction,
    }

DVCCA ¶

DVCCA(
    n_components: int,
    encoder: Module,
    decoders: list[Module],
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Deep variational CCA: a latent inferred from one view generates all of them.

With prior \(z \sim \mathcal{N}(0, I)\), decoders \(p(x_i \mid z)\) for every view and an encoder \(q(z \mid x_1) = \mathcal{N}(\mu, \sigma^2)\) of the first view, training minimises the negative ELBO

\[ \mathcal{L} = \sum_i \tfrac12 \lVert x_i - \text{decoder}_i(z) \rVert^2 + \mathrm{KL}(q(z \mid x_1) \,\|\, \mathcal{N}(0, I)). \]

Unlike the other deep models there is one encoding, so trainer.predict returns a single array, the posterior mean \(\mu\).

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoder`	`Module`	Module mapping the first view to `2 * n_components` outputs, the mean then the log-variance.	required
`decoders`	`list[Module]`	One module per view mapping the latent back to that view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Wang, W., Yan, X., Lee, H., & Livescu, K. (2016). Deep variational canonical correlation analysis. arXiv:1610.03454.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DVCCA
>>> model = DVCCA(
...     n_components=4,
...     encoder=nn.Linear(10, 8),
...     decoders=[nn.Linear(4, 10), nn.Linear(4, 6)],
... )

Source code in cca_zoo/deep/_dvcca.py

def __init__(
    self,
    n_components: int,
    encoder: nn.Module,
    decoders: list[nn.Module],
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=[encoder],
        learning_rate=learning_rate,
    )
    self.decoders = nn.ModuleList(decoders)

forward ¶

forward(views: list[Tensor]) -> list[torch.Tensor]

The posterior mean from the first view, as a one-element list.

Source code in cca_zoo/deep/_dvcca.py

def forward(self, views: list[torch.Tensor]) -> list[torch.Tensor]:
    """The posterior mean from the first view, as a one-element list."""
    return [self._posterior(views)[0]]

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The negative ELBO of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "reconstruction", "kl"}`.

Source code in cca_zoo/deep/_dvcca.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The negative ELBO of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "reconstruction", "kl"}``.
    """
    views = batch["views"]
    mu, log_var = self._posterior(views)
    z = _sample(mu, log_var)
    reconstruction = _gaussian_nll(views, [dec(z) for dec in self.decoders])
    kl = _kl(mu, log_var)
    return {
        "objective": reconstruction + kl,
        "reconstruction": reconstruction,
        "kl": kl,
    }

predict_step ¶

predict_step(
    batch: Batch, batch_idx: int
) -> list[torch.Tensor]

The posterior mean of a batch, as a one-element list.

Source code in cca_zoo/deep/_dvcca.py

def predict_step(self, batch: Batch, batch_idx: int) -> list[torch.Tensor]:
    """The posterior mean of a batch, as a one-element list."""
    return self.forward(batch["views"])

DVCCAPrivate ¶

DVCCAPrivate(
    n_components: int,
    encoder: Module,
    private_encoders: list[Module],
    decoders: list[Module],
    n_private: int,
    learning_rate: float = 0.001,
)

Bases: DVCCA

DVCCA with a private latent variable per view as well as the shared one.

The shared \(z\) is inferred from the first view, \(q(z \mid x_1)\), and each view's private \(h_i\) from that view, \(q(h_i \mid x_i)\); view \(i\) is decoded from \([z, h_i]\). All latents have standard normal priors, and training minimises the negative ELBO

\[ \mathcal{L} = \sum_i \tfrac12 \lVert x_i - \text{decoder}_i(z, h_i) \rVert^2 + \mathrm{KL}(q(z \mid x_1) \,\|\, p(z)) + \sum_i \mathrm{KL}(q(h_i \mid x_i) \,\|\, p(h_i)). \]

The private variables take up view-specific variation, leaving \(z\) the shared part. trainer.predict returns the posterior mean of \(z\); :meth:private_means gives each view's private posterior mean.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Dimension of the shared latent.	required
`encoder`	`Module`	Module mapping the first view to `2 * n_components` outputs, the mean then the log-variance of \(z\).	required
`private_encoders`	`list[Module]`	One module per view mapping it to `2 * n_private` outputs, the mean then the log-variance of its \(h_i\).	required
`decoders`	`list[Module]`	One module per view mapping `n_components + n_private` inputs, \([z, h_i]\), back to that view.	required
`n_private`	`int`	Dimension of each private latent.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Wang, W., Yan, X., Lee, H., & Livescu, K. (2016). Deep variational canonical correlation analysis. arXiv:1610.03454.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DVCCAPrivate
>>> model = DVCCAPrivate(
...     n_components=4,
...     encoder=nn.Linear(10, 8),
...     private_encoders=[nn.Linear(10, 4), nn.Linear(6, 4)],
...     decoders=[nn.Linear(6, 10), nn.Linear(6, 6)],
...     n_private=2,
... )

Source code in cca_zoo/deep/_dvcca.py

def __init__(
    self,
    n_components: int,
    encoder: nn.Module,
    private_encoders: list[nn.Module],
    decoders: list[nn.Module],
    n_private: int,
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoder=encoder,
        decoders=decoders,
        learning_rate=learning_rate,
    )
    self.private_encoders = nn.ModuleList(private_encoders)
    self.n_private = n_private

private_means ¶

private_means(views: list[Tensor]) -> list[torch.Tensor]

Each view's private posterior mean, shape (batch_size, n_private).

Source code in cca_zoo/deep/_dvcca.py

def private_means(self, views: list[torch.Tensor]) -> list[torch.Tensor]:
    """Each view's private posterior mean, shape (batch_size, n_private)."""
    return [mu for mu, _ in self._private_posteriors(views)]

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The negative ELBO of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "reconstruction", "kl", "private_kl"}`.

Source code in cca_zoo/deep/_dvcca.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The negative ELBO of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "reconstruction", "kl", "private_kl"}``.
    """
    views = batch["views"]
    mu, log_var = self._posterior(views)
    z = _sample(mu, log_var)
    private = self._private_posteriors(views)
    reconstruction = _gaussian_nll(
        views,
        [
            dec(torch.cat([z, _sample(m, lv)], dim=1))
            for dec, (m, lv) in zip(self.decoders, private)
        ],
    )
    kl = _kl(mu, log_var)
    private_kl = torch.stack([_kl(m, lv) for m, lv in private]).sum()
    return {
        "objective": reconstruction + kl + private_kl,
        "reconstruction": reconstruction,
        "kl": kl,
        "private_kl": private_kl,
    }

DTCCA ¶

DTCCA(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: _ObjectiveModel

Deep tensor CCA: maximise the cross-moment tensor of whitened encodings.

Minimises \(-\|M\|_F\) with \(M = \frac{1}{n} \sum_s H_1[s] \otimes \cdots \otimes H_M[s]\) for ridge-whitened encodings \(H_i\) (:class:~cca_zoo.deep.objectives.TCCALoss).

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Wong, H. S., Wang, L., Chan, R., & Zeng, T. (2021). Deep tensor CCA for multi-view learning. IEEE Transactions on Big Data.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DTCCA
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4), nn.Linear(6, 4)]
>>> model = DTCCA(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dtcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.objective = TCCALoss(reg_covar=reg_covar)

DMCCA ¶

DMCCA(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: _ObjectiveModel

Deep multiset CCA: the sum of pairwise deep CCA losses.

\[ \mathcal{L} = \sum_{i < j} \mathcal{L}_{\text{CCA}}(z_i, z_j) \]

(:class:~cca_zoo.deep.objectives.MCCALoss), for any number of views.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Somandepalli, K., Kumar, N., Travadi, R., & Narayanan, S. (2019). Multimodal representation learning using deep multiset canonical correlation analysis. arXiv:1904.01775.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DMCCA
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4), nn.Linear(6, 4)]
>>> model = DMCCA(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dmcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.objective = MCCALoss(reg_covar=reg_covar)

DGCCA ¶

DGCCA(
    n_components: int,
    encoders: list[Module],
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: _ObjectiveModel

Deep generalized CCA: correlate every view with a shared target.

\[ \mathcal{L} = -\sum_{d=1}^{k} \lambda_d\Bigl(\sum_i H_i H_i^\top\Bigr) \]

for ridge-whitened encodings \(H_i\) (:class:~cca_zoo.deep.objectives.GCCALoss), for any number of views.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Benton, A., Khayrallah, H., Gujral, B., Reisinger, D. A., Zhang, S., & Arora, R. (2019). Deep generalized canonical correlation analysis. RepL4NLP.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DGCCA
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4), nn.Linear(6, 4)]
>>> model = DGCCA(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dgcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.reg_covar = reg_covar
    self.objective = GCCALoss(reg_covar=reg_covar)

DPCCA ¶

DPCCA(
    n_components: int,
    encoders: list[Module],
    partial_encoder: Module | None = None,
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: BaseDeep

Deep partial CCA: correlate the views after conditioning on a shared variable.

The Eckart-Young loss of :class:DCCAEY on the views' partial covariances given a conditioning variable \(Z\) (batch["partials"]). For the stacked encodings \(F\) of a batch,

\[ \Sigma_{FF|Z} = \Sigma_{FF} - \Sigma_{FZ} \Sigma_{ZZ}^{-1} \Sigma_{ZF}, \qquad \mathcal{L} = -2 \operatorname{tr}(C_{|Z}) + \operatorname{tr}(V_{|Z} V_{|Z}), \]

with \(C_{|Z}\) the mean over views of the sum of each view's blocks of \(\Sigma_{FF|Z}\), and \(V_{|Z}\) the mean of its diagonal blocks, as :class:DCCAEY forms them from the covariances.

partial_encoder=None uses \(Z\) as given; a module encodes it, trained to explain the encodings by least squares, so that partialling removes all the encoded \(Z\) can explain. \(Z\) is needed only for training: prediction encodes the views alone.

Rotman et al. train by nonlinear orthogonal iterations, regressing each view onto the other's whitened partialled encoding, with running estimates of the covariances across batches; the EY loss needs neither whitening nor running estimates, and suits mini-batch training (Chapman et al., 2024). Their encoded variant also trains the partial encoder on the correlation loss, which rewards it for failing to explain the confound; here it minimises the partialling residual instead. With two views the model is theirs; more views are handled as in :class:DCCAEY.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`partial_encoder`	`Module \| None`	Module encoding the conditioning variable, or None to use it as given. Default is None.	`None`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

Raises:

Type	Description
`ValueError`	If a training batch has no `"partials"`.

References

Rotman, G., Vulić, I., & Reichart, R. (2018). Bridging languages through images with deep partial canonical correlation analysis. ACL.

Chapman, J., Wells, L., & Lawry Aguila, A. (2024). Unconstrained Stochastic CCA: Unifying Multiview and Self-Supervised Learning. arXiv:2310.01012.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import DPCCA
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4)]
>>> model = DPCCA(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_dpcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    partial_encoder: nn.Module | None = None,
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.partial_encoder = partial_encoder
    self.reg_covar = reg_covar

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The EY loss of the encodings' partial covariances and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors and the conditioning variable under `"partials"`.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "rewards", "penalties"}`, and `"residual"`, the
`dict[str, Tensor]`	partial encoder's loss, when there is one.

Source code in cca_zoo/deep/_dpcca.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The EY loss of the encodings' partial covariances and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors and the
            conditioning variable under ``"partials"``.

    Returns:
        ``{"objective", "rewards", "penalties"}``, and ``"residual"``, the
        partial encoder's loss, when there is one.
    """
    z = self._conditioning(batch)
    f = torch.cat([f - f.mean(dim=0) for f in self(batch["views"])], dim=1)
    n, m, k = f.shape[0], len(batch["views"]), self.n_components
    eye = self.reg_covar * torch.eye(z.shape[1], device=z.device, dtype=z.dtype)
    # The partial covariance of the encodings given Z, held fixed.
    z_fixed = z.detach()
    zf = z_fixed.T @ f
    partial = (
        f.T @ f - zf.T @ torch.linalg.solve(z_fixed.T @ z_fixed + eye, zf)
    ) / (n - 1)
    blocks = partial.reshape(m, k, m, k)
    c = blocks.sum(dim=(0, 2)) / m
    v = torch.diagonal(blocks, dim1=0, dim2=2).sum(dim=-1) / m
    rewards = torch.trace(2.0 * c)
    penalties = torch.trace(v @ v)
    terms = {"rewards": rewards, "penalties": penalties}
    objective = -rewards + penalties
    if self.partial_encoder is not None:
        # The partial encoder learns to explain the encodings, held fixed, so
        # partialling removes all it can; trained on the EY loss it would
        # instead learn to leave the confound in. Its loss is the mean
        # squared residual, the trace of the partial second moment.
        f_fixed = f.detach()
        zf = z.T @ f_fixed
        residual = torch.trace(
            f_fixed.T @ f_fixed - zf.T @ torch.linalg.solve(z.T @ z + eye, zf)
        ) / (n * k)
        terms["residual"] = residual
        objective = objective + residual
    return {"objective": objective, **terms}

SplitAE ¶

SplitAE(
    n_components: int,
    encoders: list[Module],
    decoders: list[Module],
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Split autoencoder: every view reconstructed from all views' encodings.

\[ \mathcal{L} = \sum_i \operatorname{MSE}(x_i, \text{dec}_i(z_1 \| \cdots \| z_M)). \]

A reconstruction baseline that does not maximise correlation.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Dimension of each encoding.	required
`encoders`	`list[Module]`	One module per view.	required
`decoders`	`list[Module]`	One module per view, taking `n_views * n_components` inputs.	required
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Wang, W., Arora, R., Livescu, K., & Bilmes, J. (2015). On deep multi-view representation learning. ICML.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import SplitAE
>>> model = SplitAE(
...     n_components=4,
...     encoders=[nn.Linear(10, 4), nn.Linear(8, 4)],
...     decoders=[nn.Linear(8, 10), nn.Linear(8, 8)],
... )

Source code in cca_zoo/deep/_splitae.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    decoders: list[nn.Module],
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.decoders = nn.ModuleList(decoders)

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The reconstruction loss of a batch.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective": loss}`.

Source code in cca_zoo/deep/_splitae.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The reconstruction loss of a batch.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective": loss}``.
    """
    views = batch["views"]
    joint = torch.cat(self(views), dim=1)
    return {
        "objective": _reconstruction_loss(
            views, [dec(joint) for dec in self.decoders]
        )
    }

BarlowTwins ¶

BarlowTwins(
    n_components: int,
    encoders: list[Module],
    lam: float = 0.005,
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Barlow Twins: drive the cross-correlation between views to the identity.

\[ \mathcal{L} = \sum_{a < b} \Bigl( \sum_i (1 - C^{ab}_{ii})^2 + \lambda \sum_{i \neq j} (C^{ab}_{ij})^2 \Bigr), \qquad C^{ab} = \tfrac{1}{n} Z_a^\top Z_b, \]

for batch-normalised encodings \(Z_a\), summed over pairs of views; with two views this is the original loss.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`lam`	`float`	Weight of the off-diagonal term. Default is 5e-3.	`0.005`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Zbontar, J., Jing, L., Misra, I., LeCun, Y., & Deny, S. (2021). Barlow twins: Self-supervised learning via redundancy reduction. ICML.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import BarlowTwins
>>> encoders = [nn.Linear(10, 4), nn.Linear(8, 4)]
>>> model = BarlowTwins(n_components=4, encoders=encoders)

Source code in cca_zoo/deep/_barlowtwins.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    lam: float = 5e-3,
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.lam = lam
    self.bns = nn.ModuleList(
        [nn.BatchNorm1d(n_components, affine=False) for _ in encoders]
    )

forward ¶

forward(views: list[Tensor]) -> list[torch.Tensor]

Batch-normalised encodings of each view.

Source code in cca_zoo/deep/_barlowtwins.py

def forward(self, views: list[torch.Tensor]) -> list[torch.Tensor]:
    """Batch-normalised encodings of each view."""
    return [bn(enc(v)) for enc, bn, v in zip(self.encoders, self.bns, views)]

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The Barlow Twins loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "invariance", "redundancy"}`.

Source code in cca_zoo/deep/_barlowtwins.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The Barlow Twins loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "invariance", "redundancy"}``.
    """
    representations = self(batch["views"])
    n, k = representations[0].shape
    off_diagonal = ~torch.eye(k, dtype=torch.bool, device=representations[0].device)
    invariance = redundancy = torch.zeros((), device=representations[0].device)
    for z1, z2 in itertools.combinations(representations, 2):
        cross = z1.T @ z2 / n
        invariance = invariance + torch.sum((1.0 - torch.diag(cross)) ** 2)
        redundancy = redundancy + torch.sum(cross[off_diagonal] ** 2)
    objective = invariance + self.lam * redundancy
    return {
        "objective": objective,
        "invariance": invariance,
        "redundancy": redundancy,
    }

VICReg ¶

VICReg(
    n_components: int,
    encoders: list[Module],
    sim_coeff: float = 25.0,
    std_coeff: float = 25.0,
    cov_coeff: float = 1.0,
    learning_rate: float = 0.001,
)

Bases: BaseDeep

Variance-invariance-covariance regularisation.

\[ \mathcal{L} = \gamma \sum_{a < b} \operatorname{MSE}(z_a, z_b) + \mu \sum_a \operatorname{mean}(\max(0, 1 - \sigma(z_a))) + \nu \sum_a \frac{1}{d} \sum_{i \neq j} \operatorname{Cov}(z_a)_{ij}^2, \]

with \(\gamma, \mu, \nu\) = sim_coeff, std_coeff, cov_coeff. The invariance term is summed over pairs of views; with two views this is the original loss.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`sim_coeff`	`float`	Weight of the invariance term. Default is 25.0.	`25.0`
`std_coeff`	`float`	Weight of the variance term. Default is 25.0.	`25.0`
`cov_coeff`	`float`	Weight of the covariance term. Default is 1.0.	`1.0`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`

References

Bardes, A., Ponce, J., & LeCun, Y. (2022). VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning. ICLR.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import VICReg
>>> model = VICReg(n_components=4, encoders=[nn.Linear(10, 4), nn.Linear(8, 4)])

Source code in cca_zoo/deep/_vicreg.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    sim_coeff: float = 25.0,
    std_coeff: float = 25.0,
    cov_coeff: float = 1.0,
    learning_rate: float = 1e-3,
) -> None:
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.sim_coeff = sim_coeff
    self.std_coeff = std_coeff
    self.cov_coeff = cov_coeff

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The VICReg loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "sim_loss", "var_loss", "cov_loss"}`.

Source code in cca_zoo/deep/_vicreg.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The VICReg loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "sim_loss", "var_loss", "cov_loss"}``.
    """
    representations = self(batch["views"])
    sim = torch.stack(
        [F.mse_loss(a, b) for a, b in itertools.combinations(representations, 2)]
    ).sum()
    var = torch.stack([_variance_loss(z) for z in representations]).sum()
    cov = torch.stack([_covariance_loss(z) for z in representations]).sum()
    objective = self.sim_coeff * sim + self.std_coeff * var + self.cov_coeff * cov
    return {
        "objective": objective,
        "sim_loss": sim,
        "var_loss": var,
        "cov_loss": cov,
    }

LeJEPA ¶

LeJEPA(
    n_components: int,
    encoders: list[Module],
    lam: float = 0.05,
    n_slices: int = 256,
    n_points: int = 17,
    t_max: float = 5.0,
    learning_rate: float = 0.001,
    unbiased: bool = False,
)

Bases: BaseDeep

LeJEPA: every view predicts the views' centre, and SIGReg prevents collapse.

\[ \mathcal{L} = (1 - \lambda)\,\frac{1}{K n d} \sum_{k=1}^{K} \lVert z_k - \bar z \rVert_F^2 + \lambda\,\frac{1}{K} \sum_{k=1}^{K} \operatorname{SIGReg}(z_k), \qquad \bar z = \frac{1}{K} \sum_{k=1}^{K} z_k, \]

where SIGReg is the Epps-Pulley test statistic of \(z_k\)'s projections on n_slices random unit directions against a standard normal, averaged over directions, which are resampled at every training step. Isotropic Gaussian embeddings, which SIGReg favours, cannot collapse, so no stop-gradient, predictor or teacher network is needed. Every view is a global view, so the centre is the mean over all views.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view; two or more.	required
`lam`	`float`	Weight \(\lambda \in [0, 1]\) of SIGReg against the predictive term. Default is 0.05.	`0.05`
`n_slices`	`int`	Random directions per step. Default is 256.	`256`
`n_points`	`int`	Quadrature points of the Epps-Pulley integral. Default is 17.	`17`
`t_max`	`float`	Half-width of the integration interval. Default is 5.0.	`5.0`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`unbiased`	`bool`	Whether SIGReg uses the unbiased estimate of the squared characteristic-function error (see :func:`_sigreg`) in place of the plug-in one, which is biased upward by about 1 for a standard-normal batch. Default is False, the reference implementation's.	`False`

Raises:

Type	Description
`ValueError`	If `lam` is outside `[0, 1]`.

References

Balestriero, R., & LeCun, Y. (2025). LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics. arXiv:2511.08544.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import LeJEPA
>>> model = LeJEPA(n_components=4, encoders=[nn.Linear(10, 4), nn.Linear(8, 4)])

Source code in cca_zoo/deep/_lejepa.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    lam: float = 0.05,
    n_slices: int = 256,
    n_points: int = 17,
    t_max: float = 5.0,
    learning_rate: float = 1e-3,
    unbiased: bool = False,
) -> None:
    if not 0.0 <= lam <= 1.0:
        raise ValueError(f"lam must be in [0, 1], got {lam}.")
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.lam = lam
    self.n_slices = n_slices
    self.n_points = n_points
    self.t_max = t_max
    self.unbiased = unbiased

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The LeJEPA loss of a batch and its terms.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "sim_loss", "sigreg"}`.

Source code in cca_zoo/deep/_lejepa.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The LeJEPA loss of a batch and its terms.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "sim_loss", "sigreg"}``.
    """
    representations = self(batch["views"])
    z = torch.stack(representations)  # (views, n, d)
    sim = (z.mean(dim=0) - z).square().mean()
    sigreg = torch.stack(
        [
            _sigreg(
                zk,
                self.global_step,
                self.n_slices,
                self.n_points,
                self.t_max,
                self.unbiased,
            )
            for zk in representations
        ]
    ).mean()
    return {
        "objective": (1 - self.lam) * sim + self.lam * sigreg,
        "sim_loss": sim,
        "sigreg": sigreg,
    }

NRDCCA ¶

NRDCCA(
    n_components: int,
    encoders: list[Module],
    lam: float = 1.0,
    learning_rate: float = 0.001,
    reg_covar: float = 1e-06,
)

Bases: BaseDeep

Deep CCA with noise regularisation, against model collapse.

Linear CCA's correlation of a view with independent noise does not change under an invertible linear map, \(\operatorname{Corr}(X_k, A_k) = \operatorname{Corr}(W_k X_k, W_k A_k)\); an overfitted network's does. Each batch draws Gaussian noise \(A_k\) of each view's shape, and the loss

\[ \mathcal{L} = \mathcal{L}_{DCCA}\bigl(f_1(X_1), \dots, f_K(X_K)\bigr) + \lambda \sum_{k=1}^{K} \bigl\lvert \operatorname{Corr}(f_k(X_k), f_k(A_k)) - \operatorname{Corr}(X_k, A_k) \bigr\rvert \]

holds the encoders to that property. \(\mathcal{L}_{DCCA}\) is :class:~cca_zoo.deep.objectives.MCCALoss, :class:DCCA's loss summed over pairs of views, and \(\operatorname{Corr}\) is the mean canonical correlation, which compares a view's \(d_k\) canonical correlations with its encoding's. With lam=0 this is :class:DMCCA.

Parameters:

Name	Type	Description	Default
`n_components`	`int`	Number of latent dimensions.	required
`encoders`	`list[Module]`	One module per view.	required
`lam`	`float`	Weight of the noise regularisation. Default is 1.0.	`1.0`
`learning_rate`	`float`	Adam learning rate. Default is 1e-3.	`0.001`
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

Raises:

Type	Description
`ValueError`	If `lam` is negative.

References

He, J., Du, J., Xu, S., & Ma, W. (2024). Preventing model collapse in deep canonical correlation analysis by noise regularization. NeurIPS.

Examples:

>>> import torch.nn as nn
>>> from cca_zoo.deep import NRDCCA
>>> model = NRDCCA(n_components=4, encoders=[nn.Linear(10, 4), nn.Linear(8, 4)])

Source code in cca_zoo/deep/_nrdcca.py

def __init__(
    self,
    n_components: int,
    encoders: list[nn.Module],
    lam: float = 1.0,
    learning_rate: float = 1e-3,
    reg_covar: float = 1e-6,
) -> None:
    if lam < 0.0:
        raise ValueError(f"lam must be non-negative, got {lam}.")
    super().__init__(
        n_components=n_components,
        encoders=encoders,
        learning_rate=learning_rate,
    )
    self.lam = lam
    self.reg_covar = reg_covar
    self.objective = MCCALoss(reg_covar=reg_covar)

loss ¶

loss(batch: Batch) -> dict[str, torch.Tensor]

The DCCA loss and the noise regularisation of a batch.

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	Dictionary with a `"views"` list of tensors.	required

Returns:

Type	Description
`dict[str, Tensor]`	`{"objective", "dcca", "noise_regularisation"}`.

Source code in cca_zoo/deep/_nrdcca.py

def loss(self, batch: Batch) -> dict[str, torch.Tensor]:
    """The DCCA loss and the noise regularisation of a batch.

    Args:
        batch: Dictionary with a ``"views"`` list of tensors.

    Returns:
        ``{"objective", "dcca", "noise_regularisation"}``.
    """
    views = batch["views"]
    noise = [torch.randn_like(v) for v in views]
    encodings, noise_encodings = self(views), self(noise)
    regularisation = torch.stack(
        [
            torch.abs(
                _mean_canonical_correlation(f, g, self.reg_covar)
                - _mean_canonical_correlation(x, a, self.reg_covar)
            )
            for x, a, f, g in zip(views, noise, encodings, noise_encodings)
        ]
    ).sum()
    dcca = self.objective(encodings)
    return {
        "objective": dcca + self.lam * regularisation,
        "dcca": dcca,
        "noise_regularisation": regularisation,
    }

Objectives¶

objectives ¶

Differentiable CCA losses for the deep models.

CCALoss ¶

CCALoss(reg_covar: float = 1e-06)

Bases: Module

Two-view deep CCA loss (Andrew et al., 2013).

\[ \mathcal{L} = -\| \Sigma_{11}^{-1/2} \Sigma_{12} \Sigma_{22}^{-1/2} \|_F^2, \]

minus the sum of squared canonical correlations of the batch, with ridge-regularised within-view covariances.

Parameters:

Name	Type	Description	Default
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Andrew, G., Arora, R., Bilmes, J., & Livescu, K. (2013). Deep canonical correlation analysis. ICML.

Examples:

>>> import torch
>>> from cca_zoo.deep.objectives import CCALoss
>>> loss = CCALoss(reg_covar=1e-4)([torch.randn(32, 4), torch.randn(32, 4)])

Source code in cca_zoo/deep/objectives.py

def __init__(self, reg_covar: float = 1e-6) -> None:
    super().__init__()
    self.reg_covar = reg_covar

forward ¶

forward(representations: list[Tensor]) -> torch.Tensor

The loss of two views.

Parameters:

Name	Type	Description	Default
`representations`	`list[Tensor]`	One tensor of shape (batch_size, n_components) per view.	required

Returns:

Type	Description
`Tensor`	The loss.

Raises:

Type	Description
`ValueError`	If there are not exactly two views.

Source code in cca_zoo/deep/objectives.py

def forward(self, representations: list[torch.Tensor]) -> torch.Tensor:
    """The loss of two views.

    Args:
        representations: One tensor of shape (batch_size, n_components) per
            view.

    Returns:
        The loss.

    Raises:
        ValueError: If there are not exactly two views.
    """
    if len(representations) != 2:
        raise ValueError(
            "CCALoss expects exactly 2 representations, "
            f"got {len(representations)}."
        )
    z1, z2 = representations
    n = z1.shape[0]
    d1, d2 = z1.shape[1], z2.shape[1]

    z1 = z1 - z1.mean(dim=0)
    z2 = z2 - z2.mean(dim=0)

    s11 = (z1.T @ z1) / (n - 1) + self.reg_covar * torch.eye(
        d1, device=z1.device, dtype=z1.dtype
    )
    s22 = (z2.T @ z2) / (n - 1) + self.reg_covar * torch.eye(
        d2, device=z2.device, dtype=z2.dtype
    )
    s12 = (z1.T @ z2) / (n - 1)

    s11_inv_sqrt = _inv_sqrtm(s11, self.reg_covar)
    s22_inv_sqrt = _inv_sqrtm(s22, self.reg_covar)

    t = s11_inv_sqrt @ s12 @ s22_inv_sqrt
    # Squared singular values = eigenvalues of T^T T
    tt = t.T @ t
    eigvals = torch.linalg.eigvalsh(tt)
    eigvals = torch.clamp(eigvals, min=0.0)
    return -eigvals.sum()

MCCALoss ¶

MCCALoss(reg_covar: float = 1e-06)

Bases: Module

Sum of :class:CCALoss over every pair of views.

Parameters:

Name	Type	Description	Default
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Kettenring, J. R. (1971). Canonical analysis of several sets of variables. Biometrika, 58(3), 433-451.

Examples:

>>> import torch
>>> from cca_zoo.deep.objectives import MCCALoss
>>> loss = MCCALoss(reg_covar=1e-4)([torch.randn(32, 4) for _ in range(3)])

Source code in cca_zoo/deep/objectives.py

def __init__(self, reg_covar: float = 1e-6) -> None:
    super().__init__()
    self.reg_covar = reg_covar
    self._cca_loss = CCALoss(reg_covar=reg_covar)

forward ¶

forward(representations: list[Tensor]) -> torch.Tensor

The summed pairwise loss.

Parameters:

Name	Type	Description	Default
`representations`	`list[Tensor]`	One tensor of shape (batch_size, n_components) per view.	required

Returns:

Type	Description
`Tensor`	The loss.

Source code in cca_zoo/deep/objectives.py

def forward(self, representations: list[torch.Tensor]) -> torch.Tensor:
    """The summed pairwise loss.

    Args:
        representations: One tensor of shape (batch_size, n_components) per
            view.

    Returns:
        The loss.
    """
    n_views = len(representations)
    total = torch.tensor(0.0, device=representations[0].device)
    for i in range(n_views):
        for j in range(i + 1, n_views):
            total = total + self._cca_loss([representations[i], representations[j]])
    return total

GCCALoss ¶

GCCALoss(reg_covar: float = 1e-06)

Bases: Module

Generalized CCA loss: minus the top eigenvalues of the summed projections.

\[ \mathcal{L} = -\sum_{d=1}^{k} \lambda_d\Bigl(\sum_i P_i\Bigr), \qquad P_i = \frac{H_i H_i^\top}{n - 1}, \]

for ridge-whitened, centred encodings \(H_i\), so that \(P_i\) projects onto view \(i\)'s encodings. The loss lies in \([-kM, 0]\) for \(M\) views at any batch size, and is \(-kM\) when every view's encodings agree.

Parameters:

Name	Type	Description	Default
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Benton, A., Khayrallah, H., Gujral, B., Reisinger, D. A., Zhang, S., & Arora, R. (2019). Deep generalized canonical correlation analysis. RepL4NLP.

Examples:

>>> import torch
>>> from cca_zoo.deep.objectives import GCCALoss
>>> loss = GCCALoss(reg_covar=1e-4)([torch.randn(32, 4) for _ in range(3)])

Source code in cca_zoo/deep/objectives.py

def __init__(self, reg_covar: float = 1e-6) -> None:
    super().__init__()
    self.reg_covar = reg_covar

forward ¶

forward(representations: list[Tensor]) -> torch.Tensor

The generalized CCA loss.

Parameters:

Name	Type	Description	Default
`representations`	`list[Tensor]`	One tensor of shape (batch_size, n_components) per view.	required

Returns:

Type	Description
`Tensor`	The loss.

Source code in cca_zoo/deep/objectives.py

def forward(self, representations: list[torch.Tensor]) -> torch.Tensor:
    """The generalized CCA loss.

    Args:
        representations: One tensor of shape (batch_size, n_components) per
            view.

    Returns:
        The loss.
    """
    n = representations[0].shape[0]
    whitened = []
    for z in representations:
        z_c = z - z.mean(dim=0)
        cov = (z_c.T @ z_c) / (n - 1) + self.reg_covar * torch.eye(
            z_c.shape[1], device=z_c.device, dtype=z_c.dtype
        )
        whitened.append(z_c @ _inv_sqrtm(cov, self.reg_covar))

    # The summed projections onto each view's encodings, shape (n, n).
    m = sum(h @ h.T for h in whitened) / (n - 1)
    eigvals = torch.linalg.eigvalsh(m)
    k = representations[0].shape[1]
    top_eigvals = eigvals[-k:]
    objective: torch.Tensor = -top_eigvals.sum()
    return objective

TCCALoss ¶

TCCALoss(reg_covar: float = 1e-06)

Bases: Module

Tensor CCA loss: minus the norm of the whitened cross-moment tensor.

\[ \mathcal{L} = -\| M \|_F, \qquad M = \tfrac{1}{n} \sum_s H_1[s] \otimes \cdots \otimes H_M[s], \]

for ridge-whitened, centred encodings \(H_i\).

Parameters:

Name	Type	Description	Default
`reg_covar`	`float`	Non-negative regularisation added to the diagonal of each covariance, as scikit-learn's `GaussianMixture`. Default is 1e-6.	`1e-06`

References

Luo, Y., Tao, D., Ramamohanarao, K., Xu, C., & Wen, Y. (2015). Tensor canonical correlation analysis for multi-view dimension reduction. IEEE Transactions on Knowledge and Data Engineering, 27(11), 3111-3124.

Examples:

>>> import torch
>>> from cca_zoo.deep.objectives import TCCALoss
>>> loss = TCCALoss(reg_covar=1e-4)([torch.randn(32, 4) for _ in range(3)])

Source code in cca_zoo/deep/objectives.py

def __init__(self, reg_covar: float = 1e-6) -> None:
    super().__init__()
    self.reg_covar = reg_covar

forward ¶

forward(representations: list[Tensor]) -> torch.Tensor

The tensor CCA loss.

Parameters:

Name	Type	Description	Default
`representations`	`list[Tensor]`	One tensor of shape (batch_size, n_components) per view.	required

Returns:

Type	Description
`Tensor`	The loss.

Source code in cca_zoo/deep/objectives.py

def forward(self, representations: list[torch.Tensor]) -> torch.Tensor:
    """The tensor CCA loss.

    Args:
        representations: One tensor of shape (batch_size, n_components) per
            view.

    Returns:
        The loss.
    """
    n = representations[0].shape[0]
    whitened = []
    for z in representations:
        z_c = z - z.mean(dim=0)
        cov = (z_c.T @ z_c) / (n - 1) + self.reg_covar * torch.eye(
            z_c.shape[1], device=z_c.device, dtype=z_c.dtype
        )
        whitened.append(z_c @ _inv_sqrtm(cov, self.reg_covar))

    # Build outer product tensor iteratively, shape (d, d, ..., d)
    m: torch.Tensor = whitened[0]
    for i in range(1, len(whitened)):
        el = whitened[i]
        for _ in range(len(m.shape) - 1):
            el = el.unsqueeze(1)
        m = m.unsqueeze(-1) * el

    # Average over samples
    m = m.mean(dim=0)
    objective: torch.Tensor = -torch.linalg.norm(m)
    return objective

cca_zoo.deep¶

Base class¶

BaseDeep ¶

forward ¶

loss ¶

predict_step ¶

configure_optimizers ¶

Data¶

MultiviewDataset ¶

__len__ ¶

__getitem__ ¶

Models¶

DCCA ¶

DCCAEY ¶

loss ¶

DCCANOI ¶

loss ¶

DCCASDL ¶

forward ¶

loss ¶

DCCAE ¶

loss ¶

DVCCA ¶

forward ¶

loss ¶

predict_step ¶

DVCCAPrivate ¶

private_means ¶

loss ¶

DTCCA ¶

DMCCA ¶

DGCCA ¶

DPCCA ¶

loss ¶

SplitAE ¶

loss ¶

BarlowTwins ¶

forward ¶

loss ¶

VICReg ¶

loss ¶

LeJEPA ¶

loss ¶

NRDCCA ¶

loss ¶

Objectives¶

objectives ¶

CCALoss ¶

forward ¶

MCCALoss ¶

forward ¶

GCCALoss ¶

forward ¶

TCCALoss ¶

forward ¶

len ¶

getitem ¶