InducedTransformerBlock

class sdm.nn.InducedTransformerBlock(channels: int, num_inducing_points: int, inducing_block: TransformerBlock, output_block: TransformerBlock, device: device | str | None = None, dtype: dtype | None = None)

Bases: Module

Transformer block using learned inducing points.

Introduced in the “Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks” paper, the block routes attention through a small set of \(M\) learned inducing points, reducing the cost of attending a query of size \(Q\) to a key/value context of size \(K\) from \(O(Q \cdot K)\) to \(O((Q + K) \cdot M)\):

\[H = \mathrm{InducingBlock}_1(I, \mathrm{key\_value}), \quad \mathrm{out} = \mathrm{OutputBlock}_2(\mathrm{query}, H),\]

where \(I\) are the learned inducing points and \(H\) are the inducing points after attending to the key/value elements. Passing key_value=None recovers the induced self-attention block (ISAB).

Parameters:

• channels (int) – The number of input and output channels.

• num_inducing_points (int) – The number of learned inducing points \(M\).

• inducing_block (TransformerBlock) – The TransformerBlock that updates the learned inducing points from key/value context.

• output_block (TransformerBlock) – The TransformerBlock that updates the input queries from the induced context.

• device (torch.device | str | None) – The device.

• dtype (torch.dtype | None) – The dtype.

forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: Literal[False] = False, batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → Tensor

forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: Literal[True], batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → tuple[Tensor, KVCacheEntry]

forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: bool, batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → Tensor | tuple[Tensor, KVCacheEntry]

The forward pass.

Parameters:

• query (Tensor) – The query tensor with shape [..., Q, C]. Q is the query sequence length, C is the number of channels.

• key_value (Tensor | KVCacheEntry | None) – The key/value tensor with shape [..., KV, C] or precomputed key/value projections as a KVCacheEntry. KV is the key/value sequence length. If omitted, query is used for induced self-attention.

• seqused_key_value (Tensor | None) – Valid key/value lengths with shape [...] and torch.int32 dtype.

• attn_mask (Tensor | None) – Boolean attention mask with shape [..., KV]. Entries set to True participate in attention.

• return_key_value (bool) – Whether to return the computed key and value projections for the final attention site alongside the output.

• batch_size_limit (int | Literal['auto'] | None) – Maximum number of batch elements processed at once.

• out (Tensor | None) – The output tensor.

Returns:

Tensor with shape [..., Q, C] when return_key_value is False. Otherwise, a tuple of the output tensor and a KVCacheEntry.

Return type:

Tensor | tuple[Tensor, KVCacheEntry]