Research map / Probabilistic and scientific modeling / Molecules and proteins
Translation Equivariant Transformer Neural Processes
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces translation-equivariant transformer neural processes (TE-TNP and TE-PT-TNP) via translation-equivariant attention that updates input locations, showing gains on stationary spatio-temporal regression tasks.
Problem
Transformer neural processes lack translation equivariance, so they fail to generalise when spatio-temporal inputs are shifted outside the training range, even though stationary data call for it.
Approach
Replace MHSA/MHCA with translation-equivariant versions whose attention logits depend on pairwise input differences through a learned MLP, and update input locations equivariantly per layer; pseudo-tokens with pseudo-locations give linear complexity.
Key results
- On synthetic 1-D regression at shift Delta=1.0, TE-TNP scores -0.47 and TE-PT-TNP-M32 -0.46 log-likelihood, while TNP drops to -0.57 and CNP to -1.08.
- On T-CIFAR10 image completion, TE-PT-TNP-M128 reaches 1.51 log-likelihood versus 1.44 for ConvCNP and 1.46 for PT-TNP-M128.
- On ERA5 western Europe test region, TE-PT-TNP achieves 1.27 average log-likelihood versus 1.22 for GP and -0.40 for PT-TNP, and it is the best model in all three regions.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace standard attention with translation-equivariant attention | TNP · Synthetic 1-D regression, shifted test inputs (Δ=1.0) | average log-likelihood: -0.57 ± 0.01 -> -0.47 ± 0.01 | Helped |
| Replace standard attention with translation-equivariant attention | PT-TNP-M32 · Synthetic 1-D regression, shifted test inputs (Δ=1.0) | average log-likelihood: -0.63 ± 0.01 -> -0.46 ± 0.00 | Helped |
| Add translation-equivariant attention | PT-TNP-M128 · T-MNIST | average log-likelihood: 1.14 ± 0.02 -> 1.25 ± 0.03 | Helped |
| Add translation-equivariant attention | PT-TNP-M64 · T-CIFAR10 | average log-likelihood: 1.43 ± 0.01 -> 1.50 ± 0.02 | Helped |
| Add translation-equivariant attention | PT-TNP-M128 · T-CIFAR10 | average log-likelihood: 1.46 ± 0.02 -> 1.51 ± 0.02 | Helped |
| Add translation-equivariant attention | PT-TNP · C. Europe | average log-likelihood: 0.73 ± 0.01 -> 1.35 ± 0.01 | Helped |
| Add translation-equivariant attention | PT-TNP · W. Europe | average log-likelihood: -0.40 ± 0.01 -> 1.27 ± 0.01 | Helped |
| Add translation-equivariant attention | PT-TNP · N. Europe | average log-likelihood: -0.19 ± 0.00 -> 1.44 ± 0.01 | Helped |
| Increase pseudo-token count from M8 to M32 (ablation) | TE-PT-TNP · Synthetic 1-D regression, Δ=0.0 | average log-likelihood: -0.50 ± 0.00 -> -0.46 ± 0.00 | Helped |
| Increase pseudo-token count from M8 to M32 (ablation) | TE-PT-TNP · Synthetic 1-D regression, Δ=1.0 | average log-likelihood: -0.50 ± 0.00 -> -0.46 ± 0.00 | Helped |
| Dynamically update input locations (ablation) | TE-TNP · Simple input distributions, e.g. uniform | performance | No clear effect |
| Dynamically update input locations (ablation) | TE-TNP · Complex input distributions, e.g. bimodal | performance | Helped |
Surprising findings
- Data augmentation was found empirically worse than building in translation equivariance in preliminary experiments, so it was excluded from comparisons.
- The simple RCNP uses mean rather than sum aggregation because summation caused numerical instability on large datasets.
Limitations
- Pairwise computations through MLPs scale computational cost by a large factor in practice.
- Number of pseudo-tokens in PT-TNPs must scale with the number of datapoints.
- Pseudo-token initialisation can fail on pathological datasets, though none were observed in practice.
Experimental setup
Datasets: MNIST CIFAR-10 ERA5 Kolmogorov flow simulations (Kolmogorov flow dataset from prior work) synthetic GP tasks
Benchmarks: Synthetic 1-D regression (shift Delta) T-MNIST T-CIFAR10 Kolmogorov flow ERA5 environmental data (C. Europe, W. Europe, N. Europe)
Metrics: average test log-likelihood training time
Baselines: CNP TNP PT-TNP ConvCNP RCNP GP Multi-task GP
Builds on: Transformer Neural Processes Convolutional Conditional Neural Processes Relational Conditional Neural Processes Latent Bottlenecked Attention Neural Processes Set Transformer
Closest papers by meaning
- Approximately Equivariant Neural Processes (NeurIPS 2024 · similarity 0.89)
- Gridded Transformer Neural Processes for Spatio-Temporal Data (ICML 2025 · similarity 0.89)
- Flow Equivariant Recurrent Neural Networks (NeurIPS 2025 · similarity 0.88)
- Improving Equivariant Networks with Probabilistic Symmetry Breaking (ICLR 2025 · similarity 0.86)
- On Making $SE(2)$-Invariant Networks Optimal (NeurIPS 2026 · similarity 0.86)
- From Bricks to Bridges: Product of Invariances to Enhance Latent Space Communication (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.