Students

The networks whose weights a teacher generates (CNN, transformer, and Sigma-Pi variants).

polyweave.students

Student networks whose weights a hypernetwork teacher generates.

Two families, mirroring the paper's experiments:

• :class:CNNStudent — a small CIFAR-style convolutional classifier with its first conv layer (conv1/bn1/pool1) factored out so a teacher can replace just those weights (Experiment 2), while still allowing a generated linear head (Experiment 1). Three concrete architectures (A/B/C) provide the cross-architecture seen/unseen split.
• :class:TinyTransformerStudent — an attention-only transformer for the synthetic relational-lookup task whose per-layer Q/K projections a teacher generates (Experiment 3).

CNNStudent

CNNStudent(trunk_rest: Module, feature_dim: int = 256, num_classes: int = 10, in_ch: int = 3, conv1_out: int = 32, kernel_size: int = 3)

Bases: Module

A convolutional classifier with a replaceable first conv layer.

Parameters:

Name	Type	Description	Default
trunk_rest	Module	the body of the network mapping the pooled conv1 output [B, conv1_out, H/2, W/2] to a flat feature vector of size feature_dim (typically ending in Flatten/Linear/ReLU).	required
feature_dim	int	width of the penultimate feature vector.	256
num_classes	int	number of output classes.	10
in_ch	int	input image channels (3 for CIFAR).	3
conv1_out	int	output channels of conv1 (the generated target width).	32
kernel_size	int	conv1 kernel size.	3

Source code in polyweave/students/cnn.py

def __init__(
    self,
    trunk_rest: nn.Module,
    feature_dim: int = 256,
    num_classes: int = 10,
    in_ch: int = 3,
    conv1_out: int = 32,
    kernel_size: int = 3,
) -> None:
    super().__init__()
    self.in_ch = in_ch
    self.conv1_out = conv1_out
    self.kernel_size = kernel_size
    self.padding = kernel_size // 2
    self.feature_dim = feature_dim
    self.num_classes = num_classes

    self.conv1 = nn.Conv2d(in_ch, conv1_out, kernel_size, padding=self.padding)
    self.bn1 = nn.BatchNorm2d(conv1_out)
    self.pool1 = nn.MaxPool2d(2)
    self.trunk_rest = trunk_rest
    self.fc = nn.Linear(feature_dim, num_classes)

SigmaPiStudent

SigmaPiStudent(width: int = 32, num_classes: int = 10, in_ch: int = 3, kernel_size: int = 3, img_size: int = 32)

Bases: Module

A classifier with a generated multiplicative (Sigma-Pi) target block.

Parameters:

Name	Type	Description	Default
width	int	channel width of the stem output and the Sigma-Pi block (the block is channels-preserving, so this is the generated target width).	32
num_classes	int	number of output classes.	10
in_ch	int	input image channels (3 for CIFAR-like data).	3
kernel_size	int	kernel size of both the stem and the Sigma-Pi block.	3
img_size	int	spatial size of the (square) input, used to size the head.	32

Source code in polyweave/students/sigmapi.py

def __init__(
    self,
    width: int = 32,
    num_classes: int = 10,
    in_ch: int = 3,
    kernel_size: int = 3,
    img_size: int = 32,
) -> None:
    super().__init__()
    self.width = width
    self.num_classes = num_classes
    self.in_ch = in_ch
    self.kernel_size = kernel_size
    self.img_size = img_size
    padding = kernel_size // 2

    # Fixed stem maps input channels up to the block width.
    self.stem = nn.Sequential(
        nn.Conv2d(in_ch, width, kernel_size, padding=padding),
        nn.BatchNorm2d(width),
        nn.ReLU(),
    )
    # The multiplicative target block (weights generated by the teacher).
    self.sigmapi = ConvSigmaPi2d(width, kernel_size=kernel_size, padding=padding)
    self.pool = nn.AdaptiveAvgPool2d((4, 4))
    self.fc = nn.Linear(width * 4 * 4, num_classes)

pi_scale_mean

pi_scale_mean() -> float

Recruitment diagnostic of the student's own Sigma-Pi block.

Source code in polyweave/students/sigmapi.py

@torch.no_grad()
def pi_scale_mean(self) -> float:
    """Recruitment diagnostic of the student's own Sigma-Pi block."""
    return self.sigmapi.pi_scale_mean()

TinyTransformerStudent

TinyTransformerStudent(vocab_size: int = 64, seq_len: int = 10, num_classes: int = 5, d_model: int = 64, n_heads: int = 4, n_layers: int = 2, activation: str = 'relu', emb_seed: int = 7, dropout: float = 0.0)

Bases: Module

Attention-only transformer with shared frozen embeddings.

Parameters:

Name	Type	Description	Default
vocab_size	int	token vocabulary size.	64
seq_len	int	sequence length (query token is at the last position).	10
num_classes	int	number of key slots / output classes.	5
d_model	int	model width (uniform across students so the generated Q/K target is fixed-size).	64
n_heads	int	number of attention heads.	4
n_layers	int	number of transformer blocks.	2
activation	str	per-architecture pointwise activation in each block.	'relu'
emb_seed	int	seed for the shared, frozen token/positional embeddings.	7
dropout	float	attention dropout (default 0).	0.0

Source code in polyweave/students/transformer.py

def __init__(
    self,
    vocab_size: int = 64,
    seq_len: int = 10,
    num_classes: int = 5,
    d_model: int = 64,
    n_heads: int = 4,
    n_layers: int = 2,
    activation: str = "relu",
    emb_seed: int = 7,
    dropout: float = 0.0,
) -> None:
    super().__init__()
    self.vocab_size = vocab_size
    self.seq_len = seq_len
    self.num_classes = num_classes
    self.d_model = d_model
    self.n_heads = n_heads
    self.n_layers = n_layers
    self.activation = activation

    self.token_emb = nn.Embedding(vocab_size, d_model)
    gen = torch.Generator().manual_seed(emb_seed)
    with torch.no_grad():
        self.token_emb.weight.copy_(torch.randn(vocab_size, d_model, generator=gen) * 0.02)
        pos = torch.randn(1, seq_len, d_model, generator=gen) * 0.02
    self.token_emb.weight.requires_grad_(False)
    self.register_buffer("pos_emb", pos)

    self.blocks = nn.ModuleList(
        [TinySelfAttentionBlock(d_model, n_heads, activation, dropout) for _ in range(n_layers)]
    )
    self.classifier = nn.Linear(d_model, num_classes)

make_cnn_student

make_cnn_student(arch: str = 'A', *, feature_dim: int = 256, num_classes: int = 10, in_ch: int = 3, conv1_out: int = 32, kernel_size: int = 3) -> CNNStudent

Build one CNN student of architecture arch in {"A", "B", "C"}.

The trunks differ in width/depth but all consume the same fixed conv1 output (assuming a 32x32 input that has been pooled to 16x16 then reduced to 4x4 by the trunk's two pooling stages).

Source code in polyweave/students/cnn.py

def make_cnn_student(
    arch: str = "A",
    *,
    feature_dim: int = 256,
    num_classes: int = 10,
    in_ch: int = 3,
    conv1_out: int = 32,
    kernel_size: int = 3,
) -> CNNStudent:
    """Build one CNN student of architecture ``arch`` in ``{"A", "B", "C"}``.

    The trunks differ in width/depth but all consume the same fixed ``conv1``
    output (assuming a 32x32 input that has been pooled to 16x16 then reduced to
    4x4 by the trunk's two pooling stages).
    """
    key = arch.upper()
    if key not in _TRUNKS:
        raise ValueError(f"unknown arch {arch!r}; choose from {sorted(_TRUNKS)}")
    trunk = _TRUNKS[key](conv1_out, feature_dim)
    return CNNStudent(
        trunk, feature_dim=feature_dim, num_classes=num_classes,
        in_ch=in_ch, conv1_out=conv1_out, kernel_size=kernel_size,
    )

make_cnn_students

make_cnn_students(archs: Optional[List[str]] = None, **kwargs) -> List[CNNStudent]

Build one student per architecture name (defaults to A, B, C).

Source code in polyweave/students/cnn.py

def make_cnn_students(
    archs: Optional[List[str]] = None, **kwargs
) -> List[CNNStudent]:
    """Build one student per architecture name (defaults to A, B, C)."""
    if archs is None:
        archs = ["A", "B", "C"]
    return [make_cnn_student(a, **kwargs) for a in archs]

freeze_except_qk

freeze_except_qk(model: TinyTransformerStudent) -> None

Freeze everything except every block's packed in_proj (Q/K/V).

Source code in polyweave/students/transformer.py

def freeze_except_qk(model: TinyTransformerStudent) -> None:
    """Freeze everything except every block's packed in_proj (Q/K/V)."""
    for p in model.parameters():
        p.requires_grad_(False)
    for attn in attn_layers(model):
        attn.in_proj_weight.requires_grad_(True)
        attn.in_proj_bias.requires_grad_(True)

mask_qk_grads

mask_qk_grads(model: TinyTransformerStudent) -> None

Zero the V slice of every block's packed projection gradient.

Lets an optimiser step the Q/K rows while leaving the (frozen) value projection untouched, even though Q/K/V share one packed parameter.

Source code in polyweave/students/transformer.py

def mask_qk_grads(model: TinyTransformerStudent) -> None:
    """Zero the V slice of every block's packed projection gradient.

    Lets an optimiser step the Q/K rows while leaving the (frozen) value
    projection untouched, even though Q/K/V share one packed parameter.
    """
    for attn in attn_layers(model):
        D = attn.embed_dim
        if attn.in_proj_weight.grad is not None:
            attn.in_proj_weight.grad[2 * D : 3 * D].zero_()
        if attn.in_proj_bias is not None and attn.in_proj_bias.grad is not None:
            attn.in_proj_bias.grad[2 * D : 3 * D].zero_()

reinit_qk

reinit_qk(model: TinyTransformerStudent) -> None

Xavier-reinitialise the Q and K slices of every block (V left as-is).

Source code in polyweave/students/transformer.py

@torch.no_grad()
def reinit_qk(model: TinyTransformerStudent) -> None:
    """Xavier-reinitialise the Q and K slices of every block (V left as-is)."""
    for attn in attn_layers(model):
        D = attn.embed_dim
        nn.init.xavier_uniform_(attn.in_proj_weight[:D])
        nn.init.xavier_uniform_(attn.in_proj_weight[D : 2 * D])
        attn.in_proj_bias[: 2 * D].zero_()