File size: 2,098 Bytes
d7070c6 | 1 2 3 4 5 6 7 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 | # 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]}')
|