Strictly Causal Permutation-Equivariant Attention

• Strictly causal, permutation-equivariant attention architectures are neural sequence models that ensure each token only attends to strictly prior blocks while preserving input order invariance.
• They employ a two-stream design and prefix aggregation to facilitate efficient, parallel inference and robust order-agnostic training in masked diffusion and autoregressive settings.
• Experimental results demonstrate state-of-the-art performance on language modeling tasks by combining strict causal masking with permutation equivariance to enhance training efficiency and scalability.

Strictly causal permutation-equivariant attention architectures constitute a class of neural sequence models that enforce strong causal dependencies—preventing any information leak from non-past tokens—while simultaneously maintaining permutation equivariance with respect to input token orderings, within the constraints defined by block structure. These architectures underpin state-of-the-art masked diffusion models and autoregressive frameworks, enabling efficient, parallel computation of conditional probabilities for generative modeling, and permitting flexible, order-agnostic training and inference strategies (Karami et al., 23 Jan 2026, Zuo et al., 2024).

1. Mathematical Foundations

Strictly causal, permutation-equivariant (SCPE) mappings are defined over sequence-to-sequence transformations $f: (\mathbb{R}^d)^n \to (\mathbb{R}^e)^n$ such that, for any token index permutation $\pi$ on $\{1,\dots,n\}$,

$$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$

This is permutation equivariance. Strict causality, by contrast, ensures that for each output position $n$, the computation depends only on (and can only attend to) information from strictly earlier (block-ordered) tokens.

In masked diffusion settings, blockwise causality is formalized by assigning to each token $j$ a masking time $\tau(j)\in \{1,\ldots,T\}$, then sorting to get a permutation $\pi$ such that $\tau(\pi(1)) \ge \cdots \ge \tau(\pi(N))$, and then partitioning the permuted sequence into $T$ contiguous blocks $\pi$0, where block $\pi$1 contains positions with $\pi$2. The block index function is $\pi$3.

2. Strictly Causal Self-Attention Mechanisms

A strictly-causal self-attention (SCSA) layer is parameterized so that queries for token $\pi$4 depend only on prior block information. For a sequence of $\pi$5 tokens, each layer computes:

$\pi$6

where:

• $\pi$7 denote the key and value projections of the full sequence,
• $\pi$8 is a strictly-causal query matrix (each $\pi$9 depends only on tokens from earlier blocks),
• $\{1,\dots,n\}$0 if $\{1,\dots,n\}$1, $\{1,\dots,n\}$2 otherwise, so that each position $\{1,\dots,n\}$3 can only attend to those $\{1,\dots,n\}$4 which belong to strictly earlier blocks.

By this masking, no information from the same or future block can affect a given position.

3. Blockwise Autoregressive Loss and Parallel Inference

The evidence lower bound (ELBO) for masked diffusion reduces to a blockwise autoregressive loss:

$\{1,\dots,n\}$5

where $\{1,\dots,n\}$6 is a schedule weight. Critically, for each $\{1,\dots,n\}$7 in block $\{1,\dots,n\}$8, the conditional only depends on tokens from earlier blocks (i.e., $\{1,\dots,n\}$9).

Practical implementation leverages strictly-causal and block-causal attention masks to enable parallel evaluation of all $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$0 conditional log probabilities in a single forward pass. Through strided, parallelizable inference, this supports both canonical left-to-right and arbitrary unmasking orderings, as in masked diffusion models (Karami et al., 23 Jan 2026).

4. Permutation Equivariance and Proofs

Self-attention without a mask is permutation-equivariant: permuting inputs permutes outputs accordingly. With block structure, permutation equivariance is preserved if all associated masks and tensors are correspondingly permuted. For any permutation $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$1, for masked diffusion with permutation $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$2 from masking times $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$3, it holds that $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$4, and the mask $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$5 permutes accordingly:

$$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$6

remains equivariant under permutations that preserve block structure. Marginalizing the loss over random block-order permutations further ensures order-agnostic training (Karami et al., 23 Jan 2026).

5. Two-Stream Architectures for Deep Stacks

In deep architectures, information propagation is achieved via a two-stream scheme for the first $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$7 layers:

• Causal stream: $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$8, where $$f(X_{\pi(1)},\dots,X_{\pi(n)}) = (Y_{\pi(1)},\dots,Y_{\pi(n)}), \quad \text{where } (Y_1,\dots,Y_n) = f(X_1,\dots,X_n).$$9 if $n$0, $n$1 otherwise.
• Strictly-causal stream: $n$2, where $n$3 are from the causal stream, and $n$4 is a projection of the previous strictly-causal output.

After $n$5 layers, the remaining $n$6 layers use standard block-causal attention. This structure ensures no same-block leakage early in the network, enabling re-use of efficient causal modules.

6. Prefix Aggregation and Expressive Queries

Strictly-causal query computation is enhanced by prefix-aggregation:

$n$7

where $n$8 is the positional embedding and $n$9 the inverse permutation. This module, in matrix form,

$j$0

is both strictly-causal (by $j$1) and permutation-equivariant (via dot products of permuted embeddings).

7. Algorithmic and Experimental Properties

A parallel forward pass processes all $j$2 positions:

1. Sample masking times and permutation.
2. Compute block indices, build masks, permute inputs.
3. Apply prefix-aggregation for the strictly-causal stream.
4. Alternate two-stream and block-causal layers.
5. Form logits from the final strictly-causal stream, compute cross-entropy for all blocks (Karami et al., 23 Jan 2026).

Experiments confirm that strictly causal, permutation-equivariant architectures achieve state-of-the-art performance on language modeling tasks, with substantially improved training efficiency relative to prior masked diffusion architectures.

Maintaining causality and permutation equivariance requires strict architectural discipline:

• Removal of positional encodings or causal masks leads to loss of order sensitivity and convergence failures (demonstrated via metrics and ablations on digit addition tasks) (Zuo et al., 2024).
• Residual (skip) connections are essential to preserve vertical-slice (position-to-position) consistency across layers; ablation induces significant degradation in performance and disrupts positional identity.

8. Curriculum and Progressive Permutation Training

A progressive-permutation training regime begins with canonical left-to-right sequences, gradually increases permutation complexity, and ultimately presents the model with fully random unmasking orders. Throughout, the same strictly causal, permutation-equivariant layers and masks are used, enabling the architecture to generalize across both autoregressive and non-autoregressive ordering paradigms (Karami et al., 23 Jan 2026).

---

These strictly causal, permutation-equivariant attention architectures unify the efficiency and sequential consistency of autoregressive models with the flexibility and parallelism of masked diffusion approaches, ensuring strong order handling, scalability, and performance on a broad suite of sequence modeling tasks (Karami et al., 23 Jan 2026, Zuo et al., 2024).