example.py

# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Usage example: gating a new branch (here: a residual adapter) with TetherGate.

Integration checklist:
  1. one TetherGate per gated branch; theta goes into a parameter group WITHOUT weight decay;
  2. out = x + gate(s) * branch(x), with s any per-token / per-sample statistic (or a constant);
  3. call tethergate_step(gates, step, lr, ...) right AFTER optimizer.step(), every step.

    python example.py
"""
import torch
import torch.nn as nn

from tethergate import TetherGate, tethergate_step


class GatedAdapter(nn.Module):
    def __init__(self, d, r=8):
        super().__init__()
        self.down = nn.Linear(d, r, bias=False)
        self.up = nn.Linear(r, d, bias=False)
        self.gate = TetherGate(beta=6.0, g_target=0.5)

    def forward(self, x):
        u = self.up(self.down(x))
        s = u.detach().norm(dim=-1, keepdim=True) / u.shape[-1] ** 0.5   # example signal: branch activity
        return x + self.gate(s) * u


if __name__ == '__main__':
    torch.manual_seed(0)
    d, lr, steps = 64, 3e-3, 600
    model = nn.Sequential(GatedAdapter(d), GatedAdapter(d))
    gates = [m.gate for m in model]
    gate_params = {id(g.theta) for g in gates}
    opt = torch.optim.AdamW([
        {'params': [p for p in model.parameters() if id(p) not in gate_params], 'weight_decay': 0.1},
        {'params': [g.theta for g in gates], 'weight_decay': 0.0}], lr=lr)
    target = torch.randn(d, 8) @ torch.randn(8, d) / d       # low rank: an adapter can fit it
    for step in range(steps):
        x = torch.randn(128, d)
        loss = (model(x) - (x + x @ target.t())).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        # k = lr * wd sets the spring's time constant (about 1 / (lr * wd) steps)
        f = tethergate_step(gates, step, lr=lr, wd=20.0, kappa_sleep=0.05, ramp_start=150, tau=100.0)
        if step % 150 == 149:
            print(f'step {step + 1}: loss {loss.item():.4f} | floor {f:.3f} | '
                  f'theta {[round(g.theta.item(), 3) for g in gates]}')