Skip to content

analysis

shrinkai.analysis

Representation-alignment analysis between a teacher and a student model.

FeatureAnalyzer runs both models (wrapped in shrinkai.adapters.FeatureExtractor) over a dataloader and scores, layer by layer, how well the student's internal representations match the teacher's using metrics such as CKAMetric (Centered Kernel Alignment), RSAMetric (Representational Similarity Analysis), and SpatialAttentionMetric. This is diagnostic tooling: it helps decide where a feature-based distillation loss would help most, rather than being a loss itself.

Modules:

Name Description
analyzer
metrics

Classes:

Name Description
BaseFeatureMetric

Abstract base class for all feature alignment metrics.
CKAMetric

Linear Centered Kernel Alignment (CKA) (Kornblith et al. (2019), building on the
FeatureAnalyzer

Evaluates how well a student model mimics the teacher's internal hidden states.
FeatureAnalyzerReport

Holds representation analysis results and renders a dashboard.
RSAMetric

Representational Similarity Analysis (RSA) using Pearson correlation
SpatialAttentionMetric

Spatial Attention Transfer similarity (Zagoruyko & Komodakis (2017)).

Classes

BaseFeatureMetric

Bases: ABC

Abstract base class for all feature alignment metrics.

Methods:

Name Description
compute

Computes the alignment score between two feature maps.

Attributes:

Name Type Description
name str

Display name of the metric (e.g., 'CKA', 'RSA').
Source code in src/shrinkai/analysis/metrics.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class BaseFeatureMetric(ABC):
    """Abstract base class for all feature alignment metrics."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Display name of the metric (e.g., 'CKA', 'RSA')."""
        pass

    @abstractmethod
    def compute(self, student_features: torch.Tensor, teacher_features: torch.Tensor) -> float:
        """Computes the alignment score between two feature maps.

        Args:
            student_features: Tensor of shape [Batch, ...].
            teacher_features: Tensor of shape [Batch, ...].

        Returns:
            float: A score indicating representation similarity.
        """
        pass
Attributes
name abstractmethod property 
name: str

Display name of the metric (e.g., 'CKA', 'RSA').

Methods:
compute abstractmethod 
compute(
    student_features: Tensor, teacher_features: Tensor
) -> float

Computes the alignment score between two feature maps.

Parameters:

Name Type Description Default
student_features Tensor

Tensor of shape [Batch, ...].

 required
teacher_features Tensor

Tensor of shape [Batch, ...].

 required

Returns:

Name Type Description
float float

A score indicating representation similarity.
Source code in src/shrinkai/analysis/metrics.py 
16
17
18
19
20
21
22
23
24
25
26
27
@abstractmethod
def compute(self, student_features: torch.Tensor, teacher_features: torch.Tensor) -> float:
    """Computes the alignment score between two feature maps.

    Args:
        student_features: Tensor of shape [Batch, ...].
        teacher_features: Tensor of shape [Batch, ...].

    Returns:
        float: A score indicating representation similarity.
    """
    pass

CKAMetric

Bases: BaseFeatureMetric

Linear Centered Kernel Alignment (CKA) (Kornblith et al. (2019), building on the Hilbert-Schmidt Independence Criterion of Gretton et al. (2005)).

Measures the similarity of representations across models with different architectures or channel dimensions. A score of 1.0 means identical representational geometry; 0.0 means completely orthogonal.

Equation
\[CKA(K, L) = \frac{HSIC(K, L)}{\sqrt{HSIC(K, K) \cdot HSIC(L, L)}}\]

where \(K = X_c X_c^\top\) and \(L = Y_c Y_c^\top\) are the Gram matrices of the (already mean-centered) student and teacher activations \(X_c, Y_c\), and \(HSIC(K, L) \propto \langle K, L \rangle_F = \text{tr}(KL)\) for linear kernels on centered data. The shared normalization constant of the HSIC estimator cancels out in the ratio, so this implementation computes it directly as \(\text{tr}(KL) / (\|K\|_F \|L\|_F)\).

Source code in src/shrinkai/analysis/metrics.py 
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
65
66
67
68
69
70
71
72
73
class CKAMetric(BaseFeatureMetric):
    r"""Linear Centered Kernel Alignment (CKA) (Kornblith et al. (2019), building on the
    Hilbert-Schmidt Independence Criterion of Gretton et al. (2005)).

    Measures the similarity of representations across models with different
    architectures or channel dimensions. A score of 1.0 means identical
    representational geometry; 0.0 means completely orthogonal.

    Equation:
        $$CKA(K, L) = \frac{HSIC(K, L)}{\sqrt{HSIC(K, K) \cdot HSIC(L, L)}}$$

        where $K = X_c X_c^\top$ and $L = Y_c Y_c^\top$ are the Gram matrices of the
        (already mean-centered) student and teacher activations $X_c, Y_c$, and
        $HSIC(K, L) \propto \langle K, L \rangle_F = \text{tr}(KL)$ for linear kernels
        on centered data. The shared normalization constant of the HSIC estimator
        cancels out in the ratio, so this implementation computes it directly as
        $\text{tr}(KL) / (\|K\|_F \|L\|_F)$.
    """

    @property
    def name(self) -> str:
        return "CKA (Linear)"

    def compute(self, student_features: torch.Tensor, teacher_features: torch.Tensor) -> float:
        if student_features.dim() > 2:
            student_features = student_features.view(student_features.size(0), -1)
        if teacher_features.dim() > 2:
            teacher_features = teacher_features.view(teacher_features.size(0), -1)

        s_centered = student_features - student_features.mean(dim=0, keepdim=True)
        t_centered = teacher_features - teacher_features.mean(dim=0, keepdim=True)

        s_gram = s_centered @ s_centered.t()
        t_gram = t_centered @ t_centered.t()

        hsic = torch.sum(s_gram * t_gram)
        norm_s = torch.sqrt(torch.sum(s_gram * s_gram))
        norm_t = torch.sqrt(torch.sum(t_gram * t_gram))

        if norm_s == 0 or norm_t == 0:
            return 0.0

        cka = hsic / (norm_s * norm_t)
        return cka.item()

FeatureAnalyzer

Evaluates how well a student model mimics the teacher's internal hidden states.

Methods:

Name Description
__init__

Initializes the FeatureAnalyzer.
evaluate

Runs the dataset through both models and computes alignment metrics per layer.
Source code in src/shrinkai/analysis/analyzer.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 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
141
class FeatureAnalyzer:
    """Evaluates how well a student model mimics the teacher's internal hidden states."""

    def __init__(
        self,
        teacher_extractor: FeatureExtractor,
        student_extractor: FeatureExtractor,
        device: torch.device | str = "auto",
    ) -> None:
        """Initializes the FeatureAnalyzer.

        Args:
            teacher_extractor: Teacher model wrapped in FeatureExtractor.
            student_extractor: Student model wrapped in FeatureExtractor.
            device: Computing target.
        """
        self.teacher = teacher_extractor
        self.student = student_extractor
        self.device = resolve_device(device)
        self.teacher.to(self.device)
        self.student.to(self.device)

    def evaluate(
        self,
        dataloader: DataLoader,
        metrics: list[str | BaseFeatureMetric] | None = None,
    ) -> FeatureAnalyzerReport:
        """Runs the dataset through both models and computes alignment metrics per layer.

        Args:
            dataloader: Dataloader yielding validation batches.
            metrics: List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

        Returns:
            FeatureAnalyzerReport: Formatted report ready for `.show()`.
        """
        if metrics is None:
            metrics = ["cka"]
        self.teacher.eval()
        self.student.eval()

        active_metrics = []
        for m in metrics:
            if isinstance(m, str):
                if m.lower() not in FEATURE_METRICS:
                    raise ValueError(
                        f"Unknown metric '{m}'. Available: {list(FEATURE_METRICS.keys())}"
                    )
                active_metrics.append(FEATURE_METRICS[m.lower()])
            else:
                active_metrics.append(m)

        accumulated_scores = defaultdict(lambda: defaultdict(float))
        num_batches = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(dataloader):
                inputs = batch[0].to(self.device)
                num_batches += 1

                _, s_features = self.student(inputs)
                _, t_features = self.teacher(inputs)

                common_keys = set(s_features.keys()).intersection(set(t_features.keys()))

                if batch_idx == 0 and not common_keys:
                    warnings.warn(
                        "No common alias keys between the teacher and the studen."
                        "Check `target_layers` in your FeatureExtractors.",
                        stacklevel=2,
                    )
                    break

                for key in common_keys:
                    for metric in active_metrics:
                        score = metric.compute(s_features[key], t_features[key])
                        accumulated_scores[key][metric.name] += score

        if num_batches == 0:
            raise ValueError("No batches in the provided dataloader.")

        final_scores = defaultdict(dict)
        for layer, metric_dict in accumulated_scores.items():
            for metric_name, total_score in metric_dict.items():
                final_scores[layer][metric_name] = total_score / num_batches

        return FeatureAnalyzerReport(dict(final_scores))
Methods:
__init__ 
__init__(
    teacher_extractor: FeatureExtractor,
    student_extractor: FeatureExtractor,
    device: device | str = "auto",
) -> None

Initializes the FeatureAnalyzer.

Parameters:

Name Type Description Default
teacher_extractor FeatureExtractor

Teacher model wrapped in FeatureExtractor.

 required
student_extractor FeatureExtractor

Student model wrapped in FeatureExtractor.

 required
device device | str

Computing target.

 'auto'
Source code in src/shrinkai/analysis/analyzer.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(
    self,
    teacher_extractor: FeatureExtractor,
    student_extractor: FeatureExtractor,
    device: torch.device | str = "auto",
) -> None:
    """Initializes the FeatureAnalyzer.

    Args:
        teacher_extractor: Teacher model wrapped in FeatureExtractor.
        student_extractor: Student model wrapped in FeatureExtractor.
        device: Computing target.
    """
    self.teacher = teacher_extractor
    self.student = student_extractor
    self.device = resolve_device(device)
    self.teacher.to(self.device)
    self.student.to(self.device)
evaluate 
evaluate(
    dataloader: DataLoader,
    metrics: list[str | BaseFeatureMetric] | None = None,
) -> FeatureAnalyzerReport

Runs the dataset through both models and computes alignment metrics per layer.

Parameters:

Name Type Description Default
dataloader DataLoader

Dataloader yielding validation batches.

 required
metrics list[str | BaseFeatureMetric] | None

List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

 None

Returns:

Name Type Description
FeatureAnalyzerReport FeatureAnalyzerReport

Formatted report ready for .show().
Source code in src/shrinkai/analysis/analyzer.py 
 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
141
def evaluate(
    self,
    dataloader: DataLoader,
    metrics: list[str | BaseFeatureMetric] | None = None,
) -> FeatureAnalyzerReport:
    """Runs the dataset through both models and computes alignment metrics per layer.

    Args:
        dataloader: Dataloader yielding validation batches.
        metrics: List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

    Returns:
        FeatureAnalyzerReport: Formatted report ready for `.show()`.
    """
    if metrics is None:
        metrics = ["cka"]
    self.teacher.eval()
    self.student.eval()

    active_metrics = []
    for m in metrics:
        if isinstance(m, str):
            if m.lower() not in FEATURE_METRICS:
                raise ValueError(
                    f"Unknown metric '{m}'. Available: {list(FEATURE_METRICS.keys())}"
                )
            active_metrics.append(FEATURE_METRICS[m.lower()])
        else:
            active_metrics.append(m)

    accumulated_scores = defaultdict(lambda: defaultdict(float))
    num_batches = 0

    with torch.no_grad():
        for batch_idx, batch in enumerate(dataloader):
            inputs = batch[0].to(self.device)
            num_batches += 1

            _, s_features = self.student(inputs)
            _, t_features = self.teacher(inputs)

            common_keys = set(s_features.keys()).intersection(set(t_features.keys()))

            if batch_idx == 0 and not common_keys:
                warnings.warn(
                    "No common alias keys between the teacher and the studen."
                    "Check `target_layers` in your FeatureExtractors.",
                    stacklevel=2,
                )
                break

            for key in common_keys:
                for metric in active_metrics:
                    score = metric.compute(s_features[key], t_features[key])
                    accumulated_scores[key][metric.name] += score

    if num_batches == 0:
        raise ValueError("No batches in the provided dataloader.")

    final_scores = defaultdict(dict)
    for layer, metric_dict in accumulated_scores.items():
        for metric_name, total_score in metric_dict.items():
            final_scores[layer][metric_name] = total_score / num_batches

    return FeatureAnalyzerReport(dict(final_scores))

FeatureAnalyzerReport

Holds representation analysis results and renders a dashboard.

Methods:

Name Description
__init__

Initializes the report.
show

Renders an interactive formatted analysis table in the console.
Source code in src/shrinkai/analysis/analyzer.py 
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
class FeatureAnalyzerReport:
    """Holds representation analysis results and renders a dashboard."""

    def __init__(self, layer_scores: dict[str, dict[str, float]]) -> None:
        """Initializes the report.

        Args:
            layer_scores: Nested dict `{ "stage_1": {"CKA (Linear)": 0.85}, ... }`.
        """
        self.layer_scores = layer_scores

    def show(self) -> None:
        """Renders an interactive formatted analysis table in the console."""
        if not self.layer_scores:
            print("No data to report.")
            return

        console = Console()
        table = Table(
            title="Representation Alignment Report (Feature Distillation)",
            header_style="bold magenta",
        )

        table.add_column("Layer / Stage Alias", style="bold")

        first_layer = next(iter(self.layer_scores.values()))
        metric_names = list(first_layer.keys())
        for metric in metric_names:
            table.add_column(metric, justify="right")

        for layer_name, metrics in self.layer_scores.items():
            row_data = [layer_name]
            for metric in metric_names:
                score = metrics.get(metric, 0.0)
                color = "green" if score >= 0.8 else "yellow" if score >= 0.5 else "red"
                row_data.append(f"[{color}]{score:.3f}[/{color}]")
            table.add_row(*row_data)

        console.print(table)
Methods:
__init__ 
__init__(layer_scores: dict[str, dict[str, float]]) -> None

Initializes the report.

Parameters:

Name Type Description Default
layer_scores dict[str, dict[str, float]]

Nested dict { "stage_1": {"CKA (Linear)": 0.85}, ... }.

 required
Source code in src/shrinkai/analysis/analyzer.py 
17
18
19
20
21
22
23
def __init__(self, layer_scores: dict[str, dict[str, float]]) -> None:
    """Initializes the report.

    Args:
        layer_scores: Nested dict `{ "stage_1": {"CKA (Linear)": 0.85}, ... }`.
    """
    self.layer_scores = layer_scores
show 
show() -> None

Renders an interactive formatted analysis table in the console.

Source code in src/shrinkai/analysis/analyzer.py 
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
def show(self) -> None:
    """Renders an interactive formatted analysis table in the console."""
    if not self.layer_scores:
        print("No data to report.")
        return

    console = Console()
    table = Table(
        title="Representation Alignment Report (Feature Distillation)",
        header_style="bold magenta",
    )

    table.add_column("Layer / Stage Alias", style="bold")

    first_layer = next(iter(self.layer_scores.values()))
    metric_names = list(first_layer.keys())
    for metric in metric_names:
        table.add_column(metric, justify="right")

    for layer_name, metrics in self.layer_scores.items():
        row_data = [layer_name]
        for metric in metric_names:
            score = metrics.get(metric, 0.0)
            color = "green" if score >= 0.8 else "yellow" if score >= 0.5 else "red"
            row_data.append(f"[{color}]{score:.3f}[/{color}]")
        table.add_row(*row_data)

    console.print(table)

RSAMetric

Bases: BaseFeatureMetric

Representational Similarity Analysis (RSA) using Pearson correlation (Kriegeskorte et al. (2008)).

Measures if the relative distances between samples in a batch are preserved between the teacher and the student, regardless of their hidden dimension sizes.

Equation
\[RSA = \frac{\text{cov}(R_S, R_T)}{\sigma_{R_S} \, \sigma_{R_T}}\]

where \(R_S\) and \(R_T\) are the upper-triangular entries of the student's and teacher's Representational (Dis)similarity Matrices, here, pairwise cosine similarities between samples in the batch, and the RSA score is their Pearson correlation across the batch.

Source code in src/shrinkai/analysis/metrics.py 
 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
class RSAMetric(BaseFeatureMetric):
    r"""Representational Similarity Analysis (RSA) using Pearson correlation
    (Kriegeskorte et al. (2008)).

    Measures if the relative distances between samples in a batch are preserved
    between the teacher and the student, regardless of their hidden dimension sizes.

    Equation:
        $$RSA = \frac{\text{cov}(R_S, R_T)}{\sigma_{R_S} \, \sigma_{R_T}}$$

        where $R_S$ and $R_T$ are the upper-triangular entries of the student's and
        teacher's Representational (Dis)similarity Matrices, here, pairwise cosine
        similarities between samples in the batch, and the RSA score is their
        Pearson correlation across the batch.
    """

    @property
    def name(self) -> str:
        return "RSA (Pearson)"

    def compute(self, student_features: torch.Tensor, teacher_features: torch.Tensor) -> float:
        # flatten [B, C, H, W] -> [B, Features]
        s_flat = student_features.view(student_features.size(0), -1)
        t_flat = teacher_features.view(teacher_features.size(0), -1)

        s_sim = F.cosine_similarity(s_flat.unsqueeze(1), s_flat.unsqueeze(0), dim=-1)
        t_sim = F.cosine_similarity(t_flat.unsqueeze(1), t_flat.unsqueeze(0), dim=-1)

        # top of matrix
        idx = torch.triu_indices(s_sim.size(0), s_sim.size(1), offset=1)
        s_pdist = s_sim[idx[0], idx[1]]
        t_pdist = t_sim[idx[0], idx[1]]

        # pearson correlation
        s_mean, t_mean = s_pdist.mean(), t_pdist.mean()
        s_centered, t_centered = s_pdist - s_mean, t_pdist - t_mean

        cov = (s_centered * t_centered).sum()
        var_s = (s_centered**2).sum()
        var_t = (t_centered**2).sum()

        if var_s == 0 or var_t == 0:
            return 0.0

        pearson_corr = cov / torch.sqrt(var_s * var_t)
        return pearson_corr.item()

SpatialAttentionMetric

Bases: BaseFeatureMetric

Spatial Attention Transfer similarity (Zagoruyko & Komodakis (2017)).

Collapses the channel dimension to measure if the student and teacher activate on the same spatial regions of the input (e.g., the foreground object).

Equation
\[A = \sum_{c=1}^{C} |f_c|, \qquad Spatial = \frac{A_s \cdot A_t}{\|A_s\| \, \|A_t\|}\]

where \(f_c\) is the activation map of channel \(c\), \(A\) is the resulting spatial attention map (summed absolute activations across channels, the \(\mathcal{F}_{sum}^{p=1}\) mapping of the original paper), and the metric is the cosine similarity between the student's and teacher's (L2-normalized, spatially-resized-if-needed) attention maps.

Source code in src/shrinkai/analysis/metrics.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
class SpatialAttentionMetric(BaseFeatureMetric):
    r"""Spatial Attention Transfer similarity (Zagoruyko & Komodakis (2017)).

    Collapses the channel dimension to measure if the student and teacher
    activate on the same spatial regions of the input (e.g., the foreground object).

    Equation:
        $$A = \sum_{c=1}^{C} |f_c|, \qquad Spatial = \frac{A_s \cdot A_t}{\|A_s\| \, \|A_t\|}$$

        where $f_c$ is the activation map of channel $c$, $A$ is the resulting
        spatial attention map (summed absolute activations across channels, the
        $\mathcal{F}_{sum}^{p=1}$ mapping of the original paper), and the metric is
        the cosine similarity between the student's and teacher's (L2-normalized,
        spatially-resized-if-needed) attention maps.
    """

    @property
    def name(self) -> str:
        return "Spatial Attention"

    def compute(self, student_features: torch.Tensor, teacher_features: torch.Tensor) -> float:
        # only 4D tensors (CNNs: Batch, Channel, Height, Width)
        if student_features.dim() != 4 or teacher_features.dim() != 4:
            return 0.0

        # spatial attention map
        s_att = torch.sum(torch.abs(student_features), dim=1, keepdim=True)
        t_att = torch.sum(torch.abs(teacher_features), dim=1, keepdim=True)

        # align spatial dimensions
        if s_att.shape[2:] != t_att.shape[2:]:
            s_att = F.adaptive_avg_pool2d(s_att, output_size=t_att.shape[2:])

        # normalization
        s_att = F.normalize(s_att.view(s_att.size(0), -1), p=2, dim=1)
        t_att = F.normalize(t_att.view(t_att.size(0), -1), p=2, dim=1)

        cos_sim = (s_att * t_att).sum(dim=1).mean()
        return cos_sim.item()