Fast and Simplex: 2-Simplicial Attention in Triton

TL;DR

Standard attention scores how well a query matches each key using a dot product — a two-way interaction. This paper revives a 2019 idea (the “2-simplicial Transformer”) that scores a triple — query against two different keys at once — using a three-way “trilinear” product. The headline finding is not just that this works, but how it works: 2-simplicial attention changes the exponent of the neural scaling law for reasoning tasks, not merely the offset. In plain terms, almost every architecture tweak in the last five years just shifts the loss curve down a notch; this one changes its slope, meaning the advantage grows as you scale up. The catch: the naive version costs O(n³) in sequence length, so the real contribution is an efficient Triton GPU kernel (using a sliding window + Flash-Attention-style tiling) that makes it tractable. On models from 1B to 3.5B active parameters, the larger models show consistent gains on GSM8k, MMLU, MMLU-pro, and MBPP — and the gain widens with scale.

Problem & Motivation

The pain in one sentence: we are running out of high-quality tokens, and almost nothing we do to the architecture changes the rate at which models improve with size.

Here’s the backstory. Neural scaling laws (Kaplan 2020, Chinchilla/Hoffmann 2022) say training loss falls as a power law in two things: model parameters N and training tokens D. Chinchilla’s famous conclusion was “scale tokens and parameters together” — a 70B model trained on 4× more data beat a 280B model. That advice assumes infinite data. But frontier LLMs already ingest most of the usable internet. The compute-bound era is ending and a data-bound era is beginning.

Now the brutal part. There’s a folk theorem in the scaling-laws literature (see Everett 2025, Kaplan 2020, Shen 2024): essentially every architectural and optimizer improvement merely shifts the error term E (moves the curve down) but does not change the power-law exponent (the slope). A better optimizer, a fancier normalization, a new positional scheme — they make today’s model a bit better but don’t change how fast you improve as you grow. The only lever known to bend the exponent has been data (data pruning, distribution changes — Sorscher 2022, Bahri 2024).

So the question this paper asks: is there an architecture that changes the exponent itself? If one exists, it would be uniquely valuable in a token-constrained world — its edge compounds with scale instead of being a one-time constant. The authors bet on raising the order of the attention interaction from 2-way to 3-way.

What’s New (Core Contribution)

1. Trilinear (2-simplicial) attention, made practical. Before: the 2-simplicial Transformer (Clift 2019) was a theoretical curiosity — interesting for logic toy problems but O(n³) and never scaled. Now: a sliding-window formulation (each query attends to a small w1 × w2 rectangle of key-pairs) plus an efficient Triton kernel brings the cost down to ~dot-product levels at practical context lengths.

2. A rotation-invariant trilinear form so RoPE works. Before: the plain trilinear product ⟨q, k, k'⟩ is not invariant to rotations, so you can’t bolt RoPE (rotary positional embeddings) onto it — positional encoding would corrupt the scores. Now: they introduce a determinant-based trilinear form (signed volume of the parallelepiped spanned by the three vectors) that is rotation-invariant, and prove (Theorem 5.1) it’s expressive enough to solve Match3 (find triples summing to zero mod M) with a single 7-dimensional head.

3. Evidence that it bends the scaling exponent. Before: no architecture was known to change the exponent α. Now: fitting the power law across 1B → 3.5B models, 2-simplicial attention shows a steeper slope α on reasoning/coding benchmarks (e.g. +18.5% on GSM8k, +20.2% on MMLU-pro). That’s the genuinely surprising claim.

4. A real GPU kernel hitting 520 TFLOPS in Triton, rivaling fast Flash-Attention-v3 implementations, with a non-obvious tiling trick (fold one trilinear input into an elementwise multiply on CUDA cores, do the matmul on Tensor cores) and a two-kernel backward pass to dodge expensive atomic operations.

Be precise about novelty: the idea of trilinear attention is from 2019; the scalable kernel, the RoPE-compatible determinant form, and the scaling-exponent evidence are this paper’s real contributions.

How It Works (Technically)

Step 0: recall standard attention (the 2-way baseline)

Given a sequence X ∈ ℝ^{n×d}, you project it three ways to get queries Q, keys K, values V. The score between query i and key j is a dot product:

A_ij = ⟨q_i, k_j⟩ / √d

What it does: measures how aligned query i is with key j — one number per pair. A row-wise softmax turns each query’s scores into probabilities, and the output is the probability-weighted sum of value vectors. That’s it: attention is “for each token, take a weighted average of the other tokens’ values, where the weights come from pairwise similarity.”

Step 1: go from pairs to triples

The 2-simplicial Transformer adds two extra projections: a second key K' = X·W_{K'} and a second value V' = X·W_{V'}. Now the score is a trilinear product over a triple (query i, key j, key k):

A_ijk = ⟨q_i, k_j, k'_k⟩ / √d = (1/√d) · Σ_l Q_il · K_jl · K'_kl (Eq. 5)

What it does and why it matters: a dot product can only ask “is q_i similar to k_j?” The trilinear product asks “do q_i, k_j, and k'_k all align on the same coordinates simultaneously?” — a three-way AND. The term Q_il · K_jl · K'_kl is large only when all three are large on dimension l. This is qualitatively more expressive: Sanford (2023) proved that a 2-simplicial Transformer solves problems (like Match3 — find a triple summing to zero) that a dot-product Transformer needs exponentially many layers to solve. Three-way interactions natively capture compositional/relational structure that two-way ones must approximate with depth.

The softmax now normalizes over both key axes (j, k), and the output mixes the Hadamard product v_j ∘ v'_k (elementwise product of the two value vectors):

ṽ_i = Σ_{j,k} S_ijk · (v_j ∘ v'_k) (Eq. 7)

The whole forward pass is two einsums (Algorithm 1): one to build the [i,j,k] logit tensor, one to contract it back with the values.

Step 2: the determinant trick (making RoPE work)

RoPE encodes position by rotating query and key vectors by an angle proportional to their position; relative position falls out because the dot product is invariant to a shared rotation: ⟨q_i, k_j⟩ = ⟨Rq_i, Rk_j⟩. The trilinear product breaks this — ⟨Rq, Rk, Rk'⟩ ≠ ⟨q, k, k'⟩ — so naive RoPE would scramble the scores.

Their fix: use a function that is rotation-invariant. The signed determinant of three 3-vectors is exactly such a function — geometrically it’s the signed volume of the parallelepiped they span, and volume doesn’t change when you rotate all three vectors together. They chunk each vector into groups of 3 and sum determinants:

A^{det}_{ij₁j₂} = Σ_l det([q_i^{(l)}, k_{j₁}^{(l)}, k'_{j₂}^{(l)}]) (Eq. 9)

By Sarrus’s rule a 3×3 determinant expands into two trilinear dot-product terms (six signed products), so this costs one extra einsum versus Eq. 5. Theorem 5.1 shows that with d = 7 (six dims for two determinant chunks + one “blank pair” selector dimension), a single head computes Match3 exactly — concrete proof the form is expressive. For the kernel derivations they fall back to the simpler Eq. 5 trilinear form “without loss of generality.”

Step 3: make it affordable — sliding window + GQA

Full 2-simplicial attention is O(n³): every query looks at every pair of keys. Unworkable. So each query Q_i only attends to a local w1 × w2 rectangle: w1 candidate K keys and w2 candidate K' keys. Cost drops to O(n · w1 · w2). They sweep window sizes (Table 1) and pick (w1=512, w2=32) — at this setting the FLOP cost is comparable to dot-product attention at 48k context.

The complexity comparison:

• Causal dot-product: 2n² (two matmuls, halved by the causal mask).
• 2-simplicial: 6 · n · w1 · w2 (three multiplies in the trilinear einsum).

A naive sliding window wrecks GPU throughput because each query touches w1 + w2 - 1 distinct KK′ vectors, so the usual Flash-Attention query tiling gives poor occupancy. Borrowing from Native Sparse Attention (Yuan 2025), they use a high Grouped-Query-Attention ratio of 64 (64 query heads share one KV head). That lets them tile along query heads — many heads reuse the same loaded keys — giving dense, mask-free computation.

Architecture & data flow

flowchart LR
  X[Input sequence X] --> P{5 projections}
  P --> Q[Q query]
  P --> K[K key]
  P --> KP[K' key-2]
  P --> V[V value]
  P --> VP[V' value-2]
  Q --> T[Trilinear logits<br/>A_ijk over w1 x w2 window]
  K --> T
  KP --> T
  T --> SM[softmax over j,k]
  SM --> O[Weighted sum of<br/>Hadamard V o V']
  V --> O
  VP --> O
  O --> Y[Layer output]

flowchart TB
  subgraph Model[MoE Transformer stack]
    L1[Layer 1: dot-product attn]
    L2[Layer 2: dot-product attn]
    L3[Layer 3: dot-product attn]
    L4[Layer 4: 2-simplicial attn]
    L5[Layer 5: dot-product attn]
    L1 --> L2 --> L3 --> L4 --> L5
  end
  note[Every 4th layer is 2-simplicial<br/>to balance compute under pipeline parallelism]

Schematic: pairwise vs. trilinear scoring. Move the query coordinate and watch how a dot product (2-way) and a trilinear product (3-way AND) light up differently — the trilinear term only fires when all three vectors agree on the same coordinates.

Schematic: the sliding window. Each query (row) attends to a small w1 × w2 rectangle of key-pairs rather than all O(n²) pairs. Drag the window-size sliders to see the FLOP cost (∝ n·w1·w2) change relative to dense O(n³).

The algorithm, simplified

# Forward pass for 2-simplicial (trilinear) attention, sliding-window version.
# Stubs: project() does the 5 linear maps; softmax2d normalizes over BOTH key axes.
# Shapes: q,k,k2,v,v2 are [seq, d] for one head.

def two_simplicial_attention(x, w1=512, w2=32):
    q, k, k2, v, v2 = project(x)            # 5 projections (2 extra vs vanilla)
    n, d = q.shape
    out = zeros((n, d))

    for i in range(n):                       # each query token
        # local window: w1 candidate K keys, w2 candidate K' keys, causal
        j_idx  = range(max(0, i - w1), i + 1)
        k_idx  = range(max(0, i - w2), i + 1)

        logits = {}
        for j in j_idx:
            for kk in k_idx:
                # THREE-WAY interaction: large only if q_i, k_j, k2_kk
                # all align on the SAME coordinates (a soft AND over dims)
                logits[(j, kk)] = (q[i] * k[j] * k2[kk]).sum() / d**0.5

        s = softmax2d(logits)                # normalize over the (j, kk) grid

        for (j, kk), w in s.items():
            out[i] += w * (v[j] * v2[kk])    # mix the Hadamard product of two values
    return out

The Triton reality (Appendix B) is an online-softmax Flash-Attention loop: it streams over kv1 (the K window) in the outer loop and tiles over kv2 (the K' window) inside, keeping running max/sum for numerical stability, never materializing the full [i,j,k] tensor. The key throughput trick: compute qk1 = q * k1 (elementwise, on CUDA cores) and then tl.dot(qk1, k2) (matmul, on Tensor cores), overlapping the two unit types. The backward pass is split into two kernels (one for dK, dV, one for dK', dV', dQ) because fusing all three gradient orderings into one kernel needs atomic adds whose overhead exceeds the cost of recomputing intermediates.

Built on Prior Work

Results & Evidence

Setup: MoE models, 1B active / 57B total up to 3.5B active / 176B total. Every 4th layer is 2-simplicial (interleaved to balance pipeline-parallel compute). AdamW, peak LR 4e-3, cosine decay. They report negative log-likelihood (lower = better) on GSM8k, MMLU, MMLU-pro, MBPP — chosen because they stress math/reasoning/coding.

Headline (Table 2): the advantage grows with model size.

• At 1B active: roughly a wash, even slightly worse on some metrics — no gain below ~2B.
• At 2B active: 2-simplicial is better by ~1–2% NLL.
• At 3.5B active: better by ~2.3% (GSM8k), ~1% (MMLU), ~2.2% (MMLU-pro).

The real claim (Tables 3–4): fitting −log L(N) = α·log N + β, the 2-simplicial slope α is steeper:

Fits are tight (R² ≥ 0.997). The gain is largest on the hardest, least-saturated benchmarks (GSM8k, MMLU-pro) — consistent with the “this helps reasoning” story.

Kernel: 520 TFLOPS in Triton, competitive with CUTLASS FlashAttention-v3 at large sequence lengths (Figures 3).

What the evidence does NOT establish — read this carefully:

• The scaling-law fit uses three data points (1B, 2B, 3.5B). Fitting a two-parameter line to three points gives high R² almost trivially; the “+18.5% exponent” is suggestive, not airtight. Extrapolating it to 70B+ is a leap.
• It regresses at 1B — there’s a crossover size below which it hurts. The “better exponent” story implicitly relies on extrapolation past that crossover.
• They report NLL on benchmarks, not downstream accuracy, pass@k, or actual generation quality. Lower NLL is a proxy.
• No FLOP-matched or wall-clock-matched comparison of final quality — only “similar sized.” The extra projections and trilinear compute aren’t free in a tokens-per-dollar sense.
• The authors themselves caveat: the Triton kernel is “far away from being used in production.”
• All results are from one lab on MoE models with a specific recipe; no third-party replication.

This is a promising signal with a small-N fit, not a settled result.

How You’d Use It

Honest framing for an AI services operator: you will not be writing custom trilinear attention kernels for clients in 2025. This is foundation-model R&D, not an off-the-shelf capability. But it’s strategically useful in three concrete ways:

1. Buy-side intelligence. If a base-model vendor ships “2-simplicial” or “trilinear” or “higher-order attention,” you now know what it means: better token efficiency on reasoning, advantage that grows with size, extra compute cost per layer. When a client asks “should we fine-tune Model A or Model B for our math-heavy workflow?”, architecture-level token efficiency is a real differentiator you can speak to.

2. The thesis transfers even if the kernel doesn’t. The durable lesson is: higher-order interactions buy reasoning per parameter. In your own MAS work, the analog of “going trilinear” is letting an agent condition on triples of context (query + two retrieved items jointly) rather than scoring retrieved items independently. A reranker that scores (query, doc_i, doc_j) triples can catch compositional relevance that pairwise scoring misses — same intuition, prompt/retrieval layer instead of CUDA.

3. Reasoning-task positioning. The data says the gains concentrate on GSM8k/MMLU-pro — math, multi-step logic, code. If your offering targets those (financial reasoning, code-gen agents, technical Q&A), prefer base models whose architecture is tuned for reasoning token-efficiency over raw size, and benchmark on hard sets where the exponent advantage shows up.

Where it slots in: this is a pre-training architecture decision — upstream of everything you touch. Your leverage is model selection and the transfer of the idea to your retrieval/reranking/agent-scoring layers.

Build Your Own (Minimal Recipe)

You can’t cheaply reproduce the 520-TFLOPS kernel, but you can build a working trilinear-attention layer to internalize the mechanism and prototype on small models.

Smallest version that captures ~80% of the value:

1. A reference trilinear attention layer in PyTorch. Add W_K' and W_V' projections. Compute the [batch, heads, i, j, k] logit tensor with torch.einsum("bhid,bhjd,bhkd->bhijk", q, k, k2), softmax over the last two axes, contract with einsum("bhijk,bhjd,bhkd->bhid", attn, v, v2). This is Algorithm 1 verbatim — correct, slow, fine for n ≤ 256.
2. Add the sliding window. Restrict j to the last w1 and k to the last w2 positions per query (start w1=64, w2=8 at toy scale). This is what makes memory not explode and is the single most important practical step.
3. Drop it into a small Transformer, interleaving — say every 4th layer is 2-simplicial, rest are vanilla — and train on a reasoning-flavored dataset (GSM8k-style synthetic, or Match3 itself as a unit test that proves the layer learns what theory says).
4. (Optional, hard) the Triton kernel for speed.

The 1–2 genuinely hard parts:

• The Triton kernel and its backward pass. The forward online-softmax loop is intricate; the backward pass needs the two-kernel split to avoid atomics. Budget weeks, not days. For a prototype, skip it — PyTorch autograd handles the reference layer.
• The GQA-64 tiling that gives real throughput. Without it the windowed kernel is memory-bound. This is a systems problem, not a modeling one.

Reach for: PyTorch + torch.einsum for the reference; Triton (with FlashAttention’s tutorial kernel as a template) only if you go fast; Sanford’s Match3 as a correctness oracle; a small MoE or dense model (~100M–1B) so you can actually see the crossover the paper reports.

How to Improve It

1. Settle the scaling claim with more points. The exponent rests on three model sizes. Train 4–6 sizes (and ideally a 7–13B) to see whether the steeper slope survives — and whether the 1B regression is a true crossover or noise. This is the single highest-value follow-up.
2. FLOP-matched comparison. The fair question isn’t “same param count” but “same training compute / same dollars.” Re-run with FLOP-matched baselines; if the trilinear advantage holds per-FLOP, that’s a much stronger result.
3. Learn or adapt the window. (512, 32) is a fixed sweep choice. Make w1, w2 content-dependent (route each query to a window size) or let the model attend to a sparse, learned set of key-pairs (NSA-style selection) instead of a contiguous rectangle — the relevant triples may not be local.
4. Use the determinant form everywhere, not just the proof. They train with the simpler Eq. 5 form and reserve the rotation-invariant determinant form for the expressivity theorem. Train end-to-end with the determinant form + RoPE and measure whether positional generalization (long-context) improves — that’s the form’s whole point.
5. Port to CUTLASS / a production kernel. The authors flag the Triton kernel as prototype-grade. A hand-tuned CUTLASS or a Tensor-core-native fused backward (no atomics, no recompute) would close the gap to dot-product cost and make the architecture deployable — turning a research signal into a usable lever.
6. Test on long-context reasoning. Trilinear interactions should shine on tasks needing 3-way relational binding across distance (multi-hop QA, theorem proving, code with long-range dependencies). The current benchmarks are mostly short-context; the architecture’s theoretical edge (Match3-style composition) is barely stressed.

Glossary

• Attention — mechanism where each token produces an output by taking a weighted average of other tokens’ value vectors; weights come from how well its query matches their keys.
• Dot-product / bilinear attention — standard attention: the score is a 2-way (pairwise) dot product ⟨q, k⟩.
• Trilinear / 2-simplicial attention — score is a 3-way product ⟨q, k, k'⟩ over a triple; a soft AND across coordinates, more expressive than pairwise.
• Simplex — a generalization of a triangle: a 1-simplex is an edge (2 points), a 2-simplex is a filled triangle (3 points); hence “2-simplicial” = three-way interaction.
• Scaling law — empirical power-law relating training loss to parameters N and tokens D: L = E + A/N^α + B/D^β.
• Exponent vs. offset — the exponent (α) is the slope of the log-log loss curve (how fast you improve with scale); the offset (E) just shifts it down. Bending the exponent is rare and valuable.
• Token efficiency — getting more model quality per training token; the goal when data, not compute, is scarce.
• RoPE (Rotary Position Embedding) — encodes token position by rotating query/key vectors so the dot product depends on relative distance; relies on dot-product rotation invariance.
• Rotation invariance — a function whose value is unchanged when all inputs are rotated together; the dot product has it, the plain trilinear product doesn’t, the determinant does.
• Determinant trilinear form — using the signed volume (3×3 determinant) of three vectors as the score, because volume is rotation-invariant — making RoPE compatible.
• Sarrus’s rule — the standard formula expanding a 3×3 determinant into six signed product terms.
• Hadamard product (∘) — elementwise multiplication of two equal-length vectors.
• Match3 — toy task: find triples of inputs that sum to zero mod M; provably hard for dot-product attention, easy for trilinear; used as an expressivity oracle.
• Sliding window attention — restricting each query to a local neighborhood of keys to cut cost; here a 2D w1 × w2 window over key-pairs.
• GQA (Grouped Query Attention) — multiple query heads share one key/value head, saving memory; ratio 64 means 64 query heads per KV head, enabling head-wise tiling.
• Flash Attention — IO-aware attention kernel that computes exact attention without materializing the full score matrix, using online (streaming) softmax.
• Online softmax — computes softmax in a single streaming pass by tracking a running max and sum, for numerical stability without storing all logits.
• Triton — a Python-like language for writing GPU kernels, easier than CUDA, used here to implement the custom attention.
• CUDA core vs. Tensor core — two compute units on NVIDIA GPUs; CUDA cores do general elementwise math, Tensor cores do fast matrix multiplies; the kernel overlaps both.
• TFLOPS — trillions of floating-point operations per second; a throughput measure (520 here ≈ competitive with the fastest attention kernels).
• Atomic operation — a thread-safe read-modify-write to shared memory; correct but slow under contention, which is why the backward pass is split to avoid it.
• MoE (Mixture of Experts) — model where each token is routed to a few “expert” sub-networks; “active params” = those used per token, “total params” = all experts combined.
• Negative log-likelihood (NLL) — the loss = −log(probability the model assigns to the correct answer); lower is better; a proxy for accuracy.
• Pipeline parallelism — splitting model layers across GPUs in stages; motivates interleaving the heavy 2-simplicial layers evenly to balance per-stage compute.