Skip to content

adapters

shrinkai.adapters

Model-agnostic adapters bridging teacher/student architectures for distillation.

FeatureExtractor wraps any nn.Module to capture intermediate activations via forward hooks, without modifying the model's source code. FeatureProjector and AttentionHeadSelector then reconcile dimension mismatches between a teacher's and a student's internal representations (channel counts, attention head counts, ...) so that feature-based losses (shrinkai.distillation.losses) can compare them directly.

Modules:

Name Description
extractor
projector

Classes:

Name Description
AttentionHeadSelector

Adapts teacher attention maps to match student dimensions by selecting specific heads.
FeatureExtractor

Wraps a PyTorch model to extract intermediate feature maps via forward hooks.
FeatureProjector

Projects student features to match the channel dimensions of the teacher features.

Classes

AttentionHeadSelector

Bases: Module

Adapts teacher attention maps to match student dimensions by selecting specific heads.

In Transformer distillation (e.g., TinyBERT), a student often has fewer attention heads than the teacher (e.g., 2 vs 12). This adapter slices the teacher's attention tensors [Batch, Heads, Seq, Seq] to keep only the indices corresponding to the student.

Methods:

Name Description
__init__

Initializes the AttentionHeadSelector.
forward

Slices the attention tensors along the head dimension.
Source code in src/shrinkai/adapters/projector.py 
 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
class AttentionHeadSelector(nn.Module):
    """Adapts teacher attention maps to match student dimensions by selecting specific heads.

    In Transformer distillation (e.g., TinyBERT), a student often has fewer attention
    heads than the teacher (e.g., 2 vs 12). This adapter slices the teacher's attention
    tensors [Batch, Heads, Seq, Seq] to keep only the indices corresponding to the student.
    """

    def __init__(self, heads_to_keep: list[int]) -> None:
        """Initializes the AttentionHeadSelector.

        Args:
            heads_to_keep: List of integer indices representing which teacher heads
                to retain (e.g., [0, 6] to keep the first and seventh head).
        """
        super().__init__()
        self.heads_to_keep = heads_to_keep

    def forward(self, attention_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Slices the attention tensors along the head dimension.

        Args:
            attention_dict: Dictionary of attention tensors of shape [B, Num_Heads, S, S].

        Returns:
            dict[str, torch.Tensor]: Dictionary of sliced tensors of
                shape [B, len(heads_to_keep), S, S].
        """
        projected_attention = {}
        for layer_name, tensor in attention_dict.items():
            if tensor.dim() != 4:
                raise ValueError(
                    f"Expected 4D attention tensor [Batch, Heads, Seq, Seq], "
                    f"but got shape {tensor.shape} for layer '{layer_name}'."
                )

            num_heads = tensor.shape[1]
            invalid_heads = [h for h in self.heads_to_keep if h < -num_heads or h >= num_heads]
            if invalid_heads:
                raise IndexError(
                    f"Head indices {invalid_heads} in heads_to_keep are out of bounds. "
                    f"Layer '{layer_name}' only has {num_heads} heads."
                )
            projected_attention[layer_name] = tensor[:, self.heads_to_keep, :, :]

        return projected_attention
Methods:
__init__ 
__init__(heads_to_keep: list[int]) -> None

Initializes the AttentionHeadSelector.

Parameters:

Name Type Description Default
heads_to_keep list[int]

List of integer indices representing which teacher heads to retain (e.g., [0, 6] to keep the first and seventh head).

 required
Source code in src/shrinkai/adapters/projector.py 
 99
100
101
102
103
104
105
106
107
def __init__(self, heads_to_keep: list[int]) -> None:
    """Initializes the AttentionHeadSelector.

    Args:
        heads_to_keep: List of integer indices representing which teacher heads
            to retain (e.g., [0, 6] to keep the first and seventh head).
    """
    super().__init__()
    self.heads_to_keep = heads_to_keep
forward 
forward(
    attention_dict: dict[str, Tensor],
) -> dict[str, torch.Tensor]

Slices the attention tensors along the head dimension.

Parameters:

Name Type Description Default
attention_dict dict[str, Tensor]

Dictionary of attention tensors of shape [B, Num_Heads, S, S].

 required

Returns:

Type Description
dict[str, Tensor]

dict[str, torch.Tensor]: Dictionary of sliced tensors of shape [B, len(heads_to_keep), S, S].
Source code in src/shrinkai/adapters/projector.py 
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
def forward(self, attention_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Slices the attention tensors along the head dimension.

    Args:
        attention_dict: Dictionary of attention tensors of shape [B, Num_Heads, S, S].

    Returns:
        dict[str, torch.Tensor]: Dictionary of sliced tensors of
            shape [B, len(heads_to_keep), S, S].
    """
    projected_attention = {}
    for layer_name, tensor in attention_dict.items():
        if tensor.dim() != 4:
            raise ValueError(
                f"Expected 4D attention tensor [Batch, Heads, Seq, Seq], "
                f"but got shape {tensor.shape} for layer '{layer_name}'."
            )

        num_heads = tensor.shape[1]
        invalid_heads = [h for h in self.heads_to_keep if h < -num_heads or h >= num_heads]
        if invalid_heads:
            raise IndexError(
                f"Head indices {invalid_heads} in heads_to_keep are out of bounds. "
                f"Layer '{layer_name}' only has {num_heads} heads."
            )
        projected_attention[layer_name] = tensor[:, self.heads_to_keep, :, :]

    return projected_attention

FeatureExtractor

Bases: Module

Wraps a PyTorch model to extract intermediate feature maps via forward hooks.

This wrapper does not modify the original model's source code. It dynamically attaches hooks to intercept the output of specified layers during the forward pass.

Methods:

Name Description
__init__

Initializes the FeatureExtractor.
forward

Performs a forward pass and captures intermediate features.
remove_hooks

Removes all registered hooks to prevent memory leaks.
Source code in src/shrinkai/adapters/extractor.py 
 8
 9
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
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
class FeatureExtractor(nn.Module):
    """Wraps a PyTorch model to extract intermediate feature maps via forward hooks.

    This wrapper does not modify the original model's source code. It dynamically
    attaches hooks to intercept the output of specified layers during the forward pass.
    """

    def __init__(self, model: nn.Module, target_layers: Iterable[str] | Mapping[str, str]) -> None:
        """Initializes the FeatureExtractor.

        Args:
            model: The PyTorch model to extract features from.
            target_layers: An iterable of layer names to hook (e.g., ["layer1", "layer2"]),
                or a mapping dictionary to assign common aliases for cross-model matching
                (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

        Raises:
            ValueError: If a target layer does not exist in the model.
        """
        super().__init__()
        self.model = model
        if isinstance(target_layers, Mapping):
            aliases = list(target_layers.values())
            if len(set(aliases)) != len(aliases):
                raise ValueError("Many layers have the same alias.")
            self.target_layers = target_layers
        else:
            self.target_layers = {layer: layer for layer in target_layers}
        self.features: dict[str, torch.Tensor] = {}
        self._hooks: list[torch.utils.hooks.RemovableHandle] = []

        self._register_hooks()

    def _register_hooks(self) -> None:
        """Registers PyTorch forward hooks on the target layers."""
        model_modules = dict(self.model.named_modules())

        for layer_name, alias in self.target_layers.items():
            if layer_name not in model_modules:
                raise ValueError(
                    f"Layer '{layer_name}' not found in the model. "
                    f"Available layers: {list(model_modules.keys())[:10]}..."
                )

            module = model_modules[layer_name]

            def hook_fn(
                module: nn.Module, input: tuple, output: torch.Tensor, name: str = alias
            ) -> None:
                if isinstance(output, tuple):
                    output = output[0]  # take main tensor in first position
                self.features[name] = output

            handle = module.register_forward_hook(hook_fn)
            self._hooks.append(handle)

    def forward(
        self, x: torch.Tensor, *args: Any, **kwargs: Any
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """Performs a forward pass and captures intermediate features.

        Args:
            x: Input tensor.
            *args: Additional positional arguments for the model.
            **kwargs: Additional keyword arguments for the model.

        Returns:
            tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the
                final model output (logits) and a dictionary of extracted features.
        """
        self.features.clear()
        logits = self.model(x, *args, **kwargs)
        return logits, self.features.copy()

    def remove_hooks(self) -> None:
        """Removes all registered hooks to prevent memory leaks."""
        for handle in self._hooks:
            handle.remove()
        self._hooks.clear()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.remove_hooks()
Methods:
__init__ 
__init__(
    model: Module,
    target_layers: Iterable[str] | Mapping[str, str],
) -> None

Initializes the FeatureExtractor.

Parameters:

Name Type Description Default
model Module

The PyTorch model to extract features from.

 required
target_layers Iterable[str] | Mapping[str, str]

An iterable of layer names to hook (e.g., ["layer1", "layer2"]), or a mapping dictionary to assign common aliases for cross-model matching (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

 required

Raises:

Type Description
ValueError

If a target layer does not exist in the model.
Source code in src/shrinkai/adapters/extractor.py 
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
def __init__(self, model: nn.Module, target_layers: Iterable[str] | Mapping[str, str]) -> None:
    """Initializes the FeatureExtractor.

    Args:
        model: The PyTorch model to extract features from.
        target_layers: An iterable of layer names to hook (e.g., ["layer1", "layer2"]),
            or a mapping dictionary to assign common aliases for cross-model matching
            (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

    Raises:
        ValueError: If a target layer does not exist in the model.
    """
    super().__init__()
    self.model = model
    if isinstance(target_layers, Mapping):
        aliases = list(target_layers.values())
        if len(set(aliases)) != len(aliases):
            raise ValueError("Many layers have the same alias.")
        self.target_layers = target_layers
    else:
        self.target_layers = {layer: layer for layer in target_layers}
    self.features: dict[str, torch.Tensor] = {}
    self._hooks: list[torch.utils.hooks.RemovableHandle] = []

    self._register_hooks()
forward 
forward(
    x: Tensor, *args: Any, **kwargs: Any
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]

Performs a forward pass and captures intermediate features.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required
*args Any

Additional positional arguments for the model.

 ()
**kwargs Any

Additional keyword arguments for the model.

 {}

Returns:

Type Description
tuple[Tensor, dict[str, Tensor]]

tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the final model output (logits) and a dictionary of extracted features.
Source code in src/shrinkai/adapters/extractor.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def forward(
    self, x: torch.Tensor, *args: Any, **kwargs: Any
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """Performs a forward pass and captures intermediate features.

    Args:
        x: Input tensor.
        *args: Additional positional arguments for the model.
        **kwargs: Additional keyword arguments for the model.

    Returns:
        tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the
            final model output (logits) and a dictionary of extracted features.
    """
    self.features.clear()
    logits = self.model(x, *args, **kwargs)
    return logits, self.features.copy()
remove_hooks 
remove_hooks() -> None

Removes all registered hooks to prevent memory leaks.

Source code in src/shrinkai/adapters/extractor.py 
82
83
84
85
86
def remove_hooks(self) -> None:
    """Removes all registered hooks to prevent memory leaks."""
    for handle in self._hooks:
        handle.remove()
    self._hooks.clear()

FeatureProjector

Bases: Module

Projects student features to match the channel dimensions of the teacher features.

Uses 1x1 convolutions for spatial feature maps (B, C, H, W) or Linear layers for flattened features (B, C) or sequences (B, L, D).

Methods:

Name Description
__init__

Initializes the FeatureProjector.
forward

Applies the projection layers to the student's extracted features.
Source code in src/shrinkai/adapters/projector.py 
 8
 9
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class FeatureProjector(nn.Module):
    """Projects student features to match the channel dimensions of the teacher features.

    Uses 1x1 convolutions for spatial feature maps (B, C, H, W) or Linear layers
    for flattened features (B, C) or sequences (B, L, D).
    """

    def __init__(self, mapping_config: dict[str, dict[str, Any]]) -> None:
        """Initializes the FeatureProjector.

        Args:
            mapping_config: A dictionary defining the projection parameters for each layer.
                Format: {
                    "block_1": {"in_channels": 64, "out_channels": 128, "type": "conv", "use_norm": True},
                    "block_2": {"in_channels": 256, "out_channels": 512, "type": "linear"}
                }
                - `type`: Must be either 'conv' (default) or 'linear'.
                - `use_norm`: (Optional) If True, appends a normalization layer (BatchNorm2d
                for 'conv', LayerNorm for 'linear') to stabilize gradients. Defaults to False.
        """  # noqa: E501
        super().__init__()
        self.projectors = nn.ModuleDict()

        for layer_name, config in mapping_config.items():
            if "in_channels" not in config:
                raise ValueError(
                    f"Missing required key 'in_channels' for layer '{layer_name}'in mapping_config."
                )
            if "out_channels" not in config:
                raise ValueError(
                    f"Missing required key 'out_channels' for layer '{layer_name}'"
                    "in mapping_config."
                )

            in_c = config["in_channels"]
            out_c = config["out_channels"]
            proj_type = config.get("type", "conv")
            use_norm = config.get("use_norm", False)

            layers = []
            if proj_type == "conv":
                layers.append(nn.Conv2d(in_c, out_c, kernel_size=1, bias=not use_norm))
                if use_norm:
                    layers.append(nn.BatchNorm2d(out_c))
            elif proj_type == "linear":
                layers.append(nn.Linear(in_c, out_c, bias=not use_norm))
                if use_norm:
                    layers.append(nn.LayerNorm(out_c))
            else:
                raise ValueError(f"Unsupported projection type '{proj_type}'.")

            safe_key = layer_name.replace(".", "_")
            self.projectors[safe_key] = nn.Sequential(*layers)

        self.mapping_keys = {k: k.replace(".", "_") for k in mapping_config.keys()}

    def forward(self, student_features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Applies the projection layers to the student's extracted features.

        Args:
            student_features: Dictionary of raw feature tensors from the student.

        Returns:
            dict[str, torch.Tensor]: Dictionary of projected feature tensors.
        """
        projected_features = {}

        for layer_name, feature_tensor in student_features.items():
            safe_key = self.mapping_keys.get(layer_name)

            if safe_key and safe_key in self.projectors:
                projected_features[layer_name] = self.projectors[safe_key](feature_tensor)
            else:
                warnings.warn(
                    f"No projection configured for layer '{layer_name}'. "
                    "Feature tensor is returned. Check 'mapping_config' if this is unexpected.",
                    stacklevel=2,
                )
                projected_features[layer_name] = feature_tensor

        return projected_features
Methods:
__init__ 
__init__(mapping_config: dict[str, dict[str, Any]]) -> None

Initializes the FeatureProjector.

Parameters:

Name Type Description Default
mapping_config dict[str, dict[str, Any]]

A dictionary defining the projection parameters for each layer. Format: { "block_1": {"in_channels": 64, "out_channels": 128, "type": "conv", "use_norm": True}, "block_2": {"in_channels": 256, "out_channels": 512, "type": "linear"} } - type: Must be either 'conv' (default) or 'linear'. - use_norm: (Optional) If True, appends a normalization layer (BatchNorm2d for 'conv', LayerNorm for 'linear') to stabilize gradients. Defaults to False.

 required
Source code in src/shrinkai/adapters/projector.py 
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
def __init__(self, mapping_config: dict[str, dict[str, Any]]) -> None:
    """Initializes the FeatureProjector.

    Args:
        mapping_config: A dictionary defining the projection parameters for each layer.
            Format: {
                "block_1": {"in_channels": 64, "out_channels": 128, "type": "conv", "use_norm": True},
                "block_2": {"in_channels": 256, "out_channels": 512, "type": "linear"}
            }
            - `type`: Must be either 'conv' (default) or 'linear'.
            - `use_norm`: (Optional) If True, appends a normalization layer (BatchNorm2d
            for 'conv', LayerNorm for 'linear') to stabilize gradients. Defaults to False.
    """  # noqa: E501
    super().__init__()
    self.projectors = nn.ModuleDict()

    for layer_name, config in mapping_config.items():
        if "in_channels" not in config:
            raise ValueError(
                f"Missing required key 'in_channels' for layer '{layer_name}'in mapping_config."
            )
        if "out_channels" not in config:
            raise ValueError(
                f"Missing required key 'out_channels' for layer '{layer_name}'"
                "in mapping_config."
            )

        in_c = config["in_channels"]
        out_c = config["out_channels"]
        proj_type = config.get("type", "conv")
        use_norm = config.get("use_norm", False)

        layers = []
        if proj_type == "conv":
            layers.append(nn.Conv2d(in_c, out_c, kernel_size=1, bias=not use_norm))
            if use_norm:
                layers.append(nn.BatchNorm2d(out_c))
        elif proj_type == "linear":
            layers.append(nn.Linear(in_c, out_c, bias=not use_norm))
            if use_norm:
                layers.append(nn.LayerNorm(out_c))
        else:
            raise ValueError(f"Unsupported projection type '{proj_type}'.")

        safe_key = layer_name.replace(".", "_")
        self.projectors[safe_key] = nn.Sequential(*layers)

    self.mapping_keys = {k: k.replace(".", "_") for k in mapping_config.keys()}
forward 
forward(
    student_features: dict[str, Tensor],
) -> dict[str, torch.Tensor]

Applies the projection layers to the student's extracted features.

Parameters:

Name Type Description Default
student_features dict[str, Tensor]

Dictionary of raw feature tensors from the student.

 required

Returns:

Type Description
dict[str, Tensor]

dict[str, torch.Tensor]: Dictionary of projected feature tensors.
Source code in src/shrinkai/adapters/projector.py 
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
def forward(self, student_features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Applies the projection layers to the student's extracted features.

    Args:
        student_features: Dictionary of raw feature tensors from the student.

    Returns:
        dict[str, torch.Tensor]: Dictionary of projected feature tensors.
    """
    projected_features = {}

    for layer_name, feature_tensor in student_features.items():
        safe_key = self.mapping_keys.get(layer_name)

        if safe_key and safe_key in self.projectors:
            projected_features[layer_name] = self.projectors[safe_key](feature_tensor)
        else:
            warnings.warn(
                f"No projection configured for layer '{layer_name}'. "
                "Feature tensor is returned. Check 'mapping_config' if this is unexpected.",
                stacklevel=2,
            )
            projected_features[layer_name] = feature_tensor

    return projected_features