Causal Optimal-Transport Attention: Positional Priors and Length Extrapolation

원문 초록

Causal masking obstructs the column balancing used by optimal-transport (OT) attention. We study an unbalanced OT operator that uses only current and earlier rows and characterize its induced positional prior. For bounded scores and any fixed balancing strength below one, its fixed point assigns vanishing mass to every fixed window as context grows. With constant scores, the prior has an exact scale-invariant form. We classify positional encodings by whether they prevent this dilution. We also derive finite-length window-mass bounds for log-distance priors with exponents near one. Preregistered experiments use 25M- and 100M-parameter language models. The OT operator improves loss extrapolation over softmax, while exploratory comparisons favor an explicit distance bias at a fraction of the cost. Changing only the bias beyond the training length moves long-context loss in the predicted direction on all 30 tested checkpoints. During training, learned exponents approach one even when constrained above it. A margin chosen from prior calculations controls distant-mass growth in the tested setting.