Set transformer benchmark: ISAB vs DeepSet, MLA and full attention¶

Compares equivariant set encoders on small synthetic set-learning problems and checks the exact invariances InducedSelfAttention is supposed to satisfy. All numbers below were produced on Apple Silicon CPU (JAX 0.9.2, float32).

JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_set_transformers --steps 5000 \
  --output docs/benchmarks/set_transformers.json

Measurement settings¶

Budget per arm and task: 5000 training steps, 3 seeds, batch 64, eval batch 512, Adam 3e-3 with global-norm clip 5. Models are 2 layers wide 64, four heads, ISAB/MLA arms use 32 inducing points, the in-context task uses 48 context + 16 query tokens (N = 64), both set tasks use N = 64. Runtime scaling uses batch four, width 64, two layers and 64 inducing points at N = 64 / 256 / 1024 / 2048.

The arms:

ISAB blocks use zero_init=True (the default): each MAB's attention output projection and last feedforward layer start at zero, so a block begins as the identity. They also share one norm between each MAB's query and key/value. ISAB blocks run two MAB stages each, so their parameter counts are roughly double the single-attention arms (205k vs 100k at this width).

Invariance audit¶

Checked on one fixed model and input; these are exactness checks, not trainable metrics.

The padding row is the regression guard for the kv_len bug: MAB 2 attends over the inducing bank, not over kv_len keys, so query scaling must resolve kv_len from the mask for MAB 1 only.

Task quality (mean test MSE over 3 seeds)¶

Lower is better. Baselines are computed per seed on the eval batch.

cluster_denoise (N=64, identity 0.090, oracle 0.003)¶

padded_outlier_mean (N in [16, 64], naive mean 0.040, oracle 0.002)¶

in_context_regression (48 context + 16 query, zero prediction 4.6)¶

Observations:

• At 5000 steps the best ISAB variant beats sab, mla and deepset on all three tasks, while paying O(N M) instead of O(N²).
• On in_context_regression the gap is largest: routing context through the inducing bank is an inductive bias for this task.
• Differences between ISAB variants are often within the seed spread (3 seeds), and the best query scaling varies by task.
• Zero init was compared against the default init at the same budget (not committed): it was neutral overall, winning some arm/task pairs and losing others. At 1500 steps it was clearly slower to converge.

Runtime scaling (warm ms, batch 4)¶

ISAB grows linearly with N at fixed M=64 (about 7x cheaper forward and 9x cheaper backward than sab at N=2048) and its query scaling adds no measurable overhead. mla still pays O(N²) attention over the shared latent and is the slowest quadratic arm at large N. CPU timings do not establish GPU performance.

Re-run with a smaller budget¶

JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_set_transformers \
  --steps 50 --seeds 1 --batch-size 16 --set-size 32 --num-inducing 8 \
  --model-dim 32 --tokens 32 128 --output docs/benchmarks/set_smoke.json

Write smoke results to a distinct filename so the published measurements above stay intact. See the benchmark index for the general caveats.