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 $[\mathcal X^{(1)},\ldots,\mathcal X^{(T)}]$, where block $t$ contains positions with $\tau(j)=T-t+1$. The block index function is $\mathcal{B}(n)= T-\tau(n)+1$.

2. Strictly Causal Self-Attention Mechanisms

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

$$\mathrm{SA}^{sc}(\hat Q, K, V; M^{sc}) = \mathrm{Softmax}(\hat Q K^\top + M^{sc}) V$$

where:

• $K,V \in \mathbb{R}^{N\times d}$ denote the key and value projections of the full sequence,
• $\hat Q \in \mathbb{R}^{N\times d}$ is a strictly-causal query matrix (each $\hat q_n$ depends only on tokens from earlier blocks),
• $M^{sc}_{n,i} = 0$ if $\mathcal{B}(i) < \mathcal{B}(n)$, $-\infty$ otherwise, so that each position $n$ can only attend to those $i$ 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:

$$\mathcal{L}_{\mathrm{diff}} = \mathbb{E}_{q} \biggl[ \sum_{t=1}^{T-1}\gamma(t) \sum_{n\in \mathcal{B}(t)} -\log p_\theta(x_n|\{x_i: \mathcal{B}(i)\leq T-t\}) \biggr]$$

where $\gamma(t)$ is a schedule weight. Critically, for each $n$ in block $t$, the conditional only depends on tokens from earlier blocks (i.e., $\mathcal{B}(i)<\mathcal{B}(n)$).

Practical implementation leverages strictly-causal and block-causal attention masks to enable parallel evaluation of all $N$ 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 $\sigma$, for masked diffusion with permutation $\pi$ from masking times $\tau$, it holds that $$\mathcal{B}_{\sigma\circ\pi}(n)=\mathcal{B}_\pi(\sigma^{-1}n)$$, and the mask $M^{sc}$ permutes accordingly:

$\mathrm{Softmax}(\hat Q K^\top + M^{sc})V$

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 $L^{2s}$ layers:

• Causal stream: $$h^l = \mathrm{SA}^c(h^{l-1}, h^{l-1}, h^{l-1}; M^c)$$, where $M^c_{n,i}=0$ if $\mathcal{B}(i) \leq \mathcal{B}(n)$, $-\infty$ otherwise.
• Strictly-causal stream: $$\tilde h^l = \mathrm{SA}^{sc}(\tilde Q, K^l, V^l; M^{sc})$$, where $K^l,V^l$ are from the causal stream, and $\tilde Q$ is a projection of the previous strictly-causal output.

After $L^{2s}$ layers, the remaining $L-L^{2s}$ 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:

$$\hat q_n = \sum_{i:\mathcal{B}(i)<\mathcal{B}(n)} H_{n,i} x_i, \quad H_{n,i}=\langle \mathrm{PE}_{\hat\pi(n)}, \mathrm{PE}_{\hat\pi(i)} \rangle$$

where $\mathrm{PE}$ is the positional embedding and $\hat\pi$ the inverse permutation. This module, in matrix form,

$$X^0 = \mathrm{PA}^{sc}(X) = (\mathrm{PE}\,\mathrm{PE}^\top \odot M^{sc}) X$$

is both strictly-causal (by $M^{sc}$) and permutation-equivariant (via dot products of permuted embeddings).

7. Algorithmic and Experimental Properties

A parallel forward pass processes all $N$ 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).