Attention Layers

The k3_node.layers.attention module provides standalone attention building blocks used by K3-Node's scalable graph transformers (Polynormer, SGFormer, GPSE) and can also be composed into custom architectures. Unlike k3_node.layers.conv attention layers (GATConv, TransformerConv, ...), these operate on dense [batch, nodes, channels] tensors rather than sparse edge_index message passing.

Linear / Kernelized Attention

PerformerAttention

FAVOR+ kernelized linear attention (from the Performer paper), used by Polynormer for its global attention branch. Approximates full softmax attention in linear time/memory by projecting queries/keys through random Fourier features.

`k3_node.layers.attention.PerformerAttention`

Bases: Layer

k3_node.layers.PerformerAttention

Initialization Arguments:

Parameters:

Name	Description	Default
`channels`	The number of output channels.	required
`heads`	The number of attention heads.	required
`head_channels`	The number of attention heads.	`64`
`kernel`	activation function.	`relu`
`qkv_bias`	activation function.	`False`
`attn_out_bias`	Bias in Attention Out.	`True`
`dropout`	Dropout rate.	`0.0`

Example

import numpy as np
from k3_node.layers import PerformerAttention

x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]

mask = np.ones((1, 10), dtype=bool)  # which nodes are real (not padding)
attn = PerformerAttention(channels=8, heads=2)  # linear-complexity attention
print(tuple(attn(x, mask).shape))  # (1, 10, 8)

PerformerProjection

The random-feature projection used internally by PerformerAttention to approximate the softmax kernel; useful standalone when building a custom linear-attention layer.

`k3_node.layers.attention.PerformerProjection`

Bases: Layer

Layer PerformerProjection.

Example

import numpy as np
from k3_node.layers import PerformerProjection

q = k = v = np.random.rand(1, 2, 10, 8).astype("float32")  # [batch, heads, num_nodes, head_dim]
proj = PerformerProjection(num_cols=8)  # random-feature approximation of softmax attention
print(tuple(proj(q, k, v).shape))  # (1, 2, 10, 8)

Usage:

from keras import ops
from k3_node import layers as k3_layers

attn = k3_layers.PerformerAttention(channels=64, heads=4)
mask = ops.ones((1, num_nodes))  # [batch, nodes]
out = attn(x, mask)  # x: [batch, nodes, 64]

Polynomial & Scalable Graph Attention

PolynormerAttention

Local-to-global polynomial-expressive attention from Polynormer, combining a local propagation term with a global attention term whose polynomial expansion is computed via PerformerAttention-style kernelization.

`k3_node.layers.attention.PolynormerAttention`

Bases: Layer

Layer PolynormerAttention.

Example

import numpy as np
from k3_node.layers import PolynormerAttention

x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]

mask = np.ones((1, 10), dtype=bool)
attn = PolynormerAttention(channels=8, heads=2)
print(tuple(attn(x, mask).shape))  # (1, 10, 128)

SGFormerAttention

All-pair, single-layer linear attention from SGFormer, designed to replace deep attention stacks with one global mixing layer for scalable graph transformers.

`k3_node.layers.attention.SGFormerAttention`

Bases: Layer

The simple global attention mechanism from the "SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations" <https://arxiv.org/abs/2306.10759>_ paper.

Parameters:

Name	Type	Description	Default
`channels`	`int`	Size of each input sample.	required
`heads`	`int`	Number of parallel attention heads. (default: :obj:`1`)	`1`
`head_channels`	`int`	Size of each attention head. (default: :obj:`64`)	`64`
`qkv_bias`	`bool`	If specified, add bias to query, key and value in the self attention. (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.layers import SGFormerAttention

x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]

attn = SGFormerAttention(channels=8, heads=2)
print(tuple(attn(x).shape))  # (1, 10, 64)

Usage:

attn = k3_layers.SGFormerAttention(channels=64, heads=4)
out = attn(x)  # x: [batch, nodes, 64]

Q-Former (Query Transformer)

QFormer

A BERT-style querying transformer (as used in BLIP-2 / GPSE) that distills a variable-length input sequence into a fixed set of learned query tokens via cross-attention, useful for pooling variable-size node sets into a fixed-size graph representation.

`k3_node.layers.attention.QFormer`

Bases: Layer

Layer QFormer.

Example

import numpy as np
from k3_node.layers import QFormer

x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]

model = QFormer(input_dim=8, hidden_dim=16, output_dim=4, num_heads=2, num_layers=1)
print(tuple(model(x).shape))  # (1, 10, 4)

QFormerEncoderLayer

A single self-attention + feed-forward encoder block used inside QFormer.

`k3_node.layers.attention.QFormerEncoderLayer`

Bases: Layer

Layer QFormerEncoderLayer.

Example

import numpy as np
from k3_node.layers import QFormerEncoderLayer

x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]

layer = QFormerEncoderLayer(input_dim=8, hidden_dim=16, num_heads=2)
print(tuple(layer(x).shape))  # (1, 10, 8)

Usage:

qformer = k3_layers.QFormer(input_dim=64, hidden_dim=64, output_dim=32, num_heads=4, num_layers=2)
out = qformer(x)  # x: [batch, nodes, 64] -> [batch, nodes, 32]