Prompt learning starts with tensor contracts

01 / Start reading a model with its shapes

Accepting tensors does not establish integration with a complete task. When reading prompt-learning code, first trace the meaning of inputs, shape transformations and the path from loss to parameters. The KAPT component takes node features X:[N,H] and graph assignments batch:[N], and returns updated nodes plus one prompt per graph.

For six nodes in three graphs, with H=3 and prompt width D=8, the outputs should have shapes [6,3] and [3,8]. This defines the interface only; the prediction head, labels, task loss and data splits still need to be supplied externally.

02 / Four stages in the source

The task pool selects learned tokens and applies self-attention followed by token mean pooling. A second path encodes each node with an MLP and then averages encoded features per graph. The two prompts are concatenated and fused; the graph prompt is projected and broadcast to its nodes for a gated residual update.

The order of operations matters: averaging after a nonlinear MLP generally differs from encoding an average of raw nodes. This implementation does not read edges, so this step cannot distinguish different graph connectivity for the same node-feature collection.

03 / Contracts cover more than dimensions

Node features and prompts must be finite floating-point tensors matching parameter dtype and device. batch must be long with one entry per node. Graph IDs must cover 0 through B−1, with at least one node in each graph. Task IDs may be shared or supplied per graph; prompt width must also be divisible by the attention head count.

Rejecting malformed inputs adds validation work but turns subtle broadcasting, type and empty-graph problems into explicit errors. Keeping these constraints is useful while learning. A later throughput optimization should measure validation cost before deciding which layer owns those checks.

04 / Freezing parameters does not disable gradients

This standalone simplified example freezes a linear encoder while updating an input prompt and prediction head. Because the prompt precedes the encoder, gradients must pass through the encoder to reach it; wrapping that forward pass in no_grad would break the path. eval controls behaviors such as Dropout, rather than freezing parameters.

# Simplified teaching example, not a molecular model.
import torch
from torch import nn

encoder = nn.Linear(3, 3, bias=False)
head = nn.Linear(3, 1, bias=False)
with torch.no_grad():  # Initialization only.
    encoder.weight.copy_(torch.eye(3))
    head.weight.fill_(1.0)
encoder.requires_grad_(False)
encoder.eval()
prompt = nn.Parameter(torch.zeros(1, 3))
optimizer = torch.optim.SGD([prompt, *head.parameters()], lr=0.01)
x = torch.ones(2, 3)
encoder_before = encoder.weight.detach().clone()
prompt_before = prompt.detach().clone()

optimizer.zero_grad(set_to_none=True)
output = head(encoder(x + prompt))
assert output.shape == (2, 1)
output.square().mean().backward()  # Toy loss; no task labels.
assert encoder.weight.grad is None
assert prompt.grad is not None and bool((prompt.grad != 0).any())
optimizer.step()
assert torch.equal(encoder.weight, encoder_before)
assert not torch.equal(prompt, prompt_before)
print("Frozen weights unchanged; prompt updated.")

If a prompt operates after cached features, the encoder may extract those features without gradients. Whether to retain the computation graph depends on the insertion point; one recipe does not fit every architecture.

05 / Check gradient paths and boundaries

The existing test_prompts.py covers unequal graph sizes, node permutations, the residual formula and malformed inputs, and checks gradients for selected task tokens. A non-null grad is only a first check. After optimizer integration, compare parameters before and after a step to confirm that frozen weights stay fixed and intended parameters change.

• Compare aggregation with manually computed per-graph means. Jointly permuting nodes and batch should preserve pooled outputs, while node-level outputs should reorder accordingly.
• Test empty graphs, out-of-range tasks, integer features, nonfinite values and device mismatches.
• Store architecture configuration with a model; after restoration, compare outputs on the same input rather than treating error-free loading as sufficient evidence.

06 / What remains between components and experiments

A future comparison could include a fixed encoder with a task head, task prompts alone, pooled node features alone, fusion and residual injection, holding splits and budgets constant. This is a proposed experiment. Component checks do not establish molecular prediction performance; edge-aware encoder integration, a real-data protocol and checkpoint reproduction remain missing. The KANO method and paper belong to their authors; the available evidence supports describing the public curation and engineering changes.