Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection

About

Doubly-stochastic attention has emerged as a transport-based alternative to row-softmax attention, with recent Transformer variants using it to reduce attention sinks and rank collapse while improving performance. In this family, the standard approach is Sinkhorn scaling, which trains more efficiently but still repeats matrix scaling in every inference forward pass. Sliced-transport attention removes the online iteration, but its soft sorting approximation materializes dense tensors for each slice, requiring substantially more training resources than Sinkhorn attention. We introduce ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection, a train-then-compile method that trains the doubly-stochastic layer with Sinkhorn, then replaces the iterative scaling loop at inference with a fixed sliced-dual operator. It learns a lightweight parametric map from exact one-dimensional Kantorovich potentials to the Sinkhorn query-side dual, then reconstructs the attention plan with a two-sided entropic c-transform. Across language and vision benchmarks, ASAP keeps the cheaper training setup and remains highly competitive with recent baselines. In the main frozen-layer benchmark, ASAP is 5.3 faster than the trained Sinkhorn teacher while matching its accuracy; in downstream replacements, ASAP recovers most of the teacher performance without any retraining.

Huy Tran, Max Milkert, David Hyde• 2026

Related benchmarks

TaskDatasetResultRank
Image ClassificationCIFAR-100 (test)
Accuracy63.64
295
Text ClassificationDBPedia-14
Accuracy (Acc)98.26
22
Text ClassificationYelp Review Full (YRF)
Accuracy54.51
16
Image ClassificationCats and Dogs full-data
Accuracy (%)80.44
10
Text ClassificationIMDB
Accuracy87.36
8
Text ClassificationAG-News
Accuracy90.04
8
Showing 6 of 6 rows

Other info

Follow for update