PackedMultiheadAttention

PackedMultiheadAttention(
    source_index: object,
    target_index: object,
    query_features: int,
    key_features: int,
    embed_dim: int,
    num_heads: int,
    dropout: float = 0.0,
    bias: bool = True,
    kdim: int | None = None,
    vdim: int | None = None,
    batch_first: bool = True,
    add_self_loops: bool = False,
    *,
    identity: str | None = None,
    generator: Generator | None = None,
    chunk_size: int | None = None,
)

Bases: Module

Multi-head attention in which each node attends only to its sources.

Prior knowledge decides who may attend to whom; the data decides how much. Every node is one token, a vector of embed_dim numbers that stands for the node inside your model. For each node, the layer looks at the nodes with an edge into it (its sources), weighs them by how well their tokens match its own, and returns a weighted mix of what they pass on. The weights are recomputed for every sample: with a gene-regulatory prior, a target gene can weigh its regulators differently in each cell, depending on their state and its own. Two nodes without an edge exchange nothing in this layer.

Information flows along the edge arrow, from source to target, as in PackedLinear. In attention terms, the target asks and its sources answer. The target's token gives a query (what it looks for). Each source's token gives a key (what it offers to be matched against) and a value (what it passes on). Each query-key pair, that is, each live edge, gets a score. A softmax over the edges into one node turns their scores into weights that are non-negative and sum to 1, so the sources compete: a higher weight for one leaves less for the others. num_heads such weightings run side by side, each on its own slice of the channels, so different heads can weigh the same sources differently.

Compare PackedLinear, which learns one weight per edge, fixed after training and the same for every sample. Here nothing is learned per edge. The learned parameters are four embed_dim x embed_dim projections shared by every node and edge (q_proj, k_proj, v_proj, out_proj), so their number does not depend on the graph. The graph only decides which pairs are scored.

One call moves information one edge along the graph: a node's output depends only on its own token and those of its direct sources. Stack layers in your own encoder to reach nodes further upstream. This layer is that one building block, not a Transformer. You build the tokens (for example a projection of each measured value plus a learned embedding per node) and write the rest around the layer: normalization, residual connections, a feed-forward network, and the readout.

Masking nn.MultiheadAttention with the graph computes the same mix for every node that has a source. But it scores every pair of tokens, an (n, n) matrix per head and sample (400 million scores for 20,000 nodes), and then discards the pairs without an edge. This layer scores only the live edges, so its memory and time grow with their number, nnz, rather than with n squared.

Build it with from_adjacency on a parse_adjacency spec: one token per node (per unit, when widths= gives a node several). It reads the indices and sequence lengths off the spec, so the query and key sides cannot be swapped by hand. The constructor below takes packed indices directly, for pairs that come from no spec, such as a bipartite map in which pathway tokens (queries) attend to gene tokens (keys).

Call it like nn.MultiheadAttention (it is not a subclass): layer(query, key, value), passing the same tokens three times for attention within one graph. Two defaults differ from MHA: batch_first=True and need_weights=False. With need_weights=True the layer returns one weight per live edge, not a dense (L, S) matrix.

Parameters:

Name	Type	Description	Default
source_index	torch.Tensor, numpy.ndarray, or sequence of int	Where each live edge starts. Entry i is the position in key / value of the token that edge i reads from, and must satisfy 0 <= source_index < key_features. 1-D integer indices of length nnz >= 1: an integer tensor, an integer numpy array, or a sequence of int or numpy integers. A set or a mapping is rejected, since entries pair up by position, and so is a complex tensor. Copied to an int64 buffer, so later writes to the argument do not reach this layer.	required
target_index	torch.Tensor, numpy.ndarray, or sequence of int	Where each live edge ends. Entry i is the position in query of the token that attends along edge i, and must satisfy 0 <= target_index < query_features. Same length and forms as source_index. The two arrays are paired position by position and must not repeat a (source, target) pair. Their shared order is the order of the per-edge weights forward returns; pairs added by add_self_loops follow at the end.	required
query_features	int or numpy integer	How many tokens attend: the sequence length of query. Positive integer, stored as int.	required
key_features	int or numpy integer	How many tokens can be attended to: the sequence length of key / value. Positive integer, stored as int. Equal to query_features for attention within one node set; smaller or larger for a bipartite query/key map.	required
embed_dim	int or numpy integer	How many numbers describe each token: the width of query, key, value, and the output. A wider token can carry more about its node, and the four projections hold 4 * embed_dim**2 weights. Positive integer, divisible by num_heads, stored as int.	required
num_heads	int or numpy integer	How many attention patterns run side by side. Each head scores the same live edges from its own embed_dim // num_heads channels, so it can weigh a node's sources differently from the other heads. The heads' mixes are joined and passed through out_proj. Positive integer, stored as int.	required
dropout	float	Regularization on the edge weights, in training mode only: each weight is zeroed with this probability and the rest are scaled by 1 / (1 - dropout), so the model cannot lean on any single edge. 0.0 disables it. Must be finite and in [0, 1], the range torch accepts; integers 0 and 1 are accepted, bool is rejected.	0.0
bias	bool	Whether the four projections learn an additive bias. False makes them pure linear maps. Must be a bool.	True
kdim	int, numpy integer, or None	Exists so code written for nn.MultiheadAttention runs unchanged; keys of a different width are not supported. Must be None or an int or numpy integer equal to embed_dim (True is rejected even when embed_dim is 1).	None
vdim	int, numpy integer, or None	Same as kdim, for the values.	None
batch_first	bool	Where the batch axis sits. True (kpnn2 sample-major): batched tensors are (..., seq, embed_dim). False (the nn.MultiheadAttention layout): batched tensors are (seq, batch, embed_dim). Unbatched 2-D (seq, embed_dim) ignores this flag. Must be a bool.	True
add_self_loops	bool	Let every token also attend to itself. Without a self-loop, a node's own token shapes how it weighs its sources (through its query) but is not part of the mix, and a node with no incoming edge, such as an input node of an AdjacencySpec, mixes to zeros. With True, any missing (i, i) pair is added to this layer's index buffers, after the given pairs. Existing self-loops are kept, not duplicated, and the caller's index objects are not mutated. The graph itself is not changed: other layers built from the same spec do not see these pairs. Requires query_features == key_features. Must be a bool.	False
identity	str or None	Labels the graph this layer belongs to, so load_state_dict refuses a checkpoint labeled for another graph, even one with the same index pattern, such as a graph with a renamed node. Typically spec.fingerprint. An opaque string, stored in state_dict next to index_digest as a 1-D CPU uint8 tensor of the UTF-8 bytes. load_state_dict compares it with the checkpoint's identity: the checkpoint has no identity: it loads, even with strict=True it has this layer's identity: it loads it has another identity, or this layer's identity is None: Kpnn2Error, and no weights are loaded None means this layer does not claim an identity, so it loads only checkpoints saved without one.	None
generator	Generator or None	A reproducible start without touching the global random state: the Xavier-uniform draws for the four projections use this generator. None uses the default torch generator and keeps today's path: nn.Linear kaiming-init then Xavier on that global stream. When set, those nn.Linear constructors do not advance the global stream. Not stored on the module; pass it again to reset_parameters to replay. Do not pass a seed integer.	None
chunk_size	int, numpy integer, or None	Memory control for large graphs; it does not change what is computed. None gathers all live edges at once. A positive int processes them in slices of that many edges, so training holds less memory, at the cost of time: the softmax is still over every live key of a query, but the pair gathers are slices, and backward recomputes them. With a positive int, forward and backward also check on every call that each query's packed softmax sums to 1 (to 0 only when no live key scores above -inf: none left after padding, or every score overflowed) and that its gradient sums to 0, and raise AssertionError (a kpnn2 bug) if not; each check waits for the device once. A numpy integer is stored as int. bool and values other than None or a positive integer raise Kpnn2Error. Not stored in state_dict or index_digest, so a checkpoint loads into a layer with another chunk_size.	None

Attributes:

Name	Type	Description
query_features	int	Number of tokens that attend (sequence length of query).
key_features	int	Number of tokens that can be attended to (sequence length of key / value).
embed_dim	int	Token width.
num_heads	int	Number of attention heads.
head_dim	int	Channels per head, embed_dim // num_heads.
nnz	int	Number of live edges scored, counting any pair added by add_self_loops.
dropout	float	Dropout probability on the edge weights.
batch_first	bool	Layout flag; default True.
add_self_loops	bool	Whether missing self-loops were added at construction.
source_index	Tensor	Int64 buffer: the key / value position of each live edge, length nnz. Treat it as read-only: like any PyTorch buffer it can be written to, and doing so rewires the layer without reinitializing it. Rebuild from the edgelist instead.
target_index	Tensor	Int64 buffer: the query position of each live edge, length nnz, read-only in the same sense.
q_proj, k_proj, v_proj, out_proj	Linear	The four learned embed_dim -> embed_dim projections (query, key, value, and output): separate layers, not a fused in_proj_weight.
identity	str | None	The constructor identity, or None.
chunk_size	int | None	Edge-chunk length, or None for all live edges at once.

Raises:

Type

Description

Kpnn2Error

At construction, if the indices are empty, not 1-D integers, mismatched in length, out of range, or duplicated as (source, target) pairs; if the sizes are not positive integers; if embed_dim is not divisible by num_heads; if dropout is a bool, not a number, NaN, or outside [0, 1]; if bias, batch_first, or add_self_loops is not a bool; if kdim / vdim are neither None nor an integer equal to embed_dim; if add_self_loops is set when query_features != key_features; if identity is neither a str nor None; if generator is neither a torch.Generator nor None; or if chunk_size is neither None nor a positive integer. From load_state_dict, when the checkpoint carries an index digest or identity that does not match this layer, in which case the weights are not loaded. forward documents its own rejected arguments.

See Also

PackedLinear : Same packed pairs, but one trained weight per edge, the same for every sample, instead of attention weights computed per sample. MaskedLinear : Dense (out, in) weight; the default layer whenever that square fits. AdjacencySpec : Graph-wide packed edges; pass it to from_adjacency. torch.nn.MultiheadAttention : Dense equivalent, and the reference for call shape, heads, and training.

Notes

One forward pass, for one sample and one head with d = embed_dim // num_heads channels:

1. Project every token: q = q_proj(query), k = k_proj(key), v = v_proj(value), each split into num_heads slices of d channels.
2. Score each live edge e from source s to target t: score[e] = dot(q[t], k[s]) / sqrt(d).
3. Turn the scores of the edges into each target into weights with a softmax over those edges only. Keys masked by key_padding_mask take no part. In training mode, dropout then zeroes weights at random and scales the rest by 1 / (1 - dropout).
4. Mix: mixed[t] is the sum of weight[e] * v[s] over the edges e into t.
5. Join the heads' slices of mixed back into embed_dim channels and apply out_proj.

A query with no live keys, because it has no incoming edge and no self-loop or because key_padding_mask removed all of them, mixes to zeros rather than NaN, so its output is the bias of out_proj (zeros when bias=False). Input nodes of an AdjacencySpec are exactly that case unless add_self_loops=True. With it, an input node's only key is itself, so its softmax weight is 1 and it passes on its own value.

No (n, n) or (L, S) score matrix is ever allocated and torch.sparse is not imported; cost scales with nnz, not with query_features * key_features. chunk_size=None (the default) gathers all live pairs at once, and training keeps those (..., nnz, heads, head_dim) gathers for backward. A positive chunk_size gathers chunked live pairs and rematerializes those (..., chunk, heads, head_dim) tensors in backward, so training does not save a full (..., nnz, heads, head_dim) gather. The packed softmax (..., nnz, heads) is still stored.

forward documents the accepted layouts, packed need_weights, attn_mask / is_causal, and key_padding_mask.

Index buffers stay integer after .half() / bfloat16 / .double(); projection weights follow the module floating dtype like nn.Linear. torch.autocast is unsupported: forward disables it and casts query, key, and value to the parameter dtype. They must be floating-point tensors; an integer, bool, or complex input raises instead of being cast. Cast the module with .to(dtype=...) (or .half() / .double()) instead.

state_dict adds index_digest, a 1-D CPU uint8 tensor of length 32: the SHA-256 of the packed indices plus query_features, key_features, embed_dim, and num_heads, so a reshape or a different head split cannot collide. It is not a registered buffer, and a missing digest is not an error, even with strict=True. identity is stored the same way when the constructor was given one: a rename that leaves the packed index pattern unchanged is caught when callers pass spec.fingerprint. chunk_size is not part of state_dict or index_digest. copy.deepcopy works.

Examples:

An input x feeds a two-node feedback loop between a and b, which feeds an output y. All four nodes form one four-token sequence:

>>> import pandas as pd
>>> import torch
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> attn = kpnn2.PackedMultiheadAttention.from_adjacency(
...     spec,
...     embed_dim=8,
...     num_heads=2,
...     add_self_loops=True,
...     identity=spec.fingerprint,
... )
>>> attn.nnz, attn.head_dim
(8, 4)
>>> tokens = torch.randn(2, spec.state_dim, 8)
>>> output, weights = attn(tokens, tokens, tokens)
>>> tuple(output.shape)
(2, 4, 8)
>>> weights is None
True

With need_weights=True each sample gets one weight per live edge: the four graph edges, then the four added self-loops.

>>> _, packed = attn(
...     tokens,
...     tokens,
...     tokens,
...     need_weights=True,
... )
>>> tuple(packed.shape)
(2, 8)
>>> for s, t in zip(
...     attn.source_index.tolist(),
...     attn.target_index.tolist(),
... ):
...     print(spec.nodes[s], "->", spec.nodes[t])
a -> b
a -> y
b -> a
x -> a
a -> a
b -> b
x -> x
y -> y

Methods:

from_adjacency classmethod

from_adjacency(
    spec: AdjacencySpec,
    embed_dim: int,
    num_heads: int,
    *,
    dropout: float = 0.0,
    bias: bool = True,
    batch_first: bool = True,
    add_self_loops: bool = False,
    identity: str | None = None,
    generator: Generator | None = None,
    chunk_size: int | None = None,
) -> PackedMultiheadAttention

Build attention over the nodes of an AdjacencySpec.

This is the recommended way to build this layer on an AdjacencySpec: one token per unit of the state vector, and the query at each edge's target attends to the key at its source. It reads the packed indices off spec and uses spec.state_dim as both sequence lengths, so the query and key sides cannot be swapped by hand. On this square layout a hand-written swap is always silent: the constructor accepts it and every node attends along reversed edges. The result is the layer the constructor builds from spec.source_index, spec.target_index, spec.state_dim, spec.state_dim, embed_dim, and num_heads, with the same init draws. It keeps no reference to spec. Options after num_heads are keyword-only; kdim and vdim, which can only equal embed_dim, are left out.

Parameters:

Name	Type	Description	Default
spec	AdjacencySpec	Spec from parse_adjacency.	required
embed_dim	int or numpy integer	Same as the PackedMultiheadAttention constructor.	required
num_heads	int or numpy integer	Same as the PackedMultiheadAttention constructor.	required
dropout	float	Same as the PackedMultiheadAttention constructor.	0.0
bias	bool	Same as the PackedMultiheadAttention constructor.	True
batch_first	bool	Same as the PackedMultiheadAttention constructor.	True
add_self_loops	bool	Same as the PackedMultiheadAttention constructor.	False
identity	str or None	Same as the PackedMultiheadAttention constructor. Not filled in from spec: pass spec.fingerprint yourself.	None
generator	Generator or None	Same as the PackedMultiheadAttention constructor.	None
chunk_size	int, numpy integer, or None	Same as the PackedMultiheadAttention constructor.	None

Returns:

Type	Description
PackedMultiheadAttention	New layer with query_features == key_features == spec.state_dim.

Raises:

Type	Description
Kpnn2Error	If spec is not an AdjacencySpec (a LayeredSpec or a Hop included), or if the constructor rejects any other argument.

See Also

PackedLinear.from_adjacency : One trainable scalar per edge on the same packed pairs.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> spec = kpnn2.parse_adjacency(
...     pd.DataFrame(
...         {
...             "source": ["x", "a", "b", "a"],
...             "target": ["a", "b", "a", "y"],
...         }
...     )
... )
>>> attn = kpnn2.PackedMultiheadAttention.from_adjacency(
...     spec,
...     embed_dim=8,
...     num_heads=2,
...     identity=spec.fingerprint,
... )
>>> attn.query_features, attn.key_features, attn.nnz
(4, 4, 4)

reset_parameters

reset_parameters(
    generator: Generator | None = None,
) -> None

Xavier-uniform projections, like nn.MultiheadAttention.

Parameters:

Name	Type	Description	Default
generator	Generator or None	Isolated RNG for the four Xavier draws. None uses the default torch generator. Not stored on the module. Biases, when present, are zeroed and do not consume the generator.	None

forward

forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    key_padding_mask: Tensor | None = None,
    need_weights: bool = False,
    attn_mask: Tensor | None = None,
    average_attn_weights: bool = True,
    is_causal: bool = False,
) -> tuple[torch.Tensor, torch.Tensor | None]

Let each query attend to its live keys and mix their values.

For attention within one graph, pass the same tokens as query, key, and value: every node then attends to its sources. Pass different tensors when queries and keys are different token sets, as in a bipartite map built with the constructor. The arguments are those of nn.MultiheadAttention.forward, so an encoder written for it can call this layer; the class docstring lists the differences.

With batch_first=True (the default), query, key, and value are (..., seq, embed_dim): unbatched 2-D (seq, embed_dim), batched 3-D (batch, seq, embed_dim), or more leading axes. With batch_first=False, unbatched 2-D is still (seq, embed_dim); batched 3-D is (seq, batch, embed_dim), transposed to batch-first for the packed kernel and transposed back. The batch axes of all three must match.

Parameters:

Name	Type	Description	Default
query	Tensor	The tokens that attend: query_features of them along the sequence axis, each embed_dim wide. Must be floating-point; an integer, bool, or complex tensor raises Kpnn2Error rather than being cast to the parameter dtype.	required
key	Tensor	The tokens matched against the queries: key_features of them along the sequence axis, each embed_dim wide. Floating-point, like query. Required; it is not defaulted to query.	required
value	Tensor	What each key position passes on, in the shape of key and usually the same tensor. Floating-point, like query.	required
key_padding_mask	Tensor or None	Keys to leave out of a given sample, for example a feature that sample did not measure: True means ignore that key. Its edges drop out of the softmax, so each query's remaining sources share the whole weight. A boolean tensor of shape (S,) unbatched or (N, S) batched, where S = key_features; with more leading axes, the same ones as the batch-first query. The batch axes come first even when batch_first=False. Applied in packed space, without allocating (n, n). Copied onto the scores' device and kept boolean. A float mask raises Kpnn2Error.	None
need_weights	bool	Also return the weight of every live edge, for example to draw edge thickness on the prior graph. Must be a bool.	False
attn_mask	Tensor or None	Must be None: the edgelist is the mask. Kept so that calls written for nn.MultiheadAttention still run.	None
average_attn_weights	bool	With need_weights=True: True averages the weights over the heads, False keeps one per head. Must be a bool.	True
is_causal	bool	Must be False. Kept so that calls written for nn.MultiheadAttention still run.	False

Returns:

Name	Type	Description
output	Tensor	The updated tokens, one per query position, in the shape and layout of query. A query with no live keys left mixes to zeros, not NaN, before out_proj.
weights	Tensor or None	None unless need_weights=True. Then the softmax weight of every live edge, in the order of source_index / target_index, not nn.MultiheadAttention's dense (L, S) map; this path does not allocate (L, S). Shape (..., nnz) averaged over heads, or (..., nnz, num_heads) per head, with the batch axes placed as in output: (nnz, batch) or (nnz, batch, num_heads) when batch_first=False and the input is batched. The weights into one query sum to 1, or to 0 when it has no live key. In training mode with dropout > 0 they are the weights after dropout, as the mix used them; call eval() first to read the plain softmax. nnz counts the pairs added by add_self_loops, so zip the weights with this layer's source_index / target_index, not with to_edgelist().

Raises:

Type	Description
Kpnn2Error	If query, key, or value is not a floating-point tensor of at least 2 dimensions with last axis embed_dim and the sequence length above; if their batch axes differ; if a batched input is not 3-D under batch_first=False; if key_padding_mask is not a boolean tensor of the shape above; if need_weights or average_attn_weights is not a bool; if attn_mask is not None; or if is_causal is not False.