Skip to main content

Operator-typed layers

omnibias exposes the derivative tower as typed layers. Instead of "a dense layer followed by an activation", you declare what differential operator the layer should apply, and omnibias wires the closed-form tower to compute it.

The K-bias collapse unit (OMBU)

The Operator Multi-Bias Unit (OMBU) is a trainable unit that, for K bias replicas, collapses to a specific operator role at initialization. It is a drop-in for an activation:

from omnibias.torch import OMBU
import torch

ombu = OMBU(num_channels=4, K=2, base="tanh")
out = ombu(torch.zeros(8, 4))

The K=2 "bias collapse" output is what gives several activations their operator role — for example, the huber activation collapses to a soft-shrink (the proximal operator of the L1 norm), and arctan collapses to a Cauchy IRLS weight. See the Activation dictionary.

Two senses of "collapse"

This page uses "collapse" in its founding sense: the K biases coalesce and the unit becomes the derivative σ^(K-1) (the closed-form derivative tower). That is distinct from the β → ∞ "collapsed-bias" penalty in omnibias-convex / -control / -routing, which sharpens one constraint into a 0/1 feasibility step (an indicator, not a derivative). See the Glossary.

Lemma-1 initialization

By default the unit is initialized so it behaves as the base σ at step zero, no matter what K is. You get the inductive bias of the operator without changing the network's initial function.

OperatorBlock: a typed scalar operator

OperatorBlock dispatches on an op tag. The tag selects the order of the derivative tower the block contracts:

opComputesRequires fast path of order
identityσ(z)0
gradfirst derivative1
laplaciantrace of the Hessian2
derivativea chosen order nn
banda band of ordersup to the band max
integralan antiderivative-style operator
from omnibias.torch import OperatorBlock

laplacian = OperatorBlock(channels=8, op="laplacian", base="gaussian")
grad = OperatorBlock(channels=8, op="grad", base="tanh")

If you request an op whose required derivative order is not implemented for the chosen activation, the block raises a clear error rather than silently falling back.

Composable linear / conv layers

cmbLinear and cmbConv* are nn.Linear / nn.Conv* with an inline operator block, so you can drop them into an existing architecture:

from omnibias.torch import cmbLinear

fc = cmbLinear(in_features=128, out_features=64, op="identity", base="tanh")

A growable variant (GrowableOMBU) lets you increase K during training under a scheduler. The same surface exists in Keras 3 as cmbDense.

Why type the operator?

  • Correctness: the layer knows which derivative order it needs, so it can validate that the activation supports it.
  • Speed: the operator is a closed-form contraction of the tower, not a nested-autodiff graph.
  • Portability: the op tag is backend-agnostic. The same OperatorBlock(op="laplacian", base="gaussian") means the same thing — and produces the same bits — on PyTorch, JAX, and Keras 3.

Next