Metrics Similarity#

similarity #

FUNCTION DESCRIPTION
cca

Compute CCA similarity between two embeddings.
pwcca

Compute Projection-Weighted CCA (PWCCA) similarity.
rsa_correlation

Compute Representational Similarity Analysis (RSA) correlation.

Functions#

cca(X: torch.Tensor, Y: torch.Tensor, n_eigen: int | None = None, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor #

Compute CCA similarity between two embeddings.

CCA (Canonical Correlation Analysis) finds linear projections of X and Y that maximize correlation. The CCA similarity is the mean of the top n_eigen canonical correlations [1].

PARAMETER DESCRIPTION
X

Input tensor of shape (n_samples, x_features).

TYPE: Tensor

Y

Input tensor of shape (n_samples, y_features).

TYPE: Tensor

n_eigen

Number of CCA components to use. If None, uses min(x_features, y_features).

TYPE: int, optional, by default None DEFAULT: None

eps

Regularization term added to covariance matrices for numerical stability.

TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS DESCRIPTION
Tensor

Mean canonical correlation (CCA similarity) in [0, 1].

Examples:

>>> X = torch.randn(100, 10)
>>> Y = torch.randn(100, 8)
>>> cca(X, Y, n_eigen=5)
0.45...

Perfect correlation:

>>> X = torch.randn(100, 5)
>>> cca(X, X)
1.0
Source code in spectre/metrics/similarity.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def cca(
    X: torch.Tensor,
    Y: torch.Tensor,
    n_eigen: int | None = None,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute CCA similarity between two embeddings.

    CCA (Canonical Correlation Analysis) finds linear projections of X and Y
    that maximize correlation. The CCA similarity is the mean of the top `n_eigen`
    canonical correlations [@raghu2017svcca].

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, x_features).

    Y : torch.Tensor
        Input tensor of shape (n_samples, y_features).

    n_eigen : int, optional, by default None
        Number of CCA components to use. If None, uses min(x_features, y_features).

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Regularization term added to covariance matrices for numerical stability.

    Returns
    -------
    torch.Tensor
        Mean canonical correlation (CCA similarity) in [0, 1].

    Examples
    --------
    >>> X = torch.randn(100, 10)
    >>> Y = torch.randn(100, 8)
    >>> cca(X, Y, n_eigen=5)
    0.45...

    Perfect correlation:

    >>> X = torch.randn(100, 5)
    >>> cca(X, X)
    1.0
    """
    check_2d(X)
    check_2d(Y)
    check_same_len(X, Y)
    check_same_device(X, Y)

    canonical_correlations, _, _ = _cca(
        X, Y, n_eigen=n_eigen, correlations_only=True, eps=eps
    )

    return canonical_correlations.mean()

pwcca(X: torch.Tensor, Y: torch.Tensor, n_eigen: int | None = None, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor #

Compute Projection-Weighted CCA (PWCCA) similarity.

PWCCA extends CCA by weighting canonical correlations by their importance to the original representations. This addresses the known limitation that standard CCA can give pessimistic similarity scores for neural network representations [2].

The weighting is based on how well each canonical direction explains variance in the original data, providing a more interpretable measure of representational similarity.

PARAMETER DESCRIPTION
X

Input tensor of shape (n_samples, x_features).

TYPE: Tensor

Y

Input tensor of shape (n_samples, y_features).

TYPE: Tensor

n_eigen

Number of CCA components to use. If None, uses min(x_features, y_features).

TYPE: int, optional, by default None DEFAULT: None

eps

Regularization term added to covariance matrices for numerical stability.

TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS DESCRIPTION
Tensor

PWCCA similarity score in [0, 1].

Examples:

>>> X = torch.randn(100, 10)
>>> Y = torch.randn(100, 8)
>>> pwcca(X, Y)
0.52...

Perfect correlation:

>>> X = torch.randn(100, 5)
>>> pwcca(X, X)
1.0
Source code in spectre/metrics/similarity.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def pwcca(
    X: torch.Tensor,
    Y: torch.Tensor,
    n_eigen: int | None = None,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute Projection-Weighted CCA (PWCCA) similarity.

    PWCCA extends CCA by weighting canonical correlations by their importance
    to the original representations. This addresses the known limitation that
    standard CCA can give pessimistic similarity scores for neural network
    representations [@morcos2018insights].

    The weighting is based on how well each canonical direction explains
    variance in the original data, providing a more interpretable measure
    of representational similarity.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, x_features).

    Y : torch.Tensor
        Input tensor of shape (n_samples, y_features).

    n_eigen : int, optional, by default None
        Number of CCA components to use. If None, uses min(x_features, y_features).

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Regularization term added to covariance matrices for numerical stability.

    Returns
    -------
    torch.Tensor
        PWCCA similarity score in [0, 1].

    Examples
    --------
    >>> X = torch.randn(100, 10)
    >>> Y = torch.randn(100, 8)
    >>> pwcca(X, Y)
    0.52...

    Perfect correlation:

    >>> X = torch.randn(100, 5)
    >>> pwcca(X, X)
    1.0
    """
    check_2d(X)
    check_2d(Y)
    check_same_len(X, Y)
    check_same_device(X, Y)

    n_samples = X.shape[0]

    X_centered = X - X.mean(dim=0, keepdim=True)

    # Compute CCA components
    canonical_correlations, _, canonical_projections = _cca(
        X, Y, n_eigen=n_eigen, correlations_only=True, eps=eps
    )

    # Compute importance weights based on explained variance
    projection_vars = (canonical_projections**2).sum(dim=0) / n_samples
    total_var = (X_centered**2).sum() / n_samples

    # Normalize weights to sum to 1
    alpha = projection_vars / (total_var + eps)
    alpha = alpha / (alpha.sum() + eps)

    # Weighted average of canonical correlations
    return torch.clamp((alpha * canonical_correlations).sum(), 0.0, 1.0)

rsa_correlation(X: torch.Tensor, Y: torch.Tensor, distance_fn: str = 'euclidean', correlation_type: str = 'spearman') -> torch.Tensor #

Compute Representational Similarity Analysis (RSA) correlation.

RSA compares second-order similarity by constructing Representational Dissimilarity Matrices (RDMs) for each representation, then computing their correlation. This captures how samples relate to each other rather than absolute activation values.

PARAMETER DESCRIPTION
X

First feature matrix, shape (n_samples, x_features).

TYPE: Tensor

Y

Second feature matrix, shape (n_samples, y_features).

TYPE: Tensor

distance_fn

Distance metric from PairwiseDistanceRegistry.

TYPE: PairwiseDistance | str, optional, by default "euclidean" DEFAULT: 'euclidean'

correlation_type

Type of correlation: "spearman", "pearson".

TYPE: str, optional, by default "spearman" DEFAULT: 'spearman'

RETURNS DESCRIPTION
Tensor

Correlation coefficient in [-1, 1]. Higher values indicate greater similarity in how representations organize samples relationally.

Examples:

>>> X = torch.randn(50, 20)
>>> Y = torch.randn(50, 15)
>>> rsa_correlation(X, Y, correlation_type="spearman")
0.32...

Use Pearson correlation:

>>> rsa_correlation(X, Y, correlation_type="pearson")
0.28...
Source code in spectre/metrics/similarity.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def rsa_correlation(
    X: torch.Tensor,
    Y: torch.Tensor,
    distance_fn: str = "euclidean",
    correlation_type: str = "spearman",
) -> torch.Tensor:
    """
    Compute Representational Similarity Analysis (RSA) correlation.

    RSA compares second-order similarity by constructing Representational
    Dissimilarity Matrices (RDMs) for each representation, then computing
    their correlation. This captures how samples relate to each other
    rather than absolute activation values.

    Parameters
    ----------
    X : torch.Tensor
        First feature matrix, shape (n_samples, x_features).

    Y : torch.Tensor
        Second feature matrix, shape (n_samples, y_features).

    distance_fn : PairwiseDistance | str, optional, by default "euclidean"
        Distance metric from `PairwiseDistanceRegistry`.

    correlation_type : str, optional, by default "spearman"
        Type of correlation: "spearman", "pearson".

    Returns
    -------
    torch.Tensor
        Correlation coefficient in [-1, 1]. Higher values indicate greater
        similarity in how representations organize samples relationally.

    Examples
    --------
    >>> X = torch.randn(50, 20)
    >>> Y = torch.randn(50, 15)
    >>> rsa_correlation(X, Y, correlation_type="spearman")
    0.32...

    Use Pearson correlation:

    >>> rsa_correlation(X, Y, correlation_type="pearson")
    0.28...
    """
    check_2d(X)
    check_2d(Y)
    check_same_len(X, Y)
    check_same_device(X, Y)

    distance_module = initialize_distance_fn(distance_fn)
    rdm_X = distance_module(X)
    rdm_Y = distance_module(Y)

    # Flatten upper triangular parts (excluding diagonal)
    triu_indices = torch.triu_indices(X.shape[0], X.shape[0], offset=1, device=X.device)
    rdm_X_flat = rdm_X[triu_indices[0], triu_indices[1]]
    rdm_Y_flat = rdm_Y[triu_indices[0], triu_indices[1]]

    if correlation_type == "pearson":
        corr_matrix = torch.corrcoef(torch.stack([rdm_X_flat, rdm_Y_flat]))

        return corr_matrix[0, 1]

    elif correlation_type == "spearman":
        rank_X = torch.argsort(torch.argsort(rdm_X_flat)).float()
        rank_Y = torch.argsort(torch.argsort(rdm_Y_flat)).float()
        corr_matrix = torch.corrcoef(torch.stack([rank_X, rank_Y]))

        return corr_matrix[0, 1]

    else:
        raise ValueError(
            f"Invalid correlation_type '{correlation_type}'. "
            f"Choose from: 'pearson', 'spearman'."
        )