NSC64/symmetry-transport-per-task-lora

Per-task LoRA checkpoints from a symmetry-transport hypernetwork

Twenty trained checkpoints in one PyTorch file, one per synthetic task. The hypernetwork generates a rank-two LoRA head per task, then that head is trained and frozen. Base and adapter both ran 600 steps in float32 on a CUDA GPU.

These are synthetic numbers. No language model and no real dataset is involved, and the .pt files are neither Transformers nor PEFT adapters.

Files

• per_task_lora_pack.pt: all 20 checkpoints, with per-task accuracy and validation BCE.
• manifest.json: SHA-256 of every checkpoint in the training output, plus held-out numbers.

The model definition isn't a separate file. It's inlined in the usage section below, so everything you need is in this repository.

What the tasks are

Every task is binary classification in 2D. Each example has an input x of shape [batch, 2] and a context c of shape [batch, 3].

The first two context columns are a unit direction (n_x, n_y), the third is a threshold tau. The input rotates into that frame first, so u is the component of x along the direction and v the component across it. Label is 1 when the family's margin is non-negative.

Only the margin differs between families:

Task family	Margin
linear	u - tau
curved	u - tau + 0.55 * (v**2 - 1/3) + 0.4 * sin(pi * tau) * v
axis	the curved margin plus 0.6 * x[0] * x[1]
high_frequency	u - tau + 0.25 * (v**2 - 1/3) + 0.30 * sin(4 * pi * v)

linear is one straight boundary. curved bends it with a quadratic in v and tilts it by tau. axis adds a term in the unrotated coordinates, which is the whole point of the family: the boundary no longer exists in the (u, v) frame, so the model has to transport the context somewhere the rule does apply. high_frequency is the control, oscillating four times faster in v, which 8 hidden units struggle to fit.

-task-11 through -task-15 are five seeds of one rule, not five rules. Each seed draws its own points (seed = 101 + 1009 * task_seed) over the same grid of 8 orientations by 5 thresholds. Four families times five seeds gives the 20.

Results

Test accuracy per family, mean over five seeds:

Task family	Mean test accuracy	Sample SD
Linear	99.23%	0.42
Curved	98.10%	0.27
Axis	97.09%	0.42
High-frequency	89.45%	0.89

Checkpoint selection used validation BCE. high_frequency is the hardest family by design; it's the control showing where this breaks.

One 295-parameter model trained jointly on all 20 tasks did not match the per-task checkpoints on axis. The checkpoints stayed separate rather than being distilled or averaged together.

Format

Target is a 2-8-1 tanh MLP. Each checkpoint holds 295 trainable parameters: 173 in the canonical base hypernetwork, 122 in the generated rank-two LoRA head.

per_task_lora_pack.pt is a dictionary under the key tasks, keyed by task name. Each entry carries a flat state_dict (bias, u, v, a, base_generator.*) plus family, task_seed, task_id, parameters, test, and validation BCE. Nothing is shared or dropped, so the pack costs exactly what 20 separate checkpoints cost, as one download.

Usage

Each entry under tasks holds a flat state_dict, which is a plain dict of tensor names to tensors. This snippet rebuilds the model and loads one task:

import torch

H, R, P = 8, 2, 33


class TaskLoRA(torch.nn.Module):
    def __init__(self):
        super().__init__()
        q = 2 * H * R + 3 * R
        self.bias = torch.nn.Parameter(torch.zeros(q))
        self.u = torch.nn.Parameter(torch.zeros(q, 2))
        self.v = torch.nn.Parameter(torch.zeros(2, 3))
        self.a = torch.nn.Parameter(torch.zeros(2))
        g = self.base_generator = torch.nn.Module()
        g.b = torch.nn.Parameter(torch.zeros(P))
        g.u = torch.nn.Parameter(torch.zeros(P, 4))
        g.v = torch.nn.Parameter(torch.zeros(4, 1))
        g.a = torch.nn.Parameter(torch.zeros(4))

    def forward(self, c, x):
        g = self.base_generator
        z = self.bias + torch.tanh(c @ self.v.T + self.a) @ self.u.T
        end = H * R
        a1 = z[:, :end].reshape(-1, H, R)
        b1 = z[:, end:end + 2 * R].reshape(-1, R, 2)
        a2 = z[:, end + 2 * R:end + 2 * R + end].reshape(-1, H, R)
        b2 = z[:, end + 2 * R + end:].reshape(-1, R, 1)

        theta = g.b + torch.tanh(c[:, 2:3] @ g.v.T + g.a) @ g.u.T
        w, b, o, d = (theta[:, :16].reshape(-1, H, 2), theta[:, 16:24],
                      theta[:, 24:32], theta[:, 32])

        n = c[:, :2]
        rot = torch.stack((n[:, 0], n[:, 1], -n[:, 1], n[:, 0]), -1).reshape(-1, 2, 2)
        uv = torch.bmm(rot, x[..., None]).squeeze(-1)
        o = o + torch.bmm(a2, b2).squeeze(-1)
        return (o * torch.tanh(torch.einsum("bhd,bd->bh", w, uv)
                               + torch.einsum("bhd,bd->bh", torch.bmm(a1, b1), x)
                               + b)).sum(1) + d


pack = torch.load("per_task_lora_pack.pt", map_location="cpu", weights_only=False)
task = pack["tasks"]["axis-task-11"]

model = TaskLoRA().double()
model.load_state_dict({k: v.double() for k, v in task["state_dict"].items()}, strict=True)
model.eval()

print("parameters:", sum(p.numel() for p in model.parameters()))
print("stored test accuracy:", task["test"]["accuracy"])

Expected output:

parameters: 295
stored test accuracy: 96.875

Call model(c, x) with x of shape [batch, 2] and c of shape [batch, 3] to get one logit per example. The first two context columns are the direction (n_x, n_y) and the third is tau.

Provenance

Author: NSC64.