A model's local rules and the regions in which they apply are often designed in separate steps. Fixing the regions first can freeze a useful modelling decision before the regional experts have learned anything from the objective.
Make the routing part of the model
omnibias soft partitions produce nonnegative routing weights that sum to one. These weights blend local expert outputs into a smooth prediction. On supported paths, gradients flow through split directions, thresholds, and expert parameters, making joint optimization possible.
Our notebook demonstrates this on a manufactured two-regime regression problem. It trains a split and two linear experts together, checks finite gradients for all parameter groups, and evaluates predictions on a held-out grid. The data is synthetic and intentionally simple enough to make the model's behaviour inspectable.
Hardening deserves a separate evaluation
Increasing inverse temperature makes a soft gate sharper. It also changes the represented function. In this example, hardening the trained gate increases prediction error. The notebook publishes both errors instead of assuming that a sharper decision is always a better one.
That distinction matters when designing training schedules or deployment routing. Temperature hardening is also different from bias collapse, the separate mathematical limit used to extract derivatives. A smooth surrogate does not give an ordinary derivative to an existing hard discontinuity.
Where regional structure can help
Regime-dependent responses, local surrogate models, and piecewise scientific approximations are natural questions to investigate with trainable regions. Any claimed advantage needs representative data and comparison against appropriate baselines. This small demonstration is a starting point for that investigation, not evidence of an industrial deployment.
Engineering and licensing
Derivon can help formulate the regional model, choose evaluation criteria, test hardening schedules, and integrate the resulting workflow. The partition extension is currently alpha and uses AGPL-3.0-or-later or commercial licensing, distinct from the Apache-2.0 Torch substrate. Package choice and licensing belong in the implementation decision.
Reproduce the result
The recorded notebook reports held-out rmse of 3.64e-05. Inspect the conditions and raw data before comparing it with another method.
Read and run the complete notebook · Discuss a scientific ML project