Taylor Multi-Head Self-Attention

• Taylor Multi-Head Self-Attention (T-MSA) is a reformulation of standard attention that uses a truncated Taylor series to approximate softmax for efficient global modeling.
• It factorizes computations into global sums and per-query inner products, reducing the quadratic complexity to linear time for handling long time-frequency sequences.
• When integrated with Spatial-Channel Enhancement Attention (SCEA) in the LORT architecture, T-MSA achieves competitive speech enhancement performance with minimal parameters.

Searching arXiv for the cited papers to ground the article in the current records. Taylor Multi-Head Self-Attention (T-MSA) is a Taylor-based reformulation of multi-head self-attention designed to reduce the quadratic cost of softmax attention on long sequences by replacing the exponential kernel with a truncated Taylor series and algebraically factorizing the resulting computation into global sums and per-query inner products. In the speech-enhancement architecture LORT, T-MSA is instantiated with a first-order truncation, integrated with Spatial-Channel Enhancement Attention (SCEA), and deployed in a multi-resolution U-Net-like encoder-decoder to provide efficient global modeling for time-frequency representations while maintaining a low parameter count and computational complexity (Wang et al., 28 Sep 2025).

1. Motivation and problem setting

In speech enhancement, spectrogram sequences are long, and standard softmax attention incurs quadratic complexity in sequence length. Within U-Net-like multi-resolution processing, global modeling across many time-frequency tokens becomes expensive. T-MSA addresses this by replacing the exponential kernel implicit in softmax with a low-order Taylor polynomial, yielding a linear-time attention computation while preserving global context (Wang et al., 28 Sep 2025).

The standard difficulty is the need to construct and normalize dense affinity matrices of size $N \times N$. This is both time- and memory-intensive, and it can be numerically sensitive when logits have large magnitude. T-MSA avoids the explicit softmax normalization over all query-key pairs by using a polynomial approximation that factors computations into global sums and low-rank aggregations. In the formulation reported for LORT, the resulting complexity drops from $O(N^2)$ to approximately $O(N)$ with respect to the number of tokens $N$.

For input of $t \times f$ patches and hidden dimension $D$, the reported costs are

$O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$

$O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$

Since $D \ll t f$, the Taylor-based form is more efficient as $t f$ grows. In LORT, this efficiency is part of a lightweight design that achieves only $O(N^2)$0 parameters while retaining multi-resolution global modeling (Wang et al., 28 Sep 2025).

2. Mathematical formulation

For reference, standard scaled dot-product attention for head $O(N^2)$1 is defined with

$O(N^2)$2

$O(N^2)$3

and

$O(N^2)$4

T-MSA replaces the exponential in softmax by a truncated Taylor series. For a scalar $O(N^2)$5,

$O(N^2)$6

In LORT, a first-order truncation is adopted:

$O(N^2)$7

For query $O(N^2)$8, key $O(N^2)$9, and value $O(N)$0, standard attention applies

$O(N)$1

With the first-order Taylor expansion,

$O(N)$2

The numerator and denominator then factorize as

$O(N)$3

$O(N)$4

where

$O(N)$5

$O(N)$6

$O(N)$7

Thus,

$O(N)$8

This factorization avoids the explicit $O(N)$9 matrix and enables linear-time aggregation. LORT retains the $N$0 scaling, and a stabilized form is also given as

$N$1

The same construction is applied independently across heads after linear projections

$N$2

followed by splitting into $N$3 heads and final recombination through

$N$4

A feature-mapping interpretation is also provided: the first-order approximation corresponds to polynomial features $N$5, so that softmax attention is approximated by a normalized bilinear form in $N$6 and $N$7 (Wang et al., 28 Sep 2025).

3. Integration with SCEA and multi-resolution LORT

In LORT, T-MSA is not used in isolation. It is enhanced by Spatial-Channel Enhancement Attention, a complementary branch introduced to facilitate inter-channel information exchange and alleviate the spatial attention limitations inherent in Taylor-based Transformers. The design explicitly assigns different roles: T-MSA emphasizes global attention primarily within each channel, while SCEA re-introduces inter-channel information exchange and spatial saliency (Wang et al., 28 Sep 2025).

SCEA has two branches. The channel-attention branch performs 2-D average pooling over time-frequency to obtain channel descriptors $N$8 and applies a 1-D convolution across channels with kernel size $N$9 and a sigmoid:

$t \times f$0

The spatial-attention branch pools over channels to obtain $t \times f$1 and applies a $t \times f$2 convolution with sigmoid:

$t \times f$3

These weights modulate the T-MSA output multiplicatively, and the SCEA-modulated output is fused with the T-MSA branch. The stated purpose is to alleviate limited spatial sensitivity and inject inter-channel communication missing in T-MSA.

At the architectural level, LORT uses a U-Net-like encoder-decoder. The encoder has two convolutions, expanding to $t \times f$4 channels and halving frequency resolution, followed by a Dilated DenseNet. A deformable embedding, denoted DSDCN, controls receptive field offsets. Between encoder and decoder, LORT stacks $t \times f$5 locally refined Taylor transformer blocks, each containing LayerNorm, T-MSA + SCEA, a feed-forward network, and a Locally Refined Convolution block. Features pass through T-MSA modules at multiple resolutions using alternating downsampling and upsampling, so that increasingly coarse time-frequency scales capture long-range dependencies before subsequent refinement during upsampling.

The LRC block complements T-MSA by modeling local structure. It integrates convolutional feed-forward layers, time-frequency dense local convolutions, and gated units. The paper characterizes the division of labor as follows: multi-resolution attention captures long-range temporal and spectral interactions for global denoising, while LRC captures fine local time-frequency patterns associated with transient preservation and artifact reduction. This suggests that T-MSA is intended as the global component in an explicitly hybrid global-local design (Wang et al., 28 Sep 2025).

4. Computational profile, normalization, and implementation

The per-head complexity description in LORT distinguishes standard MSA from T-MSA. Standard MSA requires $t \times f$6 to form $t \times f$7 and $t \times f$8 to apply $t \times f$9, with memory $D$0 to store the attention matrix. For first-order T-MSA, the main precomputations are $D$1, $D$2, and $D$3, after which each query needs only inner products with the pre-aggregated statistics. The resulting total is

$D$4

with memory

$D$5

and no $D$6 matrix (Wang et al., 28 Sep 2025).

The overall LORT model has only approximately $D$7 parameters. The stated reasons are the small encoder channel count of $D$8, the use of efficient T-MSA instead of quadratic softmax attention, and the lightweight local convolutions and gated units in LRC. On VCTK+DEMAND with $D$9 blocks, LORT achieves $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$0 FLOPs.

Several implementation details are specifically tied to T-MSA. The Taylor order is $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$1. LayerNorm is placed around attention modules, and residual connections are used in both attention and LRC blocks. SCEA uses 2-D pooling with 1-D convolution in the channel branch and channel-wise pooling with $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$2 convolution in the spatial branch, both with sigmoid activations and multiplicative broadcast modulation. Positional encoding is not explicitly stated; instead, positional context is said to rely on deformable embeddings and multi-resolution processing.

The paper also notes stabilization mechanisms. Because the first-order truncation neglects higher-order terms, LORT adds Multi-Scale Attention Refinement (MSAR), which learns local corrections from $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$3 and $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$4 via lightweight local operations to rectify the Taylor-approximated output. Together with LayerNorm, residual connections, and the optional addition of a small $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$5 in the denominator, these choices are presented as measures that stabilize training.

The STFT configuration relevant to the full system is FFT length $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$6, window $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$7, hop $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$8, and sample rate $O(\mathrm{MHSA}) = 4 t f D^2 + 2 t^2 f^2 D,$9. The stated experimental optimum is hop $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$0; larger hops reduce FLOPs but hurt PESQ, CSIG, and COVL. With $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$1 Taylor blocks and $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$2 channels, the reported runtime profile is approximately $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$3 FLOPs per $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$4, $O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$5 sample, which is described as feasible for modern edge devices (Wang et al., 28 Sep 2025).

5. Training objective and empirical behavior

Within LORT, T-MSA participates in a system that decodes enhanced magnitude and phase spectra independently and optimizes them with a composite loss:

$O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$6

The complex loss is

$O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$7

the magnitude loss is

$O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$8

and the phase loss is decomposed into group delay, instantaneous angular frequency, and instantaneous phase components:

$O(\mathrm{T\text{-}MSA}) = 18 t f D + 2 t f D^2.$9

$D \ll t f$0

$D \ll t f$1

with

$D \ll t f$2

The STFT consistency loss is

$D \ll t f$3

The metric-guided adversarial terms are

$D \ll t f$4

$D \ll t f$5

The paper sets $D \ll t f$6 and uses AdamW (Wang et al., 28 Sep 2025).

On VCTK+DEMAND, LORT with $D \ll t f$7 parameters and $D \ll t f$8 FLOPs achieves PESQ $D \ll t f$9, STOI $t f$0, CSIG $t f$1, CBAK $t f$2, and COVL $t f$3. On the DNS Challenge datasets under general DNS noise, it attains average PESQ $t f$4 and average STOI $t f$5. Under Babble and Factory noise, it delivers the highest average PESQ values of $t f$6 and $t f$7, and top STOI values of $t f$8 and $t f$9.

The ablation results specifically indicate the role of SCEA alongside T-MSA. Removing SCEA lowers PESQ from $O(N^2)$00 to $O(N^2)$01 and COVL from $O(N^2)$02 to $O(N^2)$03. The paper interprets this as evidence that SCEA effectively complements T-MSA by enhancing spatial saliency and inter-channel exchange. Removing CFN or TF-DLC also degrades performance, supporting the stated complementarity between attention-based global modeling and LRC-based local modeling. A plausible implication is that the efficacy of T-MSA in LORT depends not only on its linear attention approximation, but also on the surrounding corrective and local-refinement mechanisms (Wang et al., 28 Sep 2025).

6. Relation to TaylorShift and other efficient attention variants

A broader and more general Taylor-based perspective is provided by TaylorShift, which reformulates Taylor-Softmax so that full token-to-token interactions can be computed in linear time and space with respect to sequence length $O(N^2)$04 while permitting a “shift back” to standard attention for short sequences (Nauen et al., 2024). TaylorShift uses a second-order truncation, $O(N^2)$05, rather than the first-order truncation used in LORT.

In TaylorShift, the attention polynomial is

$O(N^2)$06

with unnormalized numerator and denominator

$O(N^2)$07

and output

$O(N^2)$08

The key algebraic device is the factorization of the elementwise-squared term through a “boxtimes” operator, giving

$O(N^2)$09

This yields an efficient per-head implementation,

$O(N^2)$10

with an analogous denominator computed by replacing $O(N^2)$11 by an all-ones vector or by augmenting $O(N^2)$12 with a ones channel. TaylorShift also introduces a normalization strategy based on L2-normalizing queries and keys, a per-head temperature $O(N^2)$13, scaling by $O(N^2)$14, and scaling $O(N^2)$15 by $O(N^2)$16, with an additional normalization factor proportional to $O(N^2)$17 to maintain numerical stability.

The complexity trade-off in TaylorShift differs from first-order T-MSA. For the efficient second-order implementation, the dominant per-head cost is linear in $O(N^2)$18 but cubic in $O(N^2)$19:

$O(N^2)$20

Its empirical crossover behavior is reported as memory-efficient for sequences as short as approximately $O(N^2)$21 tokens and faster than standard attention for inputs of approximately $O(N^2)$22 tokens and beyond. In full Transformer encoders, TaylorShift uses less memory than the standard Transformer at $O(N^2)$23, is faster at $O(N^2)$24, and reduces memory by approximately $O(N^2)$25 at $O(N^2)$26 and approximately $O(N^2)$27 at $O(N^2)$28 on NVIDIA A100 GPUs (Nauen et al., 2024).

The LORT paper positions first-order T-MSA relative to other efficient attention variants in more qualitative terms. Performer is described as using random features for the softmax kernel, thereby achieving linear attention through Monte Carlo approximations but introducing variance. Linformer is described as reducing complexity by projecting keys and values along the sequence dimension to a fixed rank, potentially imposing a learned low-rank bottleneck. Nyströmformer is described as using landmark points to approximate softmax attention and thus requiring landmark selection. LORT’s T-MSA is characterized as deterministic, polynomial, and based on closed-form global sums. Higher-order polynomial or Taylor attentions are said to improve accuracy at higher cost; LORT therefore adopts $O(N^2)$29 and compensates with MSAR and SCEA. This suggests a spectrum of Taylor-based attention designs in which LORT occupies the lightweight, speech-enhancement-oriented end, while TaylorShift provides a more general second-order formulation for long-sequence Transformers (Wang et al., 28 Sep 2025).

7. Limitations, misconceptions, and scope

A recurring misconception is that all linear-time attention methods discard dense token-to-token interactions. The TaylorShift formulation explicitly argues otherwise: under the polynomial surrogate, all interactions are preserved and the computation is merely reordered through pre-aggregated moments (Nauen et al., 2024). In contrast, LORT’s first-order T-MSA is described in a more task-specific way, emphasizing efficient global context rather than making a general claim about preserving the full fidelity of exact softmax.

The main limitation identified for LORT’s T-MSA is approximation error from first-order truncation. The paper states that accuracy may drop when $O(N^2)$30 has large magnitude or broad distribution. It also notes that reduced cross-channel attention in the T-MSA branch requires explicit compensatory mechanisms such as SCEA. These are not incidental details: SCEA is presented as necessary to restore inter-channel exchange and spatial saliency, while MSAR provides local corrective terms to offset the neglected higher-order Taylor terms.

Resource scaling also imposes a boundary on applicability. T-MSA is motivated for long spectrogram sequences and resource-constrained speech enhancement, especially on-device or low-memory settings. The practical guidance given in LORT is to prefer T-MSA when quadratic attention is prohibitive and when architectures already include local modules such as LRC for fine-grained modeling. Stability recommendations include retaining the $O(N^2)$31 scaling, adding $O(N^2)$32 to the denominator, using LayerNorm and residual connections, and employing MSAR and SCEA (Wang et al., 28 Sep 2025).

For TaylorShift, the corresponding limitations are different. The approximation quality depends on the polynomial order $O(N^2)$33; $O(N^2)$34 is a surrogate to softmax rather than an exact replacement. The efficient implementation scales as $O(N^2)$35 per head, so very large head dimensions reduce the practical advantage of linear-in-$O(N^2)$36 scaling. The linear-time guarantees also depend on mask structure: they hold for all-to-all, causal, and structured masks that allow reusable aggregation, but fully arbitrary per-query masks may lose those guarantees (Nauen et al., 2024).

Taken together, these formulations define T-MSA not as a single universally fixed mechanism but as a family of Taylor-based attention constructions. In LORT, the term refers specifically to a first-order Taylor multi-head self-attention module embedded in a speech-enhancement system with SCEA, LRC, deformable embedding, and multi-resolution processing. In the broader TaylorShift formulation, the same underlying principle is developed into a second-order, general-purpose Transformer mechanism with explicit switching criteria between efficient and standard attention regimes.