Claude Opus 5.5
feat: TabPFN-3.5 options for the PFN, a cosine schedule and faster meta-training
71aefe8 | """The prior-fitted network: the task prior, context-only attention, in-context | |
| 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 NODE_TYPE, build_model | |
| from gnn4colliders.pfn import ( | |
| ABSENT_LOGIT, | |
| TaskBatch, | |
| TaskLoader, | |
| context_ranks, | |
| event_observables, | |
| in_context_forward, | |
| rank_keys, | |
| ) | |
| from gnn4colliders.sequences import ENERGY, PT, SequenceBatch, collate_sequence_batch | |
| SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True} | |
| PRIOR = { | |
| "name": "pfn", | |
| "task_classes": [2, 3], | |
| "context_events": [8, 16], | |
| "query_events": 4, | |
| "steps_per_epoch": 3, | |
| "validation_steps": 2, | |
| } | |
| def _events(): | |
| events = random_events(200, classes=4) | |
| events.object_features[:, 6] = torch.randint( | |
| 0, 5, (events.object_features.shape[0],) | |
| ).float() | |
| return events | |
| def _loader(events, train=True, **prior): | |
| config = TaskConfig(**{**PRIOR, **prior}) | |
| return TaskLoader(events, torch.arange(len(events)), config, 64, train=train) | |
| #: 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}, | |
| # The options of TabPFN-3.5, with and without alternating attention. | |
| "tabpfn": { | |
| "pfn_summary_tokens": 2, | |
| "pfn_alternating": True, | |
| "pfn_concat": True, | |
| "pfn_retrieval": True, | |
| "pfn_qk_norm": True, | |
| "pfn_query_kv": True, | |
| "pfn_fourier": 3, | |
| "pfn_ranks": True, | |
| }, | |
| "retrieval": {"pfn_summary_tokens": 1, "pfn_retrieval": True, "pfn_ranks": True}, | |
| } | |
| def _pfn(**options): | |
| torch.manual_seed(0) | |
| config = ModelConfig(name="pfn", pfn_layers=2, max_classes=4, **SMALL, **options) | |
| model = build_model(config, num_features=7, num_node_types=5).eval() | |
| # Maps that start at zero (Fourier features, ranks) are tested non-zero. | |
| with torch.no_grad(): | |
| for parameter in model.parameters(): | |
| if not parameter.any(): | |
| parameter.normal_(std=0.3) | |
| return model | |
| 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]): | |
| tasks, context = batch.context_labels.shape | |
| assert 8 <= context <= 16 and tasks == 64 // (context + 4) | |
| assert ((batch.num_classes >= 2) & (batch.num_classes <= 3)).all() | |
| index = batch.events.event_index.view(tasks, -1) | |
| queries = batch.labels.view(tasks, -1) | |
| for task in range(tasks): | |
| k = int(batch.num_classes[task]) | |
| labels = torch.cat((batch.context_labels[task], queries[task])) | |
| # Every class has a context event, and each task label is one | |
| # original class out of prior_classes. | |
| assert set(batch.context_labels[task].tolist()) == set(range(k)) | |
| original = events.labels[index[task]] | |
| for label in range(k): | |
| assert len(set(original[labels == label].tolist())) == 1 | |
| assert not set(original.tolist()) - {1, 2, 3} | |
| def test_ranks_share_the_context_size_of_a_step_but_not_its_events(): | |
| events = _events() | |
| config = TaskConfig(**{**PRIOR, "context_events": [8, 40]}) | |
| first = [ | |
| next( | |
| iter(TaskLoader(events, torch.arange(200), config, 64, train=True, rank=r)) | |
| ) | |
| for r in (0, 1) | |
| ] | |
| sizes = [batch.context_labels.shape for batch in first] | |
| assert sizes[0] == sizes[1] | |
| assert not torch.equal(first[0].events.event_index, first[1].events.event_index) | |
| def _originals(events, batch): | |
| """Per task, the original labels behind each task label.""" | |
| tasks = batch.context_labels.shape[0] | |
| index = batch.events.event_index.view(tasks, -1) | |
| labels = torch.cat((batch.context_labels, batch.labels.view(tasks, -1)), 1) | |
| return [ | |
| [set(events.labels[index[t]][labels[t] == k].tolist()) for k in range(n)] | |
| for t, n in enumerate(batch.num_classes.tolist()) | |
| ] | |
| def test_held_out_labels_make_the_validation_tasks_only(): | |
| events = _events() | |
| for train, allowed in ((True, {0, 1}), (False, {2, 3})): | |
| for batch in _loader(events, train=train, holdout_classes=[2, 3]): | |
| for classes in _originals(events, batch): | |
| assert set().union(*classes) <= allowed | |
| def test_mixed_classes_draw_on_disjoint_labels(): | |
| events = _events() | |
| loader = _loader(events, task_classes=[2, 2], mixture_labels=[2, 2]) | |
| for batch in loader: | |
| for first, second in _originals(events, batch): | |
| assert 1 <= len(first) <= 2 and 1 <= len(second) <= 2 | |
| assert not first & second | |
| def test_class_proportions_survive_a_zero_random_draw(monkeypatch): | |
| rand = torch.rand | |
| monkeypatch.setattr(torch, "rand", lambda *a, **k: rand(*a, **k) * 0) | |
| loader = _loader(_events(), mixture_labels=[1, 2]) | |
| step = loader.sample(torch.Generator().manual_seed(0)) | |
| assert (step.labels >= 0).all() | |
| def test_training_tasks_vary_by_epoch_and_validation_tasks_do_not(): | |
| events = _events() | |
| def first(loader, epoch): | |
| loader.set_epoch(epoch) | |
| return next(iter(loader)).events.event_index | |
| train, validation = _loader(events), _loader(events, train=False) | |
| assert torch.equal(first(train, 0), first(train, 0)) | |
| assert not torch.equal(first(train, 0), first(train, 1)) | |
| assert torch.equal(first(validation, 0), first(validation, 1)) | |
| 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 | |
| tokens[first_query, :, :6] += 1.0 | |
| changed = replace(batch, events=replace(batch.events, tokens=tokens)) | |
| with torch.no_grad(): | |
| expected = model(batch).view(tasks, -1, 4) | |
| actual = model(changed).view(tasks, -1, 4) | |
| assert not torch.allclose(actual[0, 0], expected[0, 0]) | |
| torch.testing.assert_close(actual[0, 1:], expected[0, 1:]) | |
| torch.testing.assert_close(actual[1:], expected[1:]) | |
| absent = (torch.arange(4) >= batch.num_classes[:, None, None]).expand_as(expected) | |
| assert (expected[absent] == ABSENT_LOGIT).all() | |
| 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 = 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(ranks): | |
| model, events = _pfn(**VARIANTS["normalized"], pfn_ranks=ranks), _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_context_ranks_are_mid_ranks_per_task_and_object_type(): | |
| tokens = torch.zeros(3, 2, 7) | |
| tokens[..., PT] = torch.tensor([[1.0, 2.0], [3.0, 2.0], [2.0, 5.0]]) | |
| tokens[..., ENERGY] = 1.0 | |
| tokens[2, 1, NODE_TYPE] = 1 # the only object of its type | |
| batch = SequenceBatch( | |
| tokens=tokens, | |
| token_mask=torch.ones(3, 2, dtype=torch.bool), | |
| labels=None, | |
| weight=None, | |
| global_features=None, | |
| event_index=None, | |
| ) | |
| task = torch.tensor([0, 0, 1]) | |
| keys = rank_keys(batch, task) | |
| ranks = context_ranks(keys, keys[:2].flatten().sort().values) | |
| # pT of task 0: 1 < 2 = 2 < 3; task 1 has no context. | |
| torch.testing.assert_close( | |
| ranks[:2, :, 0], torch.tensor([[0.125, 0.5], [0.875, 0.5]]) | |
| ) | |
| torch.testing.assert_close(ranks[:2, :, 1], torch.full((2, 2), 0.5)) | |
| assert (ranks[2] == 0.5).all() | |
| 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 | |
| ) | |
| path = tmp_path / "jepa.pt" | |
| torch.save({"model_state_dict": jepa.state_dict()}, path) | |
| config = ModelConfig(name="pfn", pretrained=str(path), **SMALL) | |
| pfn = build_model(config, num_features=7, num_node_types=5) | |
| expected = jepa.context_encoder.state_dict() | |
| for name, value in pfn.encoder.state_dict().items(): | |
| assert torch.equal(value, expected[name]) | |
| def test_meta_train_on_delphes_then_classify_gamgam_in_context(tmp_path): | |
| small = [f"model.{key}={value}" for key, value in SMALL.items()] | |
| prior = [ | |
| "task.context_events=[4,16]", | |
| "task.query_events=4", | |
| "task.steps_per_epoch=2", | |
| "task.validation_steps=1", | |
| ] | |
| meta = fixture_config( | |
| "delphes_pfn", | |
| [TTH_FIXTURE, TTH_FIXTURE], | |
| [0, 1], | |
| tmp_path / "meta", | |
| "train.max_epochs=2", | |
| "model.pfn_layers=1", | |
| *small, | |
| *prior, | |
| ) | |
| path = str(save_config(meta, tmp_path / "meta.yaml")) | |
| assert main(["prepare", path]) == 0 | |
| assert main(["train", path]) == 0 | |
| history = (meta.path / "history.jsonl").read_text().splitlines() | |
| assert "accuracy" in json.loads(history[-1])["validation"] | |
| gamgam = fixture_config( | |
| "gamgam_pfn", | |
| GAMGAM_FIXTURES, | |
| list(range(5)), | |
| tmp_path / "gamgam", | |
| f"checkpoint={meta.path}/best.pt", | |
| "data.train_events=20", | |
| "model.pfn_layers=1", | |
| "split=validation", | |
| *small, | |
| ) | |
| path = str(save_config(gamgam, tmp_path / "gamgam.yaml")) | |
| assert main(["prepare", path]) == 0 | |
| assert main(["train", path]) == 2 # nothing to train on GamGam | |
| assert main(["evaluate", path]) == 0 | |
| metrics = json.loads((gamgam.path / "validation_metrics.json").read_text()) | |
| assert 0 <= metrics["metrics"]["auc"] <= 1 | |
| assert main(["predict", path]) == 0 | |
| assert main(["export", path]) == 2 | |
| def test_in_context_models_need_a_context(tmp_path): | |
| config = fixture_config( | |
| "gamgam_tabpfn", | |
| GAMGAM_FIXTURES, | |
| list(range(5)), | |
| tmp_path, | |
| "data.train_events=null", | |
| ) | |
| path = str(save_config(config, tmp_path / "tabpfn.yaml")) | |
| assert main(["prepare", path]) == 0 | |
| assert main(["train", path]) == 2 | |
| assert main(["evaluate", path]) == 2 | |