Train through a branch#
A model with an
ifon a learnable threshold has a gradient of exactly zero for that threshold. One option,jump=, gives it back, and plain gradient descent then finds the switch.
Time: ~8 min · Runs on: CPU · You need: eagle, hawk, NumPy and matplotlib; read first: Gradients through conditionals, long runs and events, Gradients through a graph: eagle.diff
The sibling page Gradients through conditionals, long runs and events differentiates a stop that falls between two steps. Here the discontinuity is in the model itself: a value that jumps when a parameter crosses a data point. The two pages are one family: a gradient that has to cross a discontinuity.
Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.
A threshold the loss cannot see#
A sensor reads x; the model answers a below a threshold t and b above it. The data were made with t = 0.3, and we start from t = -0.4. As the kernel’s author you write the branch the way you would anywhere else:
import tempfile
import numpy as np
import eagle
import hawk
from hawk import Mutable, Param, Scalar
def make(**options):
@hawk.kernel(**options)
def switch(x: Scalar, t: Param, a: Param, b: Param, y: Mutable[Scalar]):
if x < t:
y = a
else:
y = b
return switch
n = 4096
x = np.random.default_rng(0).uniform(-1.0, 1.0, n)
A, B, T_TRUE = -0.5, 0.8, 0.3
target = np.where(x < T_TRUE, A, B)
def loss_and_grad(kernel):
"""A function t -> (mean squared error, dloss/dt), from one recorded launch and its reverse."""
(plan,) = eagle.deploy([kernel], targets=("host",), cache_dir=tempfile.mkdtemp())
def run(t):
y = np.zeros(n)
graph = eagle.Graph().add(plan.bind(x=x, t=t, a=A, b=B, y=y), name="s").build(record=True)
graph.launch()
grads = eagle.diff.Reverse(graph)(graph.record, {"s.y": 2 * (y - target) / n})
return float(np.mean((y - target) ** 2)), float(grads["t"])
return run
naive = loss_and_grad(make())
print("loss, d loss / d t at t = -0.4:", naive(-0.4))
loss, d loss / d t at t = -0.4: (0.5978540039062501, 0.0)
The loss is large and the gradient of t is exactly zero. That is correct for the program as written: t only chooses a branch, the branch values a and b do not depend on it, and moving t a little changes no output until it passes a sample. A gradient step would never move it.
Give the jump a width#
jump= smooths the switch for the derivative only. A sample within a few width of the threshold passes on the jump between the two branches, a - b, weighted by a bump of that width; farther samples cost one comparison. The kernel itself, its outputs and its digest are unchanged:
jumped = loss_and_grad(make(jump=0.05))
for t in (-0.4, 0.0, 0.3, 0.5):
(l, g0), (_, g1) = naive(t), jumped(t)
print(f"t = {t:4}: loss {l:.3f} plain gradient {g0:6.3f} jump= gradient {g1:6.3f}")
t = -0.4: loss 0.598 plain gradient 0.000 jump= gradient -0.842
t = 0.0: loss 0.243 plain gradient 0.000 jump= gradient -0.781
t = 0.3: loss 0.000 plain gradient 0.000 jump= gradient 0.000
t = 0.5: loss 0.164 plain gradient 0.000 jump= gradient 0.756
Same loss, same kernel, but now the gradient points downhill: negative left of the answer (raise t), positive right of it, and zero at t = 0.3 where every sample is right. The picture shows both across a sweep of t:
import matplotlib.pyplot as plt
from themed_figure import show
ts = np.linspace(-0.8, 0.8, 81)
plain = np.array([naive(t) for t in ts])
smooth = np.array([jumped(t) for t in ts])
fig, (ax, bx) = plt.subplots(1, 2, figsize=(9, 3.4))
ax.plot(ts, plain[:, 0])
ax.axvline(T_TRUE, color="#898781", ls=":")
ax.set(xlabel="threshold t", ylabel="loss", title="the loss has a clear minimum")
bx.plot(ts, plain[:, 1], "--", label="plain")
bx.plot(ts, smooth[:, 1], label="jump=0.05")
bx.axvline(T_TRUE, color="#898781", ls=":")
bx.set(xlabel="threshold t", ylabel="d loss / d t", title="only one gradient sees it")
bx.legend()
plt.tight_layout()
show()
Train it#
Plain gradient descent on t, starting well away from the answer:
from themed_figure import show
t, history = -0.4, []
for step in range(16):
l, g = jumped(t)
history.append((t, l))
t -= 0.1 * g
print(f"t after 16 steps: {t:.4f} (data made with {T_TRUE})")
fig, ax = plt.subplots(figsize=(5, 3.2))
ax.plot([h[0] for h in history], "o-")
ax.axhline(T_TRUE, color="#898781", ls=":")
ax.set(xlabel="training step", ylabel="threshold t", title="the threshold walks to 0.3")
plt.tight_layout()
show()
t after 16 steps: 0.3003 (data made with 0.3)
The threshold crosses the data in about eight steps and settles at the generating value. Each step cost one launch and one reverse, the same as any other trained parameter.
The width is the one choice to make. A wider width reaches a threshold that is further away and gives a smoother gradient; a narrower one is sharper but sees only samples close to the switch. jump="auto" picks it from the batch, and jump="h" reads it from a Param you can change between launches without a rebuild. A comparison with ==, or a break that depends on the compared value, is refused by name. The full list is in the hawk docs; the stop-time version of the same problem is in Gradients through conditionals, long runs and events, with its reference in Event refinement.