feat: TabPFN-like PFN options: summary tokens, alternating attention, per-task momentum scale, synthetic tasks and an inference ensemble

- model.pfn_summary_tokens: learned tokens carry each event through the
encoder; a context event's label enters at the input.
- model.pfn_alternating: attention across the task's events after every
encoder block (TabPFN's alternation of feature and row attention).
- model.pfn_normalize: each task's momenta are scaled so the context's
geometric-mean object pT is 50 GeV.
- model.pfn_ensemble: in-context inference averages class probabilities
over passes with rotated, flipped events and permuted class labels.
- task.synthetic_tasks: a fraction of training tasks rank events of 1-3
processes by a random network of event observables.

Defaults keep the previous model and prior: outputs of the pooled PFN
match the previous code to 2.4e-7, and the prior draws the same tasks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/gnn4colliders/cli.py

@@ -321,19 +321,21 @@ def in_context_forward(
     checkpoint = _checkpoint(config)
     model = make_model(config, events, load_pretrained=False)
     load_weights(model, load_checkpoint(checkpoint))
-    build = (
-        meta_tabular_forward
-        if model_config.name in ("tabpfn", "tabicl")
-        else pfn_forward
-    )
-    forward = build(
-        model.to(device),
-        events,
-        context,
-        num_classes=model_config.out_size,
-        batch_size=data.batch_size,
-        device=device,
-    )
+    sizes = {"num_classes": model_config.out_size, "batch_size": data.batch_size}
+    if model_config.name in ("tabpfn", "tabicl"):
+        forward = meta_tabular_forward(
+            model.to(device), events, context, device=device, **sizes
+        )
+    else:
+        forward = pfn_forward(
+            model.to(device),
+            events,
+            context,
+            device=device,
+            ensemble=model_config.pfn_ensemble,
+            seed=config.train.seed,
+            **sizes,
+        )
     return forward, str(checkpoint)
 
 

src/gnn4colliders/config.py

@@ -120,6 +120,21 @@ class ModelConfig:
     max_classes: int = 10
     #: Attention blocks over each task's events.
     pfn_layers: int = 6
+    #: Learned tokens added to each event's objects that carry the event
+    #: across the task (labels enter at the input); 0: the mean over the
+    #: encoded objects, with the label added after the encoder.
+    pfn_summary_tokens: int = 0
+    #: Attention across the task's events after every encoder block too (as
+    #: TabPFN alternates attention across features and across rows), so the
+    #: event encodings depend on the task; needs ``pfn_summary_tokens``.
+    pfn_alternating: bool = False
+    #: Scale each task's momenta and energies so that the geometric mean pT
+    #: of the context's objects is ``PT_REFERENCE`` (pfn.py).
+    pfn_normalize: bool = False
+    #: In-context inference averages the class probabilities of this many
+    #: passes, each with its own azimuthal rotation, pseudorapidity flip and
+    #: order of the class labels.
+    pfn_ensemble: int = 1
     # TabPFN only.
     tabpfn_version: str = "v3.5"
     #: A checkpoint of that version (a path, or a file name in the TabPFN
@@ -177,6 +192,10 @@ class TaskConfig:
     #: Smallest and largest number of context events (log-uniform).
     context_events: list[int] = field(default_factory=lambda: [16, 2048])
     query_events: int = 256
+    #: Fraction of training tasks whose classes are not processes but
+    #: ranges of a random function of event observables (pfn.py), on events
+    #: of 1 to 3 processes.
+    synthetic_tasks: float = 0.0
     #: Momenta and energies of a task are scaled by a factor uniform in
     #: ``[1 - energy_jitter, 1 + energy_jitter]``.
     energy_jitter: float = 0.2

src/gnn4colliders/models.py

@@ -663,6 +663,9 @@ def build_model(
             n_heads=config.n_heads,
             ffn_size=config.ffn_size or 4 * config.hid_size,
             dropout=config.dropout,
+            summary_tokens=config.pfn_summary_tokens,
+            alternating=config.pfn_alternating,
+            normalize=config.pfn_normalize,
         )
     if config.name in ("tabpfn", "tabicl"):
         from . import tabular

src/gnn4colliders/pfn.py

@@ -2,9 +2,9 @@
 Hollmann et al. 2023) that classifies collider events in context.
 
 A task is a set of context events with labels and query events without.
-The PFN embeds every event with the object transformer (mean over objects),
-adds a label embedding to the context events, and runs attention blocks over
-the task's events in which every event attends to the context events only.
+The PFN encodes every event with the object transformer, adds a label
+embedding to the context events, and runs attention blocks over the task's
+events in which every event attends to the context events only.
 Queries therefore never see each other, and the context's keys and values
 can be computed once for any number of queries (:func:`in_context_forward`).
 
@@ -14,15 +14,26 @@ prior: each task takes ``task.task_classes`` random classes of
 proportions and a log-uniform number of context events. Training tasks are
 augmented by a random azimuthal rotation and longitudinal flip of each event
 and a scale of the task's momenta and energies (``task.energy_jitter``).
-A new dataset then needs no training: its labelled training events are the
-context (``data.train_events``).
+With ``task.synthetic_tasks``, a fraction of the training tasks instead
+classify events of 1 to 3 processes by ranges of a random function of their
+observables, as TabPFN's prior draws random functions: no two tasks repeat
+and many decision boundaries are fine. A new dataset then needs no
+training: its labelled training events are the context
+(``data.train_events``).
+
+The model options follow TabPFN too: summary tokens that carry each event
+and its label through the encoder, attention across events between the
+encoder blocks (``model.pfn_alternating``, TabPFN's alternation of
+attention over features and over rows), a per-task scale of the momenta
+(``model.pfn_normalize``) and an ensemble of transformed passes at
+inference (``model.pfn_ensemble``).
 """
 
 from __future__ import annotations
 
 import math
 from collections.abc import Callable, Iterator
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 import torch
@@ -30,8 +41,17 @@ from torch import nn
 
 from .config import TaskConfig
 from .data import PackedEvents, map_tensors
-from .models import EncoderBlock, ObjectTransformer
-from .sequences import SequenceBatch, augment_tokens, collate_sequence_batch
+from .models import NODE_TYPE, EncoderBlock, ObjectTransformer, pair_features
+from .sequences import (
+    BTAG,
+    ENERGY,
+    ETA,
+    PHI,
+    PT,
+    SequenceBatch,
+    augment_tokens,
+    collate_sequence_batch,
+)
 
 #: Logit of the classes a task does not have.
 ABSENT_LOGIT = -1e4
@@ -59,7 +79,18 @@ class TaskBatch:
 
 class EventPFN(nn.Module):
     """Query logits ``[tasks * queries, max_classes]`` of a :class:`TaskBatch`;
-    the logits of classes a task does not have are ``ABSENT_LOGIT``."""
+    the logits of classes a task does not have are ``ABSENT_LOGIT``.
+
+    Each event is encoded by the object transformer ``encoder``. Without
+    summary tokens an event is the mean of its encoded objects plus, for a
+    context event, its label embedding. With ``summary_tokens`` each event
+    carries that many learned tokens through the encoder, and a context
+    event's label embedding is added to them at the input; with
+    ``alternating`` the summary tokens also attend across the task's events
+    after every encoder block, so the encoding of an event depends on the
+    task. Then ``n_layers`` blocks attend across events. Across events,
+    every event attends to the context events only.
+    """
 
     def __init__(
         self,
@@ -70,62 +101,199 @@ class EventPFN(nn.Module):
         n_heads: int,
         ffn_size: int,
         dropout: float = 0.0,
+        summary_tokens: int = 0,
+        alternating: bool = False,
+        normalize: bool = False,
     ) -> None:
         super().__init__()
+        if alternating and not summary_tokens:
+            raise ValueError("model.pfn_alternating needs model.pfn_summary_tokens")
         hid_size = encoder.hid_size
-        self.max_classes = max_classes
+        self.max_classes, self.normalize = max_classes, normalize
         self.encoder = encoder
         self.label_embedding = nn.Embedding(max_classes, hid_size)
-        self.blocks = nn.ModuleList(
-            EncoderBlock(hid_size, n_heads, ffn_size, dropout) for _ in range(n_layers)
+        self.summary_tokens = summary_tokens
+        self.summary = (
+            nn.Parameter(0.02 * torch.randn(summary_tokens, hid_size))
+            if summary_tokens
+            else None
         )
+
+        def blocks(count: int) -> nn.ModuleList:
+            return nn.ModuleList(
+                EncoderBlock(hid_size, n_heads, ffn_size, dropout) for _ in range(count)
+            )
+
+        self.event_blocks = blocks(len(encoder.blocks)) if alternating else None
+        self.blocks = blocks(n_layers)
         self.norm = nn.LayerNorm(hid_size)
         self.classifier = nn.Linear(hid_size, max_classes)
 
-    def embed(self, batch: SequenceBatch) -> torch.Tensor:
-        """``[events, hid_size]``: the mean encoding of each event's objects."""
-        mask = batch.token_mask
-        encoded = self.encoder(batch.tokens, mask)
-        counts = mask.sum(dim=1, keepdim=True).clamp_min(1)
-        return (encoded * mask.unsqueeze(-1)).sum(dim=1) / counts
-
-    def context(
-        self, embedded: torch.Tensor, labels: torch.Tensor
-    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
-        """Each block's keys and values of the context events
-        ``embedded [tasks, events, hid_size]`` with ``labels``."""
-        x = embedded + self.label_embedding(labels)
-        memory = []
-        for block in self.blocks:
-            q, k, v = block.project(x)
-            memory.append((k, v))
-            x = block.attend(x, q, k, v)
-        return memory
-
-    def queries(
+    def _labelled(self, x: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        """``x [events, tokens, hid_size]`` plus the label embedding of the
+        events whose ``labels`` are not -1 (the queries')."""
+        known = (labels >= 0).unsqueeze(-1)
+        return x + (self.label_embedding(labels.clamp_min(0)) * known).unsqueeze(1)
+
+    def _object_mask(self, batch: SequenceBatch) -> torch.Tensor:
+        """The attention mask of :class:`ObjectTransformer`, with the summary
+        tokens first (no pair bias)."""
+        count, mask = self.summary_tokens, batch.token_mask
+        mask = nn.functional.pad(mask, (count, 0), value=True)[:, None, None, :]
+        if self.encoder.pair_encoder is None:
+            return mask
+        bias = self.encoder.pair_encoder(pair_features(batch.tokens))
+        bias = nn.functional.pad(bias.permute(0, 3, 1, 2), (count, 0, count, 0))
+        return bias.masked_fill(~mask, float("-inf"))
+
+    @staticmethod
+    def _across(
+        block: EncoderBlock,
+        x: torch.Tensor,
+        tasks: int,
+        context: int | None,
+        memory: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        """``x [events, tokens, hid_size]`` of ``tasks`` equal tasks after
+        ``block`` attends across events: to the tokens of the first
+        ``context`` events of each task, or to the keys and values
+        ``memory``. Also returns those keys and values."""
+        events, tokens, hid_size = x.shape
+        x = x.reshape(tasks, -1, hid_size)
+        q, k, v = block.project(x).unbind(0)
+        memory = memory or (k[:, :, : context * tokens], v[:, :, : context * tokens])
+        x = block.attend(x, q, *memory)
+        return x.view(events, tokens, hid_size), memory
+
+    def encode(
         self,
-        embedded: torch.Tensor,
-        memory: list[tuple[torch.Tensor, torch.Tensor]],
-        num_classes: torch.Tensor,
-    ) -> torch.Tensor:
-        """Logits ``[tasks, events, max_classes]`` of the query events
-        ``embedded`` given the context ``memory`` of :meth:`context`."""
-        x = embedded
-        for block, (k, v) in zip(self.blocks, memory):
-            x = block.attend(x, block.project(x)[0], k, v)
-        logits = self.classifier(self.norm(x))
+        batches: list[SequenceBatch],
+        labels: list[torch.Tensor],
+        *,
+        tasks: int = 1,
+        context: int | None = None,
+        memory: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+    ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]:
+        """The final tokens ``[events, tokens, hid_size]`` of the events of
+        ``batches`` (in order, ``tasks`` equal tasks) with ``labels`` (-1:
+        a query), and the keys and values of every block across events.
+        Across events, they attend to the first ``context`` events of each
+        task or to ``memory`` (of an earlier call with the context)."""
+        count = self.summary_tokens
+        hidden = [self.encoder.embedding(batch.tokens) for batch in batches]
+        if count:
+            hidden = [
+                self._labelled(
+                    torch.cat((self.summary.expand(len(x), -1, -1), x), 1), part
+                )
+                for x, part in zip(hidden, labels)
+            ]
+        # The pair bias of one batch is kept; that of many (a large context)
+        # is recomputed by every block.
+        masks = [self._object_mask(batches[0])] if len(batches) == 1 else None
+        memory = iter(memory or [])
+        kept = []
+        for layer, block in enumerate(self.encoder.blocks):
+            hidden = [
+                block(x, masks[0] if masks else self._object_mask(batch))
+                for x, batch in zip(hidden, batches)
+            ]
+            if self.event_blocks is None:
+                continue
+            summary = torch.cat([x[:, :count] for x in hidden])
+            summary, kv = self._across(
+                self.event_blocks[layer], summary, tasks, context, next(memory, None)
+            )
+            kept.append(kv)
+            parts = summary.split([len(x) for x in hidden])
+            hidden = [torch.cat((p, x[:, count:]), 1) for p, x in zip(parts, hidden)]
+        if count:
+            x = self.encoder.norm(torch.cat([x[:, :count] for x in hidden]))
+        else:
+            pooled = []
+            for x, batch in zip(hidden, batches):
+                mask = batch.token_mask.unsqueeze(-1)
+                encoded = self.encoder.norm(x) * mask
+                pooled.append(encoded.sum(1) / mask.sum(1).clamp_min(1))
+            x = self._labelled(torch.cat(pooled).unsqueeze(1), torch.cat(labels))
+        for block in self.blocks:
+            x, kv = self._across(block, x, tasks, context, next(memory, None))
+            kept.append(kv)
+        return x, kept
+
+    def logits(self, x: torch.Tensor, num_classes: torch.Tensor) -> torch.Tensor:
+        """Logits ``[tasks, events, max_classes]`` of the final tokens of
+        :meth:`encode`, for tasks of ``num_classes`` classes."""
+        logits = self.classifier(self.norm(x.mean(1)))
+        logits = logits.view(len(num_classes), -1, self.max_classes)
         classes = torch.arange(self.max_classes, device=logits.device)
         absent = classes >= num_classes.reshape(-1, 1, 1)
         return logits.masked_fill(absent, ABSENT_LOGIT)
 
     def forward(self, batch: TaskBatch) -> torch.Tensor:
-        tasks, size = batch.context_labels.shape
-        embedded = self.embed(batch.events).view(tasks, -1, self.encoder.hid_size)
-        memory = self.context(embedded[:, :size], batch.context_labels)
-        logits = self.queries(embedded[:, size:], memory, batch.num_classes)
+        tasks, context = batch.context_labels.shape
+        events = batch.events
+        size = len(events.tokens) // tasks
+        if self.normalize:
+            sums, counts = log_pt_sums(events)
+            scale = normalizing_scale(
+                sums.view(tasks, -1)[:, :context].sum(1),
+                counts.view(tasks, -1)[:, :context].sum(1),
+            )
+            events = transformed(events, scale.repeat_interleave(size))
+        labels = nn.functional.pad(batch.context_labels, (0, size - context), value=-1)
+        x, _ = self.encode([events], [labels.flatten()], tasks=tasks, context=context)
+        logits = self.logits(x, batch.num_classes)[:, context:]
         return logits.flatten(0, 1)
 
 
+#: The geometric-mean object pT (in the units of the tokens, 10 GeV) that
+#: ``model.pfn_normalize`` gives the context of every task: about that of
+#: the Open Data prior (48 GeV), GamGam (52 GeV) and Delphes (63 GeV).
+PT_REFERENCE = 5.0
+
+
+def log_pt_sums(batch: SequenceBatch) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per event, the sum of ``log pT`` of its objects with pT > 0, and
+    their number."""
+    pt = batch.tokens[..., PT]
+    real = batch.token_mask & (pt > 0)
+    return torch.where(real, pt.clamp_min(1e-12).log(), 0).sum(1), real.sum(1)
+
+
+def normalizing_scale(sums: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+    """The scale of momenta and energies that makes the geometric-mean pT
+    of the objects behind ``sums`` and ``counts`` (:func:`log_pt_sums`)
+    ``PT_REFERENCE``."""
+    mean = sums / counts.clamp_min(1)
+    return torch.exp(math.log(PT_REFERENCE) - mean)
+
+
+def transformed(
+    batch: SequenceBatch,
+    scale: torch.Tensor | float = 1.0,
+    rotation: float = 0.0,
+    flip: float = 1.0,
+) -> SequenceBatch:
+    """``batch`` with momenta and energies scaled by ``scale`` (one number,
+    or one per event), azimuths rotated by ``rotation`` and pseudorapidities
+    multiplied by ``flip``."""
+    events = len(batch.tokens)
+    device = batch.tokens.device
+
+    def per_event(value: torch.Tensor | float) -> torch.Tensor:
+        return torch.as_tensor(value, device=device, dtype=torch.float).expand(events)
+
+    tokens = augment_tokens(
+        batch.tokens,
+        batch.token_mask,
+        scale=per_event(scale),
+        rotation=per_event(rotation),
+        flip=per_event(flip),
+    )
+    return replace(batch, tokens=tokens)
+
+
 @torch.inference_mode()
 def in_context_forward(
     model: EventPFN,
@@ -135,25 +303,60 @@ def in_context_forward(
     num_classes: int,
     batch_size: int,
     device: torch.device,
+    ensemble: int = 1,
+    seed: int = 0,
 ) -> Callable[[SequenceBatch], torch.Tensor]:
     """A function from a batch of events to their logits ``[events,
-    num_classes]``, given the labelled ``context`` events (one task)."""
+    num_classes]``, given the labelled ``context`` events (one task).
+
+    The context is encoded once. With ``ensemble`` passes, pass ``m`` rotates
+    every event by ``2 pi m / ensemble`` in azimuth, flips the
+    pseudorapidities of odd passes and numbers the classes in a random order
+    (seeded by ``seed``; the first pass keeps them); the logits are then the
+    logarithm of the mean class probabilities."""
     if num_classes > model.max_classes:
         raise ValueError(f"the PFN tells apart at most {model.max_classes} classes")
     model.eval()
-    embedded = torch.cat(
-        [
-            model.embed(collate_sequence_batch(events, part).to(device))
-            for part in context.split(batch_size)
-        ]
-    )
-    labels = events.labels[context].to(device, torch.long)
-    memory = model.context(embedded.unsqueeze(0), labels.unsqueeze(0))
+    batches = [
+        collate_sequence_batch(events, part).to(device)
+        for part in context.split(batch_size)
+    ]
+    labels = events.labels[context].to(device, torch.long).split(batch_size)
+    scale = 1.0
+    if model.normalize:
+        sums, counts = zip(*map(log_pt_sums, batches))
+        scale = normalizing_scale(torch.cat(sums).sum(), torch.cat(counts).sum())
+    generator = torch.Generator().manual_seed(seed)
+    passes = []
+    for member in range(ensemble):
+        order = (
+            torch.randperm(num_classes, generator=generator).to(device)
+            if member
+            else torch.arange(num_classes, device=device)
+        )
+        rotation = 2 * math.pi * member / ensemble
+        flip = -1.0 if member % 2 else 1.0
+        _, memory = model.encode(
+            [transformed(batch, scale, rotation, flip) for batch in batches],
+            [order[part] for part in labels],
+            context=len(context),
+        )
+        passes.append((order, rotation, flip, memory))
     classes = torch.tensor([num_classes], device=device)
 
+    def logits(batch: SequenceBatch, order, rotation, flip, memory) -> torch.Tensor:
+        batch = transformed(batch, scale, rotation, flip)
+        queries = torch.full((len(batch.tokens),), -1, device=device)
+        x, _ = model.encode([batch], [queries], memory=memory)
+        # Class c of the dataset is class order[c] of this pass.
+        return model.logits(x, classes)[0, :, order]
+
+    @torch.inference_mode()
     def forward(batch: SequenceBatch) -> torch.Tensor:
-        logits = model.queries(model.embed(batch).unsqueeze(0), memory, classes)
-        return logits[0, :, :num_classes]
+        if len(passes) == 1:
+            return logits(batch, *passes[0])
+        probabilities = [logits(batch, *one).softmax(-1) for one in passes]
+        return torch.stack(probabilities).mean(0).log()
 
     return forward
 
@@ -161,6 +364,45 @@ def in_context_forward(
 # --- The prior ----------------------------------------------------------------
 
 
+def event_observables(events: PackedEvents) -> torch.Tensor:
+    """``[events, 4 * types + 1]`` observables of every event: per object
+    type, the number of objects and ``log(1 + x)`` of their scalar sum pT,
+    leading pT and invariant mass; then the number of b-tagged objects."""
+    features = events.object_features
+    counts = events.object_offsets.diff()
+    event = torch.repeat_interleave(torch.arange(len(counts)), counts)
+    types = int(features[:, NODE_TYPE].max()) + 1
+    key = event * types + features[:, NODE_TYPE].long()
+    pt, eta, phi = features[:, PT], features[:, ETA], features[:, PHI]
+
+    def total(value: torch.Tensor) -> torch.Tensor:
+        sums = torch.zeros(len(counts) * types, dtype=torch.float64)
+        return sums.index_add_(0, key, value.double()).view(-1, types)
+
+    leading = torch.zeros(len(counts) * types).scatter_reduce_(0, key, pt, "amax")
+    mass2 = total(features[:, ENERGY]).square() - sum(
+        total(p).square()
+        for p in (pt * torch.cos(phi), pt * torch.sin(phi), pt * torch.sinh(eta))
+    )
+    return torch.cat(
+        (
+            total(torch.ones_like(pt)).float(),
+            total(pt).float().log1p(),
+            leading.view(-1, types).log1p(),
+            mass2.clamp_min(0).sqrt().float().log1p(),
+            torch.zeros(len(counts)).index_add_(0, event, features[:, BTAG])[:, None],
+        ),
+        dim=1,
+    )
+
+
+def _dirichlet(count: int, generator: torch.Generator) -> torch.Tensor:
+    """Dirichlet(1) proportions of ``count`` parts."""
+    # rand can return exactly 0, whose -log is infinite.
+    share = -torch.log(torch.rand(count, generator=generator).clamp_min(1e-12))
+    return share / share.sum()
+
+
 @dataclass(frozen=True)
 class _Step:
     """The sampled events and augmentations of one meta-training step."""
@@ -263,6 +505,13 @@ class TaskLoader:
         ]
         if len(self.pools) < 2 * config.mixture_labels[0]:
             raise ValueError("the PFN prior needs events of at least two classes")
+        # Standardized observables of the events, for synthetic tasks.
+        self.observables = None
+        if train and config.synthetic_tasks:
+            observables = event_observables(events)
+            used = observables[torch.cat(self.pools)]
+            mean, std = used.mean(0), used.std(0).clamp_min(1e-6)
+            self.observables = (observables - mean) / std
 
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
@@ -291,15 +540,22 @@ class TaskLoader:
         largest = min(config.task_classes[1], len(self.pools) // fewest, context)
         smallest = min(config.task_classes[0], largest)
 
-        def dirichlet(count: int) -> torch.Tensor:  # Dirichlet(1) proportions
-            # rand can return exactly 0, whose -log is infinite.
-            draw = torch.rand(count, generator=generator).clamp_min(1e-12)
-            share = -torch.log(draw)
-            return share / share.sum()
+        def dirichlet(count: int) -> torch.Tensor:
+            return _dirichlet(count, generator)
 
         indices, labels, num_classes = [], [], []
         for _ in range(tasks):
             k = int(torch.randint(smallest, largest + 1, (), generator=generator))
+            # No draw without synthetic tasks: earlier priors are unchanged.
+            if (
+                self.observables is not None
+                and float(torch.rand((), generator=generator)) < config.synthetic_tasks
+            ):
+                task_indices, task_labels = self._synthetic(k, size, generator)
+                indices.append(task_indices)
+                labels.append(task_labels)
+                num_classes.append(k)
+                continue
             # Each class mixes 1 to ``most`` labels, no label in two classes.
             # Without mixing, no draws are added: earlier priors are unchanged.
             mixed = (
@@ -356,6 +612,48 @@ class TaskLoader:
             flip=flip,
         )
 
+    def _synthetic(
+        self, k: int, size: int, generator: torch.Generator
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """The events and labels of a synthetic task of ``k`` classes:
+        events of 1 to 3 processes (Dirichlet proportions), classified by
+        ranges of a random one-hidden-layer network (16 tanh units) of 1 to
+        4 random observables, the ranges holding Dirichlet fractions of the
+        task's events. The first ``k`` events have the ``k`` labels."""
+
+        def dirichlet(count: int) -> torch.Tensor:
+            return _dirichlet(count, generator)
+
+        parts = int(torch.randint(1, 4, (), generator=generator))
+        chosen = torch.randperm(len(self.pools), generator=generator)[:parts]
+        part = torch.multinomial(dirichlet(parts), size, True, generator=generator)
+        task_indices = torch.empty(size, dtype=torch.long)
+        for which, pool_index in enumerate(chosen.tolist()):
+            pool, members = self.pools[pool_index], part == which
+            draw = torch.randint(len(pool), (int(members.sum()),), generator=generator)
+            task_indices[members] = pool[draw]
+        observables = self.observables[task_indices]
+        inputs = min(4, observables.shape[1])
+        inputs = int(torch.randint(1, inputs + 1, (), generator=generator))
+        columns = torch.randperm(observables.shape[1], generator=generator)[:inputs]
+        hidden = torch.randn(inputs, 16, generator=generator)
+        bias, output = torch.randn(2, 16, generator=generator)
+        value = torch.tanh(observables[:, columns] @ hidden + bias) @ output
+        # Ranks, ties (as of object counts) broken at random.
+        tie = torch.rand(size, generator=generator)
+        order = torch.argsort(value + 1e-6 * value.std() * tie)
+        rest = torch.multinomial(dirichlet(k), size - k, True, generator=generator)
+        sizes = torch.bincount(rest, minlength=k) + 1
+        task_labels = torch.empty(size, dtype=torch.long)
+        task_labels[order] = torch.repeat_interleave(
+            torch.randperm(k, generator=generator), sizes
+        )
+        first = torch.stack([(task_labels == c).nonzero()[0, 0] for c in range(k)])
+        others = torch.ones(size, dtype=torch.bool)
+        others[first] = False
+        arranged = torch.cat((first, others.nonzero().flatten()))
+        return task_indices[arranged], task_labels[arranged]
+
     def __iter__(self) -> Iterator[Any]:
         generator = self._generator()
         steps = _Steps(self.events, [self.sample(generator) for _ in range(len(self))])

tests/test_pfn.py

@@ -4,13 +4,21 @@ inference and the Delphes-to-GamGam workflow."""
 import json
 from dataclasses import replace
 
+import pytest
 import torch
 from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events
 
 from gnn4colliders.cli import main
 from gnn4colliders.config import ModelConfig, TaskConfig, save_config
+from gnn4colliders.data import pack_events
 from gnn4colliders.models import build_model
-from gnn4colliders.pfn import ABSENT_LOGIT, TaskLoader, in_context_forward
+from gnn4colliders.pfn import (
+    ABSENT_LOGIT,
+    TaskBatch,
+    TaskLoader,
+    event_observables,
+    in_context_forward,
+)
 from gnn4colliders.sequences import collate_sequence_batch
 
 SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True}
@@ -37,12 +45,35 @@ def _loader(events, train=True, **prior):
     return TaskLoader(events, torch.arange(len(events)), config, 64, train=train)
 
 
-def _pfn():
+#: The mean-pooled PFN, summary tokens, alternating attention and the
+#: per-task scale of momenta.
+VARIANTS = {
+    "pooled": {},
+    "summary": {"pfn_summary_tokens": 2},
+    "alternating": {"pfn_summary_tokens": 2, "pfn_alternating": True},
+    "normalized": {"pfn_summary_tokens": 1, "pfn_normalize": True},
+}
+
+
+def _pfn(**options):
     torch.manual_seed(0)
-    config = ModelConfig(name="pfn", pfn_layers=2, max_classes=4, **SMALL)
+    config = ModelConfig(name="pfn", pfn_layers=2, max_classes=4, **SMALL, **options)
     return build_model(config, num_features=7, num_node_types=5).eval()
 
 
+def _task(events, context, queries):
+    """One task: the ``context`` events with their labels, then ``queries``."""
+    indices = torch.cat((context, queries))
+    return TaskBatch(
+        events=collate_sequence_batch(events, indices),
+        context_labels=events.labels[context][None].long(),
+        num_classes=torch.tensor([4]),
+        labels=events.labels[queries],
+        weight=torch.ones(len(queries)),
+        event_index=queries,
+    )
+
+
 def test_tasks_follow_the_prior():
     events = _events()
     for batch in _loader(events, prior_classes=[1, 2, 3]):
@@ -112,8 +143,10 @@ def test_training_tasks_vary_by_epoch_and_validation_tasks_do_not():
     assert torch.equal(first(validation, 0), first(validation, 1))
 
 
-def test_queries_see_only_the_context():
-    model, batch = _pfn(), next(iter(_loader(_events(), train=False)))
+@pytest.mark.parametrize("variant", VARIANTS)
+def test_queries_see_only_the_context(variant):
+    model = _pfn(**VARIANTS[variant])
+    batch = next(iter(_loader(_events(), train=False)))
     tasks, context = batch.context_labels.shape
     tokens = batch.events.tokens.clone()
     first_query = context  # of the first task
@@ -129,21 +162,92 @@ def test_queries_see_only_the_context():
     assert (expected[absent] == ABSENT_LOGIT).all()
 
 
-def test_in_context_inference_matches_a_task():
-    model, events = _pfn(), _events()
+@pytest.mark.parametrize("variant", VARIANTS)
+def test_in_context_inference_matches_a_task(variant):
+    model, events = _pfn(**VARIANTS[variant]), _events()
     context, queries = torch.arange(0, 40), torch.arange(40, 60)
+    # The context in three parts, the queries in two.
     forward = in_context_forward(
         model, events, context, num_classes=4, batch_size=16, device=torch.device("cpu")
     )
     with torch.no_grad():
-        actual = forward(collate_sequence_batch(events, queries))
-        everything = collate_sequence_batch(events, torch.cat((context, queries)))
-        embedded = model.embed(everything).unsqueeze(0)
-        memory = model.context(embedded[:, :40], events.labels[context][None])
-        expected = model.queries(embedded[:, 40:], memory, torch.tensor([4]))[0]
+        actual = torch.cat(
+            [forward(collate_sequence_batch(events, q)) for q in queries.split(16)]
+        )
+        expected = model(_task(events, context, queries))
     torch.testing.assert_close(actual, expected)
 
 
+def test_normalized_pfn_ignores_the_energy_scale():
+    model, events = _pfn(**VARIANTS["normalized"]), _events()
+    scaled = replace(events, object_features=events.object_features.clone())
+    scaled.object_features[:, [0, 3]] *= 1.3
+    context, queries = torch.arange(0, 40), torch.arange(40, 60)
+    with torch.no_grad():
+        expected = model(_task(events, context, queries))
+        actual = model(_task(scaled, context, queries))
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
+
+
+def test_ensemble_averages_probabilities_over_transformed_passes():
+    model, events = _pfn(**VARIANTS["alternating"]), _events()
+    context, queries = torch.arange(0, 40), torch.arange(40, 60)
+
+    def logits(ensemble):
+        forward = in_context_forward(
+            model,
+            events,
+            context,
+            num_classes=4,
+            batch_size=64,
+            device=torch.device("cpu"),
+            ensemble=ensemble,
+        )
+        return forward(collate_sequence_batch(events, queries))
+
+    single, ensemble = logits(1), logits(4)
+    torch.testing.assert_close(ensemble.exp().sum(1), torch.ones(20))
+    assert not torch.allclose(ensemble, single.log_softmax(1))
+    torch.testing.assert_close(logits(4), ensemble)
+
+
+def test_synthetic_tasks_rank_events_of_few_labels():
+    events = _events()
+    loader = _loader(events, synthetic_tasks=1.0, prior_classes=None)
+    for batch in loader:
+        for task, classes in enumerate(_originals(events, batch)):
+            k = int(batch.num_classes[task])
+            assert set(batch.context_labels[task].tolist()) == set(range(k))
+            assert len(set().union(*classes)) <= 3
+    # Without synthetic tasks the prior draws exactly as before.
+    plain, zero = _loader(events), _loader(events, synthetic_tasks=0.0)
+    first = plain.sample(torch.Generator().manual_seed(1))
+    assert torch.equal(
+        first.indices, zero.sample(torch.Generator().manual_seed(1)).indices
+    )
+
+
+def test_event_observables_per_object_type():
+    # Two photons back to back (pT 3 and 4, eta 0), then MET of pT 2.
+    objects = [
+        torch.tensor(
+            [
+                [3.0, 0.0, 0.0, 3.0, 0.0, 0.0, 3.0],
+                [4.0, 0.0, torch.pi, 4.0, 1.0, 0.0, 3.0],
+                [2.0, 0.0, 1.0, 2.0, 0.0, 0.0, 4.0],
+            ]
+        )
+    ]
+    observables = event_observables(pack_events(objects, torch.tensor([0])))[0]
+    counts, total, leading, mass = observables[:20].view(4, 5)
+    assert counts.tolist() == [0, 0, 0, 2, 1]
+    torch.testing.assert_close(total[3], torch.log1p(torch.tensor(7.0)))
+    torch.testing.assert_close(leading[3], torch.log1p(torch.tensor(4.0)))
+    # m^2 = (3 + 4)^2 - (3 - 4)^2 = 48
+    torch.testing.assert_close(mass[3], torch.log1p(torch.tensor(48.0).sqrt()))
+    assert observables[20] == 1  # b-tagged objects
+
+
 def test_encoder_starts_from_a_jepa_checkpoint(tmp_path):
     jepa = build_model(
         ModelConfig(name="jepa", **SMALL), num_features=7, num_node_types=5