Skip to content

Hypernets

Full weight-generating teachers (vector- and map-head).

polyweave.hypernets

Hypernetwork teachers — modules that generate a student's target weights.

A teacher reads a prototype [1, C, H, W] and emits the weights for a target layer. Two head topologies, both available in vanilla and Sigma-Pi variants (toggle sigma_pi=True):

  • :class:ConvFilterTeachervector head: encode, global-average-pool, then a linear layer emits a flat parameter vector that a :class:~polyweave.targets. TargetSpec unpacks. Used for conv-filter generation (Experiment 2).
  • :class:FCMapTeacher / :class:QKMapTeacherspatial map head: the encoder preserves the prototype's spatial resolution and a conv weight-head emits weight maps aligned with that resolution, with a small pooled bias head. Used for FC head generation (Experiment 1) and Q/K generation (Experiment 3), where the prototype's spatial structure is the signal and must not be pooled away.

Every Sigma-Pi teacher exposes :meth:pi_scale_mean returning the diagnostic exp(pi_scale).mean() (None for vanilla teachers).

ConvFilterTeacher

ConvFilterTeacher(spec: TargetSpec, proto_channels: int = 4, width: int = 64, sigma_pi: bool = False, dropout: float = 0.0, center_product: bool = False)

Bases: _TeacherBase

Generate a flat parameter vector, unpacked by a :class:TargetSpec.

The encoder output is globally average-pooled to a single width-vector, then a linear head emits spec.num_params values that spec.unpack reshapes into the target's structured weights.

Parameters:

Name Type Description Default
spec TargetSpec

target specification (e.g. Conv2dTargetSpec) defining num_params and unpack.

required
proto_channels int

channels of the input prototype.

4
width int

encoder width.

64
sigma_pi bool

use a Sigma-Pi encoder (default False).

False
Source code in polyweave/hypernets/teachers.py
def __init__(
    self,
    spec: TargetSpec,
    proto_channels: int = 4,
    width: int = 64,
    sigma_pi: bool = False,
    dropout: float = 0.0,
    center_product: bool = False,
) -> None:
    super().__init__()
    self.spec = spec
    self.encoder = self._build_encoder(
        proto_channels, width, sigma_pi, dropout, center_product
    )
    self.pool = nn.AdaptiveAvgPool2d(1)
    self.head = nn.Linear(width, spec.num_params)

FCMapTeacher

FCMapTeacher(num_classes: int, feature_dim: int, proto_channels: int = 4, width: int = 64, sigma_pi: bool = False, dropout: float = 0.0, center_product: bool = False)

Bases: _TeacherBase

Generate a linear head {"weight": [K, F], "bias": [K]} (Experiment 1).

The prototype's spatial dims are (num_classes, feature_dim); the weight head emits a single map of exactly that shape, and a small pooled branch emits the per-class bias.

Parameters:

Name Type Description Default
num_classes int

rows of the linear head (and prototype height).

required
feature_dim int

input width of the linear head (and prototype width).

required
proto_channels int

channels of the input prototype.

4
width int

encoder width.

64
sigma_pi bool

use a Sigma-Pi encoder (default False).

False
Source code in polyweave/hypernets/teachers.py
def __init__(
    self,
    num_classes: int,
    feature_dim: int,
    proto_channels: int = 4,
    width: int = 64,
    sigma_pi: bool = False,
    dropout: float = 0.0,
    center_product: bool = False,
) -> None:
    super().__init__()
    self.num_classes = num_classes
    self.feature_dim = feature_dim
    self.encoder = self._build_encoder(
        proto_channels, width, sigma_pi, dropout, center_product
    )
    self.weight_head = nn.Conv2d(width, 1, 1)
    self.bias_head = nn.Sequential(
        nn.AdaptiveAvgPool2d((num_classes, 1)),
        nn.Conv2d(proto_channels, 1, 1),
    )

QKMapTeacher

QKMapTeacher(d_model: int, n_layers: int, proto_channels: int = 4, width: int = 64, sigma_pi: bool = False, out_scale: float = 0.1, dropout: float = 0.0, center_product: bool = False)

Bases: _TeacherBase

Generate per-layer query/key projections (Experiment 3).

The prototype is a stack of D x D cross-moment matrices; the weight head emits 2 * n_layers weight maps of shape D x D (a Wq and Wk per layer) and a pooled branch emits the matching biases. out_scale keeps the generated projections small at initialisation.

Returns a list of n_layers dicts {"q_weight", "q_bias", "k_weight", "k_bias"} — the format consumed by :class:~polyweave.targets.AttentionQKTargetSpec and the transformer student.

Source code in polyweave/hypernets/teachers.py
def __init__(
    self,
    d_model: int,
    n_layers: int,
    proto_channels: int = 4,
    width: int = 64,
    sigma_pi: bool = False,
    out_scale: float = 0.1,
    dropout: float = 0.0,
    center_product: bool = False,
) -> None:
    super().__init__()
    self.d_model = d_model
    self.n_layers = n_layers
    self.encoder = self._build_encoder(
        proto_channels, width, sigma_pi, dropout, center_product
    )
    self.weight_head = nn.Conv2d(width, 2 * n_layers, 3, padding=1)
    self.bias_head = nn.Linear(width, 2 * n_layers * d_model)
    self.out_scale = nn.Parameter(torch.tensor(float(out_scale)))