Causal Optimal-Transport Attention: Positional Priors and Length Extrapolation

Ming Lian

cs.LGNov 1, 2026

Causal Optimal-Transport Attention: Positional Priors and Length Extrapolation

Authors: Ming Lian

Organizations: Independent Researcher

Abstract

Causal masking obstructs the column balancing used by optimal-transport (OT) attention. We study a prefix-consistent unbalanced operator and characterize the positional prior it induces. Under bounded scores, for any fixed balancing strength $ρ<1$ , its fixed point dilutes attention: every fixed window receives a vanishing share as the context grows. For constant scores, we derive an exact scale-invariant prior. We then classify positional encodings by whether their distant attention mass remains bounded away from one. Discrete window-mass bounds quantify the finite-length dilution of summable log-distance priors near exponent one. Preregistered experiments with 25M- and 100M-parameter language models establish improved loss extrapolation over softmax; exploratory comparisons favor an explicit distance bias at a fraction of the cost. CABLE outperforms RoPE with evaluation-time YaRN on every seed at three training lengths. Changing only the bias beyond the training length moves long-context loss in the predicted direction on all 30 tested checkpoints. Even with exponents constrained above one, many heads approach the summability boundary during training. A margin selected by prior-only calculations for a specified window, length range and growth tolerance controls distant-mass growth in the tested setting. These results connect loss extrapolation to positional decay over the lengths actually evaluated.

Figures & tables

Figure 1: The implicit prior with constant scores. Left: the last row of the fixed point at

t=2048

(circles, every 100th key) against the limit shape

f(s)

of Theorem 7 (lines). Right: the same row near the diagonal for

\rho=0.5

, multiplied by

\Gamma(1-\rho)\,d^{\rho}

so that the power law of Theorem 8 (2) is the line

1

. The fixed point (circles) follows the discrete profile of Theorem 8 (1) at small

d

, the power law for

1\ll d\ll t

, and the limit shape

f

, here

(d/t)^{\rho}(\log\frac{t}{t-d})^{-\rho}

, as

d

approaches

t

Figure 2: Attention mass beyond a window

L

at row

t

for the prior

(1+d)^{-c}

with constant content (Corollary 16 ), against the context length; the shaded range is the one used in our experiments;

c=1.001

(dashed) nearly overlaps

c=1

. The linear bias

-m_{h}d

(dotted) uses ALiBi’s slopes and keeps a bounded far mass (Lemma 13 ); its far mass is averaged over eight heads.

Arm	val@1024	val@2048	s/step
softmax	$3.0439\pm 0.0023$	$3.5970\pm 0.0285$	0.42
log- $n$ softmax	$3.0459\pm 0.0032$	$3.4456\pm 0.0596$	0.45
OT, $\rho=0.25$	$3.0388\pm 0.0038$	$3.3379\pm 0.0471$	0.98
OT, $\rho=0.5$	$3.0344\pm 0.0010$	$3.1423\pm 0.0544$	1.54
softmax + learnable $-c_{h}\log(1+d)$	$\mathbf{3.0254\pm 0.0023}$	$\mathbf{3.0147\pm 0.0075}$	0.62

Table 1: OT attention at 100M: validation loss in nats per token (mean

\pm

sd over three seeds) at the training length (1024) and at twice the training length. The last row is an exploratory follow-up with the same configuration.

Figure 3: Exploratory: far mass (top) and validation loss (bottom) against the evaluation length at 25M, means over two seeds, one column per case of Proposition 14 ; each row shares its vertical axis. Values for every arm, including the small differences within the bounded case, are in Table 7 .

Training length (seeds)	gap in $L_{\mathrm{ext}}$	gap at $1\times$	gap beyond $L$
1024 (5)	$0.059$ ( $0.055$ – $0.063$ )	$0.019$	$0.040$
2048 (3)	$0.065$ ( $0.062$ – $0.071$ )	$0.030$	$0.035$
4096 (3)	$0.074$ ( $0.070$ – $0.078$ )	$0.045$	$0.029$

Table 2: Retest at 100M on new seeds: B0 + YaRN-t minus CABLE in

L_{\mathrm{ext}}

(range over seeds), B0 minus CABLE at the training length (

1\times

), and their difference, the gap beyond

L

measured from each arm’s own loss at

L

. Absolute losses are in Table 9 .

Figure 4: Test-time interventions beyond the training length: loss change at the longest length (

16\times

at 25M,

8\times

at 100M) for each of the 30 checkpoints, with the mean of each group (bar). Colours are the case of the arms each intervention is applied to; the panels have separate linear axes (values in Table 10 , loss curves of three checkpoints in Figure 5 ).

Arm	$1\times$	$32\times$	far-mass growth
P1, $\delta_{0}=0$	3.444	3.427	0.032
P1, $\delta_{0}=0.25$	3.445	3.395	0.015
P1, $\delta_{0}=0.5$	3.448	3.385	0.004
ALiBi (P4)	3.441	3.355	$-0.006$

Table 3: Training-time margin at 25M,

L=512

: validation loss and far-mass growth from

4L

16L

(preregistered tolerance

0.02

), means over three seeds. Far mass and the share of heads at the bound are in Appendix E.5 .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Result	Limit	Finite length
Thm. 6	$t\lambda_{t}\to{a^{*}}$	decreases towards ${a^{*}}$ ; about $10^{-3}$ above it at $t=2048$
Thm. 7	$tP_{t,\lfloor st\rfloor}\to f(s)$	within $2\%$ at $n=1024$ , error halves from $n=256$ to $1024$
Thm. 8	discrete profile, then $d^{-\rho}$	within $1\%$ for $d\leq 8$ at $n=1024$
Remark 22	window mass $\approx(L/t)^{1-\rho}/\Gamma(2-\rho)$	$L=64$ : doubling $n$ from $512$ to $1024$ multiplies it by $2^{-(1-\rho)}$ to $2\%$ ; within $5\%$ of the limit
Thm. 9	$tP^{(k)}_{t,\lfloor st\rfloor}\to f(s)$ , $k\geq 2$	$k=2$ : ratio to $f(s)$ constant in $s$ to $1\%$ at $n=1024$ (the amplitude: Appendix D.3 )
Prop. 10	$tP\kappa/(Kf)\to 1$	binned mean within $2\%$ of one

Appendix

Table 4: Limit statements against finite lengths. Tests in test_unbalanced_scaling , except Theorem 6 ( test_unbalanced_recursion ) and the last row ( test_potential_framework ).

	Screening (25M)	Confirmation (100M)
Width / depth / heads	512 / 8 / 8	768 / 14 / 12
MLP ratio, vocabulary	4, 50,304 (tied)	4, 50,304 (tied)
Non-embedding parameters	25.2M	99.2–99.5M
Training length	512	1024 (also 2048, 4096)
Micro-batch $\times$ accumulation	$32\times 2$	$8\times 4$ ( $4\times 4$ , $2\times 4$ )
Tokens / steps	0.5B / 15,258	2B / 61,035

Appendix

Table 5: Training configurations; 32,768 tokens per step at both scales. Arms differ only in their position encoding and normalizer; their settings are the overrides listed in the accompanying configuration files ( configs/exp0002/arms.yaml ).

Arm	Position term	Bounded far mass
B0	RoPE, no bias	no (dilutes)
P1	RoPE $+$ learnable $-c_{h}\log(1+d)$ , one exponent per head	iff $c_{h}>1$
P2	as P1 with fixed exponents (log-spaced $0.1$ – $1$ )	no (dilutes)
P3	RoPE $+$ KERPLE log kernel $-p_{h}\log(1+a_{h}d)$ [ 4 ]	iff $p_{h}>1$
P4 (ALiBi)	RoPE $+$ $-m_{h}d$ [ 32 ]	yes
P5	P1 without RoPE	iff $c_{h}>1$

Appendix

Table 6: Positional schemes, with

d=t-j

. “Bounded far mass” is the guarantee of Proposition 14 under bounded inputs, where the conditional cases are stated in full. As in their papers, GRAPE-A, GRAPE-AP, FoX and CABLE are used without RoPE; so is P5. Every other arm with a bias adds it to RoPE.

Case	Arm	$1\times$	$L_{\mathrm{ext}}$	far mass $4\times$	$16\times$	$c>1$
dilutes	B0	3.474	3.865	0.679	0.889	–
	GRAPE-M	3.471	3.930	0.701	0.899	–
	O1	3.469	3.573	0.428	0.646	–
	P2	3.461	3.460	0.424	0.627	0.00
learned	P5	3.489	3.498	0.253	0.480	0.34
	P6	3.460	3.440	0.364	0.561	0.18

Appendix

Table 7: Screening at 25M,

L=512

, means over two seeds, grouped by the cases of Proposition 14 . Loss in nats per token over 2M predictions; far mass from the probe of Section 5 ; last column, the share of learned exponents above one.

Arm	$1\times$	$L_{\mathrm{ext}}$	$8\times$	$L_{\mathrm{ext}}$ , YaRN-t	ms/step
B0 (RoPE)	$3.0440\pm 0.0019$	$4.1136\pm 0.0584$	$5.468\pm 0.106$	$3.0397\pm 0.0016$	257
log- $n$ softmax	$3.0464\pm 0.0038$	$3.9505\pm 0.0827$	$5.369\pm 0.339$	$3.0541\pm 0.0054$	347
GRAPE-M	$3.0433\pm 0.0014$	$4.1982\pm 0.1076$	$5.576\pm 0.183$	–	262
GRAPE-AP	$3.0418\pm 0.0008$	$3.0105\pm 0.0013$	$3.004\pm 0.001$	–	482
ALiBi (P4)	$3.0312\pm 0.0007$	$2.9932\pm 0.0006$	$2.979\pm 0.001$	$3.0585\pm 0.0019$	260
O2 (P4 + OT)	$3.0351\pm 0.0011$	$2.9972\pm 0.0011$	$2.983\pm 0.001$	$3.0458\pm 0.0025$	1582

Appendix

Table 8: Confirmation at 100M: validation loss in nats per token (mean

\pm

sd over three seeds). YaRN-t is defined only for arms with RoPE. CABLE uses the learning-rate variant selected by the preregistered rule.

	$1\times$		$L_{\mathrm{ext}}$
Length (seeds)	B0	CABLE	B0 + YaRN-t	CABLE
1024 (5)	$3.0445\pm 0.0039$	$3.0258\pm 0.0012$	$3.0391\pm 0.0039$	$2.9801\pm 0.0012$
2048 (3)	$3.0233\pm 0.0028$	$2.9932\pm 0.0012$	$3.0302\pm 0.0039$	$2.9649\pm 0.0012$
4096 (3)	$3.0201\pm 0.0045$	$2.9749\pm 0.0009$	$3.0343\pm 0.0038$	$2.9602\pm 0.0009$

Appendix

Table 9: Retest at 100M on new seeds: validation loss in nats per token (mean

\pm

sd) at the training length (

1\times

) and

L_{\mathrm{ext}}

Figure 5: Loss against the evaluation length for one 25M checkpoint per intervention, before (dark grey) and after; the curves coincide up to the training length by construction.

	Checkpoints (seeds each)	Scale	Loss change	Far mass
				before	after
cap	ALiBi, GRAPE-A, FoX, CABLE (2)	25M	$+0.060$ ( $+0.008$ to $+0.161$ )	0.007–0.058	0.036–0.201
cap	ALiBi, CABLE, CABLE (default lr) (3)	100M	$+0.020$ ( $+0.012$ to $+0.064$ )	0.003–0.045	0.036–0.093
repair	P1, P2, P6 (2)	25M	$-0.264$ ( $-0.312$ to $-0.198$ )	0.55–0.63	0.32–0.39
repair	LC (2)	100M	$-0.009$ ( $-0.011$ to $-0.008$ )	0.36–0.37	0.26–0.27
repair	log bias, exploratory (3)	100M	$-0.182$ ( $-0.210$ to $-0.166$ )	0.58–0.61	0.39–0.44

Appendix

Table 10: Test-time interventions beyond the training length (Figure 4 ). Loss change in nats at

16\times

(25M) or

8\times

(100M), mean and range over checkpoints; far mass at the same length before and after. The

1\times

losses agree within

10^{-4}

in every case.

Exp.	Test	Observed	Outcome
0001	best $\rho$ below softmax at $L$ by $0.0100$ , every seed	$0.0095$	fails
0001	best $\rho$ below softmax / log- $n$ at $2L$ by $0.057$ , every seed	$0.455$ , $0.303$	holds
0002	CABLE below B0 + YaRN-t in $L_{\mathrm{ext}}$ by $\delta_{e}=0.117$ , every seed	$0.058$	fails
0002	GRAPE or O2 beats CABLE (switching rules)	no	not triggered
0003	CABLE below B0 + YaRN-t in $L_{\mathrm{ext}}$ by $\delta_{A}\geq 0.01$ , every seed	$0.059$ , $0.065$	holds
0003	gt1 below lt1 at $32\times$ by $0.01$ , every seed	$0.488$ – $0.678$	holds

Appendix

Table 11: Preregistered tests and their outcomes. Loss thresholds apply to mean paired gains; “every seed” and “every checkpoint” require consistent signs. Far-mass thresholds apply to seed means; all thresholds were fixed before the runs.

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 197 papers

Jul 21, 2026

Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention

Transformers with relative positional encodings often extrapolate to sequences longer than those seen during training, whereas transformers with learned absolute encodings typically do not. This is a robust empirical regularity, and the explanations offered for it so far are chiefly about expressivity, that is, about whether a length-generalizing solution exists. We give an optimization explanation. On a minimal fixed-offset retrieval task that isolates positional selection, the gap is governed by the implicit bias of the trained attention head: among the many solutions that fit short sequences, which one gradient descent actually selects. We prove that rotary encodings make the attention logit a function of relative offset alone, an exact equivariance, so whatever selection rule is learned at training lengths is reproduced verbatim at every longer length. Learned absolute encodings instead leave out-of-range positions unconstrained, and the trained head pins to a fixed absolute position inside the training range. We characterize the learned rotary rule as a low-rank ``carrier'' kernel aligned with the target offset, and we derive the resulting graceful accuracy decay as an attention-dilution law; both predictions are confirmed across seeds and offsets. A linear-attention control shows the mechanism is specific to softmax: without normalization, training selects a min-norm interpolant that does not extrapolate. The phenomenon, the equivariance, and the carrier all transfer to a multi-layer, multi-head transformer trained on a full-sequence length-generalization task. The account connects the implicit bias of attention, implicit bias for extrapolation in recurrent models, and the learning side of the RASP-L conjecture.

Subham Singh, Ashutosh Mishra, Subha Raut

Department of Mathematics and Statistics, Mississippi State University · Independent Researcher

Softmax Attention Rotary Positional Embeddings Length Generalization Implicit Bias Relative Positional Encoding

Ask PDF

Jun 9, 2026

Express Language Modeling

We introduce a new tool, Express, for converting a non-causal attention approximation into a causal approximation with matching approximation guarantees. When combined with the state-of-the-art Thinformer approximation, Express improves upon the best known causal attention guarantees, delivering

\log^{3/2}(n)/s

approximation error with only

O(s)

memory and

O(s^2 \log^2(n))

compression overhead for a sequence of length

n

. We pair these developments with an efficient I/O-aware Triton implementation, demonstrate substantial speedups over FlashAttention 2, and use Express to overcome four resource bottlenecks in the language modeling pipeline: long-context prefill, KV cache compression, long-form memory-constrained decoding, and long-form compute-constrained decoding.

Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi +1

Cornell Tech · University of Cambridge · Microsoft Research

Long-Context Language Model Inference Attention Mechanisms Efficient Language Model Inference Attention Approximation Memory-Efficient Attention Causal Attention

Ask PDF

Aug 9, 2026

Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling

Self-attention models content-dependent interactions between tokens but does not by itself encode token order. Position encoding addresses this limitation by introducing absolute coordinates, relative distances, or position-dependent rotations into Transformer representations and attention scores. This technical survey develops a unified account of sinusoidal and learned absolute position embeddings, Shaw-style relative position representations, Transformer-XL, T5 relative position bias, ALiBi, and Rotary Position Embeddings (RoPE). We derive how RoPE converts absolute position indices into relative phase differences in Query-Key inner products and compare these methods in terms of where position is injected, computational cost, compatibility with KV caching, and length extrapolation. We then examine long-context extensions, including Position Interpolation, RoPE scaling laws, NTK-aware scaling, Dynamic NTK, NTK-by-parts, YaRN, LongRoPE, and LongRoPE2, with emphasis on frequency allocation, attention rescaling, training length, and target context length. We also summarize implementation considerations, evaluation protocols, and position-encoding choices in representative large language models. A central conclusion is that the ability to compute positional features beyond the training length does not imply reliable long-context generalization; context extension must be evaluated through short-context retention, position-wise perplexity, retrieval, reasoning, and long-context code tasks.

Jiguo Li

Rotary Positional Embeddings Long-Context Modeling Positional Encoding LLM Context Management Relative Positional Encoding

Ask PDF

Abstract

Long-Context Language Modeling

Explore similar work

Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention

Express Language Modeling

Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling