pio-arch 0.2.3

Perceiver IO-style permutation-invariant set models in PyTorch.

pio-arch

Perceiver IO-style permutation-invariant set models in PyTorch, for tabular, mixed-row, and sentence/set-structured data.

pio-arch provides a small set of composable building blocks for problems where each sample is a variable-length bag of heterogeneous features (numeric, categorical, free text) and the model must produce one or more predictions per sample.

The flagship model — PIOAttentionModel — is structurally Perceiver IO:

• learnable latent queries cross-attend to the variable-length context
• optional self-attention on the latent stack
• learnable task queries cross-attend to the latents
• per-task or shared MLP heads produce one prediction per task

All stage dimensions (context, latent, task query) are independently configurable. Permutation invariance over the input set is preserved end-to-end and is enforced by tests.

---

Install

Base install (PyTorch + numpy only):

pip install pio-arch

With the data-science / notebook extras (LightGBM, sentence-transformers, Optuna, polars, etc.):

pip install "pio-arch[model-dev]"

Python 3.12+ required.

---

Quick start

import torch
from pio_arch.models.pio_attention import PIOAttentionModel
from pio_arch.models.sentence_encoder import TaskSpec, masked_multitask_loss

tasks = [
    TaskSpec("fraud", "binary"),
    TaskSpec("amount", "regression", weight=0.5),
    TaskSpec("time_to_fraud", "survival_discrete", output_dim=10),  # 10 day bins
]

model = PIOAttentionModel(
    input_dim=384,          # raw context-row dim (e.g. sentence-transformer output)
    tasks=tasks,
    model_dim=64,           # context stage
    latent_dim=64,          # latent bottleneck (defaults to model_dim)
    task_dim=32,            # task query stage (defaults to latent_dim)
    num_latents=8,
    num_context_self_attn_blocks=1,
    num_latent_self_attn_blocks=1,
    num_heads=4,
)

context = torch.randn(4, 12, 384)               # [batch, n_rows, input_dim]
padding_mask = torch.zeros(4, 12, dtype=torch.bool)  # True at padded rows

predictions = model(context, padding_mask)
# predictions = {"fraud": [4, 1], "amount": [4, 1], "time_to_fraud": [4, 10]}

from pio_arch.models.survival import build_discrete_hazard_targets
fraud_targets = torch.rand(4)
amount_targets = torch.randn(4)
ttf_targets, ttf_mask = build_discrete_hazard_targets(
    torch.tensor([2, 5, 0, 9]), torch.tensor([True, False, True, False]), num_bins=10
)
targets = torch.cat([fraud_targets.unsqueeze(1), amount_targets.unsqueeze(1), ttf_targets], dim=1)
target_mask = torch.cat([torch.ones(4, 2, dtype=torch.bool), ttf_mask], dim=1)
loss, logs = masked_multitask_loss(predictions, targets, target_mask, tasks)
loss.backward()

A native PyTorch training loop is available in pio_arch.utils.train:

from pio_arch.utils.train import Trainer, make_masked_multitask_loss_fn

trainer = Trainer(
    model=model,
    optimizer=torch.optim.AdamW(model.parameters(), lr=3e-4),
    loss_fn=make_masked_multitask_loss_fn(tasks),
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)
trainer.fit(train_loader, val_loader, num_epochs=5)

---

What's in the box

Module	What it provides
pio_arch.models.embedder	UniversalFeatureTransformer — type-aware raw-feature embedder (RFF for numeric, learned embedding for discrete).
pio_arch.models.context_encoder	ContextRowEncoder — mixed numeric / discrete / projected-text row encoder with feature-type embeddings.
pio_arch.models.sentence_encoder	SentenceSetEncoder (masked-mean baseline), SentenceSetMultiTaskModel, MultiTaskHead, TaskSpec, masked_multitask_loss.
pio_arch.models.pio_attention	PIOAttentionModel (Perceiver IO with task queries), SelfAttentionBlock, CrossAttentionBlock.
pio_arch.models.contrastive	ContrastiveTaskSpec, ContrastiveHead, contrastive_infonce_loss — SimCLR-style InfoNCE auxiliary on pooled latents (v0.2.0).
pio_arch.models.decorrelation	DecorrelationSpec, feature_decorrelation_loss — optional two-sided covariance penalty that decorrelates the entries of each task's output-query features within the batch (v0.2.2).
pio_arch.models.survival	build_discrete_hazard_targets, survival_from_hazards, weibull_nll, weibull_survival, and friends — discrete-time hazard and Weibull AFT primitives backing the "survival_discrete" and "survival_weibull" task kinds (v0.2.1).
pio_arch.utils.rff	Shared Random Fourier Feature helpers.
pio_arch.utils.collate	sentence_set_collate — tensor-only collate.
pio_arch.utils.data	SentenceSetDataset + collate_sentence_set_batch for trainer-ready dict batches.
pio_arch.utils.train	Trainer, train_one_epoch, evaluate, make_masked_multitask_loss_fn (with epoch forwarding for loss-term annealing, v0.2.2).
pio_arch.utils.input_dropout	InputRowDropout — training-time input row-dropout with plain per-row, grouped (group_ids, co-missing related rows), and custom drop_fn modes; plus sample_row_drop_mask / combine_row_drops helpers (v0.2.3).
pio_arch.utils.text_dropout	TextRowDropout — plain per-row dropout for robustness (now a thin back-compat wrapper over InputRowDropout).
pio_arch.utils.sentence_pool	pool_sentence_embeddings — DeepSets-style aggregator collapsing [B, N, D] → [B, D].

---

Notebooks

End-to-end walkthroughs live under notebooks/:

• sentence_only_pio.ipynb — sentence-only PIO walkthrough using a frozen MiniLM sentence-transformer. Compares the masked-mean baseline, PIOAttentionModel, and a DeepSets-style pooled baseline.
• full_pio_walkthrough.ipynb — full mixed-input pipeline with numeric, categorical, list-of-codes, and list-of-reason-sentences features.
• optuna_sweep.ipynb — Optuna hyperparameter sweep over the PIO architecture, with typical Perceiver IO dim guidance.
• contrastive_multitask_walkthrough.ipynb — v0.2.0 walkthrough for the contrastive auxiliary task, multi-target heads, Kendall σ trajectory plot, and three freeze helpers for fine- tuning workflows.
• discrete_hazard_walkthrough.ipynb — v0.2.1 walkthrough for the "survival_discrete" task kind: discrete- time hazard math, target/mask construction, plotting S(t) curves, and joint contrastive + survival training.
• weibull_aft_walkthrough.ipynb — v0.2.1 walkthrough for the "survival_weibull" task kind: parametric AFT with right-censoring and continuous-time inference grids.
• decorrelation_walkthrough.ipynb — v0.2.2 walkthrough for the optional feature-decorrelation regularizer: the two-sided ‖C − I‖² covariance penalty, the variance-floor rationale, the B ≤ task_dim rank-deficiency limitation, linear-warmup annealing through the trainer, configurable task subsets, and the optional Kendall σ slot.
• input_dropout_walkthrough.ipynb — v0.2.3 walkthrough for InputRowDropout, contrasting its three modes: plain per-row dropout, grouped co-drop via group_ids (related rows that go missing together; negative ids = independent), and a custom drop_fn callable. Covers why row dropout needs no 1/(1-p) rescaling, the ≥1-survivor guarantee, and composing it with the Trainer via a custom loss function.

A synthetic 50k-row dataset (notebooks/pio_synthetic_50k.parquet) ships with the repo so every notebook can run end-to-end out of the box.

---

Reading the code

LEARNING_GUIDE.md is a staged tour of the source — which files to read in what order, what to note in each, and self-check exercises. Designed for ~4 hours start-to-finish with the exercises, or ~1 hour as a skim.

---

Design notes

• context.md — reference material on the broader permutation-invariant architecture design space (DeepSets, Set Transformer, Perceiver, Slot Attention). Historical / background only — not all of it is implemented in this repo.
• notebooks/pio_encoder_design.md — current encoder + attention design rationale.
• notebooks/pio_text_embedding_plan.md — sentence-transformer integration, list-valued field semantics, deployment considerations.
• AGENT_GUIDE.shared.md — tool-neutral project rules for any AI coding agent or human reading the repo for the first time.

---

Development

# clone, then
python3.12 -m venv .venv
.venv/bin/pip install -e ".[model-dev]" --group dev
.venv/bin/pre-commit install

# run tests + coverage
.venv/bin/pytest

# lint + format
.venv/bin/ruff check pio_arch tests
.venv/bin/ruff format --check pio_arch tests

# run the full pre-commit suite
.venv/bin/pre-commit run --all-files

The package currently has 162 unit + integration tests at 100% line coverage on the runtime code.

To rebuild the walkthrough notebooks after editing them:

.venv/bin/python scripts/build_walkthrough_notebooks.py

To smoke-test the notebook code paths on a 2k-row subsample without running the full epochs:

.venv/bin/python scripts/smoke_test_notebooks.py

---

License

MIT. See LICENSE.