omnibias.torch.architectures.attention
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.architectures.attentionModule contract & mathematical conventions
Non-local attention as a closed-form block inside the jet chain (torch).
Every architecture in omnibias.torch.architectures.pinn is *local* in the
same narrow sense: the value at x is a chain of elementwise activations
applied to affine maps of x. This module adds the first **non-local** block
that keeps the closed-form tower intact.
u(x) = W_o\,\Big[\underbrace{\mathrm{softmax}\big(\beta\, q(x) K^\top\big) V} _{\text{global mixture over the memory}} + q(x)\Big] + b_o, \qquad q(x) = \mathrm{MLP}(x).
The softmax couples every memory slot to every other one through a shared
denominator, so u(x) depends on the *whole* memory at once -- the value at a
point is no longer a function of a purely pointwise feature stack.
What is new here, and what is not
omnibias.hopfield already carries the closed-form log-sum-exp Jacobian and
Hessian of this same block, but **with respect to the scores**. A PDE residual
needs d/dx. omnibias.torch.jet_mv.jet_attention supplies exactly that
missing coordinate story by pushing the query *jet* through the block, and this
module is the trainable network that uses it: one
~omnibias.torch.jet_mv.mlp_jet_mv call for the encoder, one attention
block, one affine readout, and every mixed partial D^alpha u(x) falls out of
that single jet. There is no torch.autograd.grad in the differential
operator, at any order.
The block is exact because softmax factors into primitives that are each exact:
exp (a registered activation with the tower exp^(k) = exp), a linear sum,
a reciprocal (tower (-1)^k k! u^{-(k+1)}), and a jet product. That is a
genuine extension of the reachable class -- compose_jet_mv alone handles only
elementwise maps.
Because AttentionJetMLP is a
~omnibias.torch.architectures.pinn._JetMLPCore, it inherits the exact
readouts (jet / gradient / hessian / partials) and drops into the
omnibias PINN field substrate via omnibias.pinn.torch.fields.attention.
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
AttentionJetMLP
omnibias.torch.architectures.AttentionJetMLPSource ↗class AttentionJetMLP(_JetMLPCore)Deep encoder + one non-local attention block, with exact input derivatives.
u(x) = W_o [ softmax(beta q(x) K^T) V (+ q(x)) ] + b_o where q is an
ordinary MLP encoder. The memory (K, V) is a set of trainable slots that
do **not** depend on x, which is what keeps the block closed form: the
only x-dependence enters through the query, and the jet of the query is
exact.
Interpretation for PINNs: the softmax is a learned soft partition of the
input domain, and V holds one local model per region -- a differentiable,
globally-coupled relative of a domain decomposition, and the reason this
field type is worth having next to the local ones. Note the temperature
collapse this shares with the rest of the repo: as beta -> inf the
mixture hardens into a crisp nearest-slot assignment (the *feasibility*
sense of collapse, not the founding delta -> 0 one).
Parameters
in_dim, out_dim: Input coordinate count and output component count.
hidden: Encoder width. Also the query / key width d_key.
depth: Number of hidden (activated) encoder layers >= 1.
base: Encoder activation; must have a closed-form derivative fast path.
memory: Number of memory slots n.
value_dim: Width d_val of the value slots; defaults to hidden. Must equal hidden when residual is set.
beta: Initial inverse temperature beta > 0. Larger is sharper.
learnable_temperature: Train beta alongside the weights. It enters the jet as a scale on the keys, so the derivatives stay closed form while it moves.
residual: Add the query to the attention output (the usual transformer skip). Keeps a *local* path alongside the non-local one, which trains far better.
seed: Optional seed for the memory draw; None leaves the global RNG alone.
dtype: Parameter dtype, defaulting to the framework default.
Parameters
| Name | Type & default | Contract |
|---|---|---|
in_dimpositional-or-keyword | intRequired | Input coordinate count and output component count. |
hiddenpositional-or-keyword | intRequired | Encoder width. Also the query / key width d_key. |
out_dimpositional-or-keyword | intDefault: 1 | See the operation contract and implementation below. |
depthpositional-or-keyword | intDefault: 2 | Number of hidden (activated) encoder layers >= 1. |
basepositional-or-keyword | str | ActivationSpec[Tensor]Default: 'tanh' | Encoder activation; must have a closed-form derivative fast path. |
memorykeyword-only | intDefault: 16 | Number of memory slots n. |
value_dimkeyword-only | int | NoneDefault: None | Width d_val of the value slots; defaults to hidden. Must equal hidden when residual is set. |
betakeyword-only | floatDefault: 1.0 | Initial inverse temperature beta > 0. Larger is sharper. |
learnable_temperaturekeyword-only | boolDefault: False | Train beta alongside the weights. It enters the jet as a scale on the keys, so the derivatives stay closed form while it moves. |
residualkeyword-only | boolDefault: True | Add the query to the attention output (the usual transformer skip). Keeps a *local* path alongside the non-local one, which trains far better. |
seedkeyword-only | int | NoneDefault: None | Optional seed for the memory draw; None leaves the global RNG alone. |
dtypekeyword-only | torch.dtype | NoneDefault: None | Parameter dtype, defaulting to the framework default. |
Fields & methods
beta_rawfield
beta_raw: TensorDeclared Tensor field on AttentionJetMLP. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗__init__method
__init__(self, in_dim: int, hidden: int, out_dim: int=1, depth: int=2, base: str | ActivationSpec[Tensor]='tanh', *, memory: int=16, value_dim: int | None=None, beta: float=1.0, learnable_temperature: bool=False, residual: bool=True, seed: int | None=None, dtype: torch.dtype | None=None) -> NoneSee the class contract above and the exact method implementation below.
Parameters
| Name | Type & default | Contract |
|---|---|---|
in_dimpositional-or-keyword | intRequired | See the operation contract and implementation below. |
hiddenpositional-or-keyword | intRequired | See the operation contract and implementation below. |
out_dimpositional-or-keyword | intDefault: 1 | See the operation contract and implementation below. |
depthpositional-or-keyword | intDefault: 2 | See the operation contract and implementation below. |
basepositional-or-keyword | str | ActivationSpec[Tensor]Default: 'tanh' | See the operation contract and implementation below. |
memorykeyword-only | intDefault: 16 | See the operation contract and implementation below. |
value_dimkeyword-only | int | NoneDefault: None | See the operation contract and implementation below. |
betakeyword-only | floatDefault: 1.0 | See the operation contract and implementation below. |
learnable_temperaturekeyword-only | boolDefault: False | See the operation contract and implementation below. |
residualkeyword-only | boolDefault: True | See the operation contract and implementation below. |
seedkeyword-only | int | NoneDefault: None | See the operation contract and implementation below. |
dtypekeyword-only | torch.dtype | NoneDefault: None | See the operation contract and implementation below. |
Returns: None
Validation & failure cases
- When
in_dim < 1, raisesValueError(f'in_dim must be >= 1, got {in_dim}'). - When
hidden < 1, raisesValueError(f'hidden must be >= 1, got {hidden}'). - When
out_dim < 1, raisesValueError(f'out_dim must be >= 1, got {out_dim}'). - When
depth < 1, raisesValueError(f'depth (number of hidden layers) must be >= 1, got {depth}'). - When
memory < 1, raisesValueError(f'memory must be >= 1, got {memory}'). - When
beta <= 0.0, raisesValueError(f'beta must be > 0, got {beta}'). - When
d_val < 1, raisesValueError(f'value_dim must be >= 1, got {d_val}'). - When
residual and d_val != hidden, raisesValueError(f'residual needs value_dim == hidden, got {d_val} != {hidden}; pass residual=False for an asymmetric block').
Method implementation
def __init__(
self,
in_dim: int,
hidden: int,
out_dim: int = 1,
depth: int = 2,
base: str | ActivationSpec[Tensor] = "tanh",
*,
memory: int = 16,
value_dim: int | None = None,
beta: float = 1.0,
learnable_temperature: bool = False,
residual: bool = True,
seed: int | None = None,
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
if in_dim < 1:
raise ValueError(f"in_dim must be >= 1, got {in_dim}")
if hidden < 1:
raise ValueError(f"hidden must be >= 1, got {hidden}")
if out_dim < 1:
raise ValueError(f"out_dim must be >= 1, got {out_dim}")
if depth < 1:
raise ValueError(f"depth (number of hidden layers) must be >= 1, got {depth}")
if memory < 1:
raise ValueError(f"memory must be >= 1, got {memory}")
if beta <= 0.0:
raise ValueError(f"beta must be > 0, got {beta}")
d_val = hidden if value_dim is None else int(value_dim)
if d_val < 1:
raise ValueError(f"value_dim must be >= 1, got {d_val}")
if residual and d_val != hidden:
raise ValueError(
f"residual needs value_dim == hidden, got {d_val} != {hidden}; "
"pass residual=False for an asymmetric block"
)
self.in_dim = in_dim
self.out_dim = out_dim
self.depth = depth
self.memory = memory
self.value_dim = d_val
self.residual = bool(residual)
self.base = base if isinstance(base, ActivationSpec) else get_activation(base)
if seed is not None:
torch.manual_seed(seed)
linears: list[nn.Linear] = []
prev = in_dim
for _ in range(depth):
linears.append(nn.Linear(prev, hidden, dtype=dtype))
prev = hidden
self.encoder = nn.ModuleList(linears)
# Keys scaled by 1/sqrt(d_key) so the initial scores are O(1) and the
# softmax starts near-uniform rather than saturated.
self.keys = nn.Parameter(
torch.randn(memory, hidden, dtype=dtype) / math.sqrt(hidden)
)
self.values = nn.Parameter(torch.randn(memory, d_val, dtype=dtype))
self.readout = nn.Linear(d_val, out_dim, dtype=dtype)
beta_t = torch.tensor(float(beta), dtype=dtype or torch.get_default_dtype())
if learnable_temperature:
self.beta_raw = nn.Parameter(beta_t)
else:
self.register_buffer("beta_raw", beta_t)betaproperty
beta(self) -> TensorCurrent inverse temperature.
Returns: Tensor
Method implementation
def beta(self) -> Tensor:
"""Current inverse temperature."""
return self.beta_rawquerymethod
query(self, x: Tensor) -> TensorEncoder output q(x), shape (..., hidden).
Parameters
| Name | Type & default | Contract |
|---|---|---|
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
Returns: Tensor
Method implementation
def query(self, x: Tensor) -> Tensor:
"""Encoder output ``q(x)``, shape ``(..., hidden)``."""
h = x
for w, b, spec in self._encoder_specs():
h = h @ w.t()
if b is not None:
h = h + b
assert spec is not None
h = spec.forward(h)
return hattention_weightsmethod
attention_weights(self, x: Tensor) -> TensorSoftmax weights over the memory slots, shape (..., memory).
The diagnostic that makes the block readable: each row is a partition of
unity over the memory, so it says which slots the model is consulting at
x. omnibias.torch.jet_mv.jet_softmax differentiates exactly
this quantity with respect to x.
Parameters
| Name | Type & default | Contract |
|---|---|---|
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
Returns: Tensor
Method implementation
def attention_weights(self, x: Tensor) -> Tensor:
"""Softmax weights over the memory slots, shape ``(..., memory)``.
The diagnostic that makes the block readable: each row is a partition of
unity over the memory, so it says which slots the model is consulting at
``x``. :func:`omnibias.torch.jet_mv.jet_softmax` differentiates exactly
this quantity with respect to ``x``.
"""
scores = self.query(x) @ (self.keys * self.beta).t()
return torch.softmax(scores, dim=-1)valuemethod
value(self, x: Tensor) -> TensorPlain forward value u(x), shape (..., out_dim) (no jet needed).
Parameters
| Name | Type & default | Contract |
|---|---|---|
xpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
Returns: Tensor
Method implementation
def value(self, x: Tensor) -> Tensor:
"""Plain forward value ``u(x)``, shape ``(..., out_dim)`` (no jet needed)."""
q = self.query(x)
h = self.attention_weights(x) @ self.values
if self.residual:
h = h + q
out: Tensor = self.readout(h)
return outextra_reprmethod
extra_repr(self) -> strSee the class contract above and the exact method implementation below.
Returns: str
Method implementation
def extra_repr(self) -> str:
return (
f"in_dim={self.in_dim}, out_dim={self.out_dim}, depth={self.depth}, "
f"memory={self.memory}, value_dim={self.value_dim}, "
f"beta={float(self.beta_raw.detach()):g}, residual={self.residual}"
)out_dimattribute
out_dim: int = out_dimInput coordinate count and output component count.
value_dimattribute
value_dim = d_valWidth d_val of the value slots; defaults to hidden. Must equal hidden when residual is set.
residualattribute
residual = bool(residual)Add the query to the attention output (the usual transformer skip). Keeps a *local* path alongside the non-local one, which trains far better.
baseattribute
base = base if isinstance(base, ActivationSpec) else get_activation(base)Encoder activation; must have a closed-form derivative fast path.
encoderattribute
encoder = nn.ModuleList(linears)Public instance attribute assigned by __init__. The source assignment is encoder = nn.ModuleList(linears); it describes instance state, not a literal constructor default.
keysattribute
keys = nn.Parameter(torch.randn(memory, hidden, dtype=dtype) / math.sqrt(hidden))Keys scaled by 1/sqrt(d_key) so the initial scores are O(1) and the softmax starts near-uniform rather than saturated.
valuesattribute
values = nn.Parameter(torch.randn(memory, d_val, dtype=dtype))Public instance attribute assigned by __init__. The source assignment is values = nn.Parameter(torch.randn(memory, d_val, dtype=dtype)); it describes instance state, not a literal constructor default.
readoutattribute
readout = nn.Linear(d_val, out_dim, dtype=dtype)Public instance attribute assigned by __init__. The source assignment is readout = nn.Linear(d_val, out_dim, dtype=dtype); it describes instance state, not a literal constructor default.
Read the implementation
class AttentionJetMLP(_JetMLPCore):
r"""Deep encoder + one non-local attention block, with exact input derivatives.
``u(x) = W_o [ softmax(beta q(x) K^T) V (+ q(x)) ] + b_o`` where ``q`` is an
ordinary MLP encoder. The memory ``(K, V)`` is a set of trainable slots that
do **not** depend on ``x``, which is what keeps the block closed form: the
only ``x``-dependence enters through the query, and the jet of the query is
exact.
Interpretation for PINNs: the softmax is a learned soft partition of the
input domain, and ``V`` holds one local model per region -- a differentiable,
globally-coupled relative of a domain decomposition, and the reason this
field type is worth having next to the local ones. Note the temperature
collapse this shares with the rest of the repo: as ``beta -> inf`` the
mixture hardens into a crisp nearest-slot assignment (the *feasibility*
sense of collapse, not the founding ``delta -> 0`` one).
Parameters
----------
in_dim, out_dim:
Input coordinate count and output component count.
hidden:
Encoder width. Also the query / key width ``d_key``.
depth:
Number of hidden (activated) encoder layers ``>= 1``.
base:
Encoder activation; must have a closed-form derivative fast path.
memory:
Number of memory slots ``n``.
value_dim:
Width ``d_val`` of the value slots; defaults to ``hidden``. Must equal
``hidden`` when ``residual`` is set.
beta:
Initial inverse temperature ``beta > 0``. Larger is sharper.
learnable_temperature:
Train ``beta`` alongside the weights. It enters the jet as a scale on the
keys, so the derivatives stay closed form while it moves.
residual:
Add the query to the attention output (the usual transformer skip). Keeps
a *local* path alongside the non-local one, which trains far better.
seed:
Optional seed for the memory draw; ``None`` leaves the global RNG alone.
dtype:
Parameter dtype, defaulting to the framework default.
"""
beta_raw: Tensor
def __init__(
self,
in_dim: int,
hidden: int,
out_dim: int = 1,
depth: int = 2,
base: str | ActivationSpec[Tensor] = "tanh",
*,
memory: int = 16,
value_dim: int | None = None,
beta: float = 1.0,
learnable_temperature: bool = False,
residual: bool = True,
seed: int | None = None,
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
if in_dim < 1:
raise ValueError(f"in_dim must be >= 1, got {in_dim}")
if hidden < 1:
raise ValueError(f"hidden must be >= 1, got {hidden}")
if out_dim < 1:
raise ValueError(f"out_dim must be >= 1, got {out_dim}")
if depth < 1:
raise ValueError(f"depth (number of hidden layers) must be >= 1, got {depth}")
if memory < 1:
raise ValueError(f"memory must be >= 1, got {memory}")
if beta <= 0.0:
raise ValueError(f"beta must be > 0, got {beta}")
d_val = hidden if value_dim is None else int(value_dim)
if d_val < 1:
raise ValueError(f"value_dim must be >= 1, got {d_val}")
if residual and d_val != hidden:
raise ValueError(
f"residual needs value_dim == hidden, got {d_val} != {hidden}; "
"pass residual=False for an asymmetric block"
)
self.in_dim = in_dim
self.out_dim = out_dim
self.depth = depth
self.memory = memory
self.value_dim = d_val
self.residual = bool(residual)
self.base = base if isinstance(base, ActivationSpec) else get_activation(base)
if seed is not None:
torch.manual_seed(seed)
linears: list[nn.Linear] = []
prev = in_dim
for _ in range(depth):
linears.append(nn.Linear(prev, hidden, dtype=dtype))
prev = hidden
self.encoder = nn.ModuleList(linears)
# Keys scaled by 1/sqrt(d_key) so the initial scores are O(1) and the
# softmax starts near-uniform rather than saturated.
self.keys = nn.Parameter(
torch.randn(memory, hidden, dtype=dtype) / math.sqrt(hidden)
)
self.values = nn.Parameter(torch.randn(memory, d_val, dtype=dtype))
self.readout = nn.Linear(d_val, out_dim, dtype=dtype)
beta_t = torch.tensor(float(beta), dtype=dtype or torch.get_default_dtype())
if learnable_temperature:
self.beta_raw = nn.Parameter(beta_t)
else:
self.register_buffer("beta_raw", beta_t)
# -- the graph is encoder -> attention -> readout, not one chain ------------- #
def _layer_specs(self) -> list[LayerSpec]:
raise NotImplementedError(
"AttentionJetMLP is an encoder + attention block + readout, not a "
"single layer chain; use _encoder_specs()."
)
def _encoder_specs(self) -> list[LayerSpec]:
"""``(W, b, spec)`` chain of the query encoder ``q(x)`` (all layers activated)."""
specs: list[LayerSpec] = []
for lin in self.encoder:
assert isinstance(lin, nn.Linear)
specs.append((lin.weight, lin.bias, self.base))
return specs
@property
def beta(self) -> Tensor:
"""Current inverse temperature."""
return self.beta_raw
def _check_fastpath(self, max_order: int) -> None:
"""Reject an encoder activation without a closed-form tower of ``max_order``.
The attention block itself needs no check: ``exp`` and the reciprocal
tower are closed form at every order by construction.
"""
spec = self.base
if spec.fastpath is None:
raise ValueError(
f"{type(self).__name__} requires activations with a closed-form "
f"derivative kernel; activation {spec.name!r} has none."
)
try:
spec.fastpath(torch.zeros(1), max_order)
except NotImplementedError as e:
raise ValueError(
f"Activation {spec.name!r} fast-path does not support order "
f"{max_order}: {e}"
) from None
def _point_hidden_jet(self, xi: Tensor, order: int) -> Tensor:
"""Single-point pre-readout jet (encoder + attention [+ residual])."""
q_jet = mlp_jet_mv(xi, self._encoder_specs(), order)
h_jet = jet_attention(
q_jet, self.keys, self.values, self.in_dim, order, beta=self.beta
)
if self.residual:
h_jet = h_jet + q_jet
return h_jet
def _apply_readout_jet(self, hidden_jet: Tensor) -> Tensor:
"""Push a (possibly batched) pre-readout jet through the live affine head."""
w = self.readout.weight
b = self.readout.bias
out: Tensor = torch.tensordot(hidden_jet, w, dims=([-1], [-1]))
if b is not None:
if hidden_jet.ndim == 2:
out = out.clone()
out[0] = out[0] + b
else:
out = out.clone()
out[:, 0, :] = out[:, 0, :] + b
return out
def _point_jet(self, xi: Tensor, order: int) -> Tensor:
"""Single-point jet of the whole block, shape ``(M, out_dim)``."""
return self._apply_readout_jet(self._point_hidden_jet(xi, order))
def query(self, x: Tensor) -> Tensor:
"""Encoder output ``q(x)``, shape ``(..., hidden)``."""
h = x
for w, b, spec in self._encoder_specs():
h = h @ w.t()
if b is not None:
h = h + b
assert spec is not None
h = spec.forward(h)
return h
def attention_weights(self, x: Tensor) -> Tensor:
"""Softmax weights over the memory slots, shape ``(..., memory)``.
The diagnostic that makes the block readable: each row is a partition of
unity over the memory, so it says which slots the model is consulting at
``x``. :func:`omnibias.torch.jet_mv.jet_softmax` differentiates exactly
this quantity with respect to ``x``.
"""
scores = self.query(x) @ (self.keys * self.beta).t()
return torch.softmax(scores, dim=-1)
def value(self, x: Tensor) -> Tensor:
"""Plain forward value ``u(x)``, shape ``(..., out_dim)`` (no jet needed)."""
q = self.query(x)
h = self.attention_weights(x) @ self.values
if self.residual:
h = h + q
out: Tensor = self.readout(h)
return out
def extra_repr(self) -> str:
return (
f"in_dim={self.in_dim}, out_dim={self.out_dim}, depth={self.depth}, "
f"memory={self.memory}, value_dim={self.value_dim}, "
f"beta={float(self.beta_raw.detach()):g}, residual={self.residual}"
)Other public import paths
omnibias.torch.architectures.AttentionJetMLPomnibias.torch.architectures.attention.AttentionJetMLP