Skip to main content

A regional model that learns its boundaries

A soft partition assigns nonnegative weights to regions. These weights sum to one, so local expert predictions can be blended into one model. At finite inverse temperature, gradients can reach gates, thresholds, and experts.

Install​

python -m pip install --pre "omnibias-torch==0.5.0rc1" "omnibias-partition[torch]==0.1.0a2"

The current partition package is AGPL-3.0-or-later or commercial. Its license differs from the Apache-2.0 Torch backend. See licensing.

Trainable tensors​

Use partition_weights_arrays for raw trainable tensors. The convenience partition_weights interface converts a PartitionParams object and is not interchangeable with the raw-tensor training path.

soft_regions.py
"""Train the split and its experts together. Requires omnibias-partition."""
import torch
from omnibias.partition.torch.weights import combine, partition_weights_arrays

x = torch.linspace(-1, 1, 12, dtype=torch.float64).reshape(-1, 1)
W = torch.tensor([[0.8]], dtype=torch.float64, requires_grad=True)
threshold = torch.tensor([0.1], dtype=torch.float64, requires_grad=True)
slopes = torch.tensor([0.5, -0.2], dtype=torch.float64, requires_grad=True)
weights = partition_weights_arrays(W, threshold, x, beta=3.0, depth=1)
experts = x * slopes[None, :]
prediction = combine(weights, experts)
assert weights.shape == (12, 2)
torch.testing.assert_close(weights.sum(-1), torch.ones(12, dtype=x.dtype))
assert bool((weights >= 0).all())
loss = (prediction - x[:, 0].square()).square().mean()
loss.backward()
for parameter in (W, threshold, slopes):
    assert parameter.grad is not None and torch.isfinite(parameter.grad).all()
print("Partition of unity; gradients reach gates, thresholds, and experts.")

Download the complete script.

Save it as soft_regions.py and run python soft_regions.py. Expected output: Partition of unity; gradients reach gates, thresholds, and experts. The checks cover nonnegative normalized weights and finite gradients into all three parameter groups.

Can a piecewise model learn its boundary?​

At finite temperature, the soft gate provides a differentiable transition between regional experts. Its threshold and split direction can therefore learn through the blended loss. This does not differentiate a fixed hard discontinuity. Check the gate temperature and gradient magnitudes before interpreting a sharper partition as a better fitted model.

With depth one, there are two regions. weights has shape (12, 2); each row is a partition of unity. combine blends the two scalar expert outputs into one prediction per input.

Increasing beta sharpens routing away from ties. Very large inverse temperatures can reduce gradients away from boundaries. This is temperature hardening, distinct from the bias-collapse limit that extracts derivatives.

Read the weight and combination contracts before extending to vector-valued experts.

Solve a complete problem​

Continue with the executed notebook, including full initialization, training or optimization, numerical assertions, plots, and recorded results.