TetherGate — learned gates that the optimizer cannot close

Code only — no model weights.

Model card summary

What this repository is	A single-file PyTorch module: a learned gate for a network branch that the optimiser cannot push to zero. No model weights.
Developed by	Jakub Ćwirlej
Framework / hardware	Pure PyTorch; CPU and GPU
Status	Experimental: used in our own training runs, early observations promising; no published results
Intended use	Research on gated branches (adapters, mixture-of-experts style paths, auxiliary modules) whose gates tend to collapse during training
Out of scope	Commercial use (licence); drop-in use without checking the parameters on your own model (see Parameters and Pitfalls)
Licence	CC BY-NC 4.0

In plain words

Neural networks often add an optional path — an expert, an adapter, a new layer — scaled by a learned "gate" between 0 (off) and 1 (on). When the path does not help the training loss right away, the optimizer tends to close the gate slowly until the path contributes nothing, even if it would have been useful later. TetherGate keeps such a gate alive: early in training it is held open, then it is released gradually, and a spring keeps its average opening near a target you choose, while the gate can still open more for some inputs and less for others.

Status

Experimental. Tested in our own training runs; early observations are promising (the gate does not collapse under optimizer pressure). It is open in which architectures gating a branch this way improves model quality. Our long production run used wd = 0.5; the default wd = 20 in this file comes from its demo and example. Not a validated method.

What it does

g_raw = sigmoid( clamp( beta · (s − theta), arg_min, +8 ) )
g     = F(t) + (1 − F(t)) · g_raw              F(t): floor scheduled by the training STEP, not by the loss
after every optimizer.step():
    theta ← theta − k(t) · (theta − theta*)
    theta* = EMA(s) − logit(g_target) / beta   relative anchor: follows the signal's own drift
k = kappa_sleep while F = 1; afterwards k = lr · wd

• TetherGate(beta=6.0, g_target=0.5, theta_init=0.0, arg_min=-8.0, ema_momentum=0.99, detach_signal=True) — an nn.Module; g = gate(s) broadcasts like s.
• floor_schedule(step, ramp_start=3000, f_min=0.05, tau=1800, mode='exp', ramp_end=4500) — the floor F(t).
• tethergate_step(gates, step, lr, wd=20.0, kappa_sleep=5e-3, anchor='global', **floor_kw) — call after every optimizer.step(): sets the floor and pulls each theta towards its anchor. Returns the current floor.

s is any statistic you choose (per token, per sample or per layer) or a constant — with a constant, the gate becomes an anchored learnable scalar (ReZero / LayerScale style). Uses: expert gates and scales (MoE), adapters (LoRA), ReZero/LayerScale scalars, cross-attention / memory / retrieval branches, new blocks that hurt at initialisation.

Requirements

Pure PyTorch (no Triton, no CUDA required); CPU and GPU. Tested on PyTorch 2.12; uses only basic operations (the CPU tests also pass on 2.2).

Install

Copy tethergate.py next to your code.

git clone https://huggingface.co/JakCwi/tethergate

Quick start

import torch, torch.nn as nn
from tethergate import TetherGate, tethergate_step

d = 64
branch = nn.Linear(d, d)
gate = TetherGate(beta=6.0, g_target=0.5)
opt = torch.optim.AdamW([{"params": branch.parameters(), "weight_decay": 0.1},
                         {"params": [gate.theta], "weight_decay": 0.0}], lr=1e-3)   # theta: no weight decay

for step in range(1000):
    x = torch.randn(32, d)
    u = branch(x)
    s = u.abs().mean(-1, keepdim=True)            # control signal: branch activity per token
    out = x + gate(s) * u
    loss = (out - x.flip(-1)).pow(2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    tethergate_step([gate], step, lr=1e-3, wd=20.0, ramp_start=200, tau=200.0)   # AFTER opt.step()

python tethergate.py runs a minimal demo; example.py gates two residual adapters; demo_collapse.py is a toy test of gate collapse (no gate / free gate / TetherGate, CPU, seconds).

Parameters and how to choose them

parameter	default	how to choose
g_target	0.5	target mean opening, in (0, 1). 0.3–0.7 biases a branch towards closed / open.
beta	6.0	fixed sharpness; 4–8 for a signal of unit scale. Normalise s if its scale is very different (beta · std(s) of about 1–3). Do not learn it.
arg_min	−8	lower clamp of the sigmoid argument; leave it.
ema_momentum	0.99	smoothing of the signal mean used by the anchor.
detach_signal	True	s treated as data: the branch cannot learn to steer its own gate.
ramp_start, ramp_end	3000, 4500	in steps: until ramp_start the gate is fully open (F = 1); scale to your run — the branch should learn first.
f_min	0.05	minimum floor.
tau, mode	1800, 'exp'	exp: asymptotic decay with time constant tau steps (never reaches f_min); cos: reaches f_min at ramp_end.
wd	20	spring strength after the floor has decayed: k = lr · wd, time constant ≈ 1/(lr·wd) steps. Requires lr · wd ≪ 1 (at lr · wd ≥ 1 the spring overshoots). Under a steady closing push theta settles about 1/wd from its anchor, i.e. the gate argument shifts by about beta/wd: keep that well below 1 if the mean opening must stay near g_target. Smaller wd leaves the loss more freedom — our long production run used 0.5; 20 is the demo default.
kappa_sleep	5e-3	centring speed of theta while F = 1 (theta gets no loss gradient then).
anchor	'global'	global: one anchor for all gates, layers can differentiate; local: each gate to its own target, same opening everywhere.

Pitfalls and limitations

• Clamp saturation. The gradient to theta is exactly 0 whenever beta·(s − theta) lies outside the clamp range [arg_min, +8] (default −8 … +8). A gate pushed there gets no gradient from the loss; theta then returns only through the spring (this works, but more slowly).
• Order: tethergate_step must run AFTER optimizer.step(), every step; otherwise floor and spring do nothing.
• Weight decay on theta breaks the equilibrium: put theta in a group with weight_decay=0.
• The EMA starts empty (NaN): a gate that has not run a forward pass in train() mode is skipped by tethergate_step; its anchor exists after the first forward.
• Resuming: floor and ema_s are buffers in state_dict; load them with the model.
• Reopen a closed gate only with a ramp (raise the floor gradually): a step change on a mature model is a loss shock.
• The spring must keep up with signal drift: schedule k(t) smoothly, never in jumps.
• The anchor cannot restore information: if s loses its variance, the gate stays open but stops being selective. Monitor the variance of s.
• Spring too strong: with a large learning rate check lr · wd; reduce wd as it approaches 1.
• Performance: tethergate_step reads each EMA as a Python float — one GPU sync per gate per step. With hundreds of gates, vectorise it.
• Approaches that did not work for us (each fails on its own): a floor in the backward pass only; a self-referential mean in the forward; a hard lower clip (ratchet); a floor driven by the loss; a weak quadratic penalty around the target.

How to test

python test_tethergate.py   # unit tests, CPU, seconds
python demo_collapse.py    # toy collapse test

How it works in detail: TECHNICAL.md.

Files

tethergate.py (module and demo) · test_tethergate.py · demo_collapse.py (a toy demonstration of a plain gate collapsing vs TetherGate; not a benchmark) · example.py · TECHNICAL.md · LICENSE

Licence

Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0), © 2026 Jakub Ćwirlej — free for non-commercial use with attribution. Full text in LICENSE. Dependency (not included): PyTorch (BSD-3-Clause).

Citation and contact

Jakub Ćwirlej, TetherGate: relative-anchor spring gating for learned gates, 2026, https://huggingface.co/JakCwi/tethergate. Contact: https://jakubcwirlej.com.