Attention Is All You Need (to Avoid Spurious Oscillations)

In this paper, we introduce Exact Flow Linear Attention~(EFLA), an exact-flow formulation of delta-rule linear attention. We show that the delta-rule update can be interpreted as an explicit Euler discretization of an underlying continuous-time system. EFLA replaces this first-order update with the exact closed-form flow. By exploiting the rank-1 structure of the dynamics matrix, both the matrix exponential and the input integral collapse to a simple update that preserves delta-rule linear attention's algebraic structure, parameter count, linear-time complexity, and chunkwise parallelism. This attention mechanism removes the Euler discretization error of the delta-rule dynamics without introducing additional parameters. Experiments on robustness tests, language modeling benchmarks, and the MAD synthetic benchmark show that EFLA improves stability under corrupted and high-energy inputs, reduces perplexity, and achieves stronger downstream performance compared to SSM and Euler-style baselines. These results establish exact-flow integration as a principled and scalable update mechanism for delta-rule linear attention.

A second order finite volume scheme rests on two local quantities: a gradient reconstructed in each cell, and a limiter which scales it down where the reconstruction would overshoot. Both are set by fixed formulas, and on coarse unstructured meshes a small network can supply better values. But a network is free to output anything, and the usual safeguard is a penalty in the training loss, which discourages inadmissible states without preventing them. We replace the penalty by a hard constraint. The network still sets both quantities, and every value it can produce lies inside safe bounds: its stencil weights cannot cancel a neighbour, and its limiter is capped by the local flow. The flux, the wall treatment and the time step are not learned and carry their own guarantees. Admissibility therefore holds for every value of the weights rather than as an outcome of training, and no negative density or pressure occurred in any computation reported here. Because the scheme is safe whatever the network does, we could ask what the network contributes. We test it on supersonic channel flow over an obstacle, including the forward facing step of Woodward and Colella. Learning lowers the error by 38% on an unseen geometry and 29% on an unseen obstacle topology, measured against the same scheme with the network switched off. The method aims at the accuracy of a fine mesh for the cost of a coarse one, and refining once improves the error fourfold while multiplying the run time by eight. Learning secures half of this improvement for a sixth of this time. All of this comes from one of the two quantities the network sets. The gradient reconstruction reproduces the full effect on its own, and the limiter accounts for about a tenth as much. This also explains why the gain fades beyond the Mach numbers the weights were trained on.

FP8 (E4M3) acceleration for attention computation offers significant throughput gains, but the 3-bit mantissa introduces precision challenges when the softmax probability matrix~$P$ is cast to FP8 before the $P \cdot V$ matrix multiplication. We analyze two implementation choices that affect output precision under the \emph{Attention Sink} phenomenon: (1)~the KV block iteration order, and (2) the static scaling factor applied to $P$ before casting. We show that forward KV iteration causes \emph{P-collapse} -- to leading order a fraction $Φ(Δ+ δ_k - 6.93 - \ln S)$ of non-sink $P$ values underflow to zero, where the small shift $δ_k \approx 1$ (for $k_{\text{sink}}{=}4$) is the expected within-sink-block score maximum -- and that reverse iteration removes it, with a zero-underflow guarantee when reverse is combined with $S{=}256$. We further give a constructive characterization of $S = 256 = 2^8$ as the static scale that simultaneously satisfies (i)~bit-exact IEEE 754 scaling, (ii) the lower envelope of a sawtooth function $dp(S)$ over the E4M3 number line ($dp = 2^{-4}$, the minimum worst-case quantization step), and (iii)~the maximum normal-range coverage \emph{among bit-exact ($2^k$) scales} (a non-bit-exact scale such as $448$ attains slightly higher coverage; sec.5}). Both optimizations are already deployed in FlashAttention-3/4 on engineering grounds; our contribution is a quantitative account of \emph{why} these choices are good and a closed-form threshold $Δ_c = 6.93 + \ln S - δ_k$ for predicting kernel-level precision loss. Kernel-faithful experiments ($Q, K, V$ in FP32 to isolate the P-cast effect) show $3$-$10\times$ MSE improvement at moderate sink strengths, and paired tests confirm both fixes saturate to the same precision floor when combined -- which motivated updating the hpc-ops kernel from $S{=}1$ to $S{=}256$.