Hard decisions, explicit surrogate gradients
Use omnibias-binary when the forward computation must make a discrete decision, while training uses a deliberately chosen smooth surrogate. This differs from a soft regional model, whose forward values remain smooth at finite temperature.
binarize(z, beta) returns a hard sign, including +1 at zero. Its backward rule uses the derivative of tanh(beta*z). That supplied gradient is not the classical derivative of the hard sign. Increasing beta narrows the region carrying gradients; it does not change the hard forward values. Very large beta can suppress useful gradients away from the boundary.
Install the alpha and your backend
Install omnibias-binary==0.1.0a2 together with omnibias-torch==0.5.0rc1 or omnibias-jax==0.5.0rc1. The scripts below include exact commands. This package uses Apache-2.0; it is an alpha API.
Choose the representation
| Operation | Forward output | Reverse-mode rule |
|---|---|---|
binarize | −1 or +1; zero maps to +1 | Tanh surrogate slope |
binarize01, heaviside | 0 or 1; zero maps to 1 | Sigmoid surrogate slope |
ternarize | −1, 0, +1 with a closed dead zone | Average of two shifted tanh slopes |
kbit_quantize | Rounded, clipped uniform levels | Sum of shifted tanh slopes |
binarize_curvature | Same hard sign | Truncated window-average slope; beta fixed |
For k-bit quantization, require hi > lo, integer bits >= 1, and a positive finite beta. The backward path processes 2**bits - 1 thresholds; increasing bits is not a constant-cost operation. Not all invalid-domain inputs are rejected explicitly by the alpha wrapper.
Run the complete example
The exact downloadable files are displayed here. Both were executed against the pinned installed distributions on CPU in float64. Each checks hard values, shapes, finite input gradients, and the defining analytic surrogate formula. The scripts also exercise every quantizer and surrogate helper shown in the table and text.
- PyTorch
- JAX
"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.
python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
binarize, binarize01, heaviside, ternarize, kbit_quantize,
riccati_tanh_derivative, riccati_sigmoid_derivative,
surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)
z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0
t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))
schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")
"""JAX hard quantization with explicitly tested surrogate cotangents.
python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
binarize, binarize01, heaviside, ternarize, kbit_quantize,
riccati_tanh_derivative, riccati_sigmoid_derivative,
surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)
z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0
t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")
JAX uses custom VJPs for the hard quantizers. These examples use reverse-mode jax.grad; a custom VJP does not establish support for a general forward-mode jax.jvp transformation.
Both scripts report quantizer shape (5,) and tower/jet shape (4, 5) after the assertions pass.
Raw derivatives and Taylor coefficients
surrogate_tower(z, beta, order) returns raw derivatives of tanh(beta*z). surrogate_jet returns the same derivatives divided by their factorials. For every row k, jet[k] * k! == tower[k] up to floating-point tolerance. In particular, jet[2] is half the raw second derivative.
riccati_tanh_derivative expects tanh values, not raw z. Its sigmoid counterpart expects sigmoid values. Both evaluate polynomials in those activation values; passing preactivations changes the mathematical quantity being computed.
Anneal deliberately
BetaAnnealScheduler.step() returns the current beta and then advances its counter. value(step) is a stateless lookup; positions outside the schedule interval are clamped to endpoint values. reset() restarts only the scheduler counter. The PyTorch script asserts these semantics, but the scheduler itself is pure Python and works with either backend.
Choose whether one call represents a batch, optimizer update, or epoch. Track loss, gradient scale, and saturation while changing beta. Annealing does not guarantee convergence, and a curvature-corrected surrogate is an approximation to a window average rather than an exact Hessian update.
Continue with the binary package reference and scope and guarantees.