Residual-Post-Norm & Cosine Attention

• The paper demonstrates that Residual-Post-Norm repositions layer normalization after residual sums to counteract runaway activation growth in deep vision models.
• Scaled Cosine Attention replaces dot-product computations with bounded cosine similarity, ensuring stable and informative attention weights during training.
• Empirical results on Swin Transformer V2 show improved top-1 accuracy and convergence stability across model sizes, enabling efficient training of billion-parameter architectures.

Residual-Post-Norm and Cosine Attention are joint optimization techniques developed to address training instability in large-scale vision Transformers, especially as model depth and input resolution are scaled. These methods, as introduced in the Swin Transformer V2 architecture, reconfigure normalization and the self-attention mechanism to enhance gradient stability, dampen runaway activation growth, and facilitate convergence in deep or wide models. The core innovations are (1) the Residual-Post-Norm arrangement, which positions layer normalization after each residual addition, and (2) Scaled Cosine Attention, which replaces unbounded dot-product attention with bounded cosine similarity logits modulated by a learnable temperature. Empirical results on large vision models demonstrate that this combination enables the robust and efficient training of billion-parameter models at unprecedented input resolutions (Liu et al., 2021).

1. Transformer Block Normalization: Pre-Norm versus Residual-Post-Norm

Standard Transformer blocks incorporate two principal sub-layers: multi-head self-attention (MSA) and a feed-forward network (MLP), each with residual connections and layer normalization (LN). The conventional "Pre-Norm" scheme applies LN before each sub-layer, resulting in activations that are normalized prior to attention or MLP operations. Formally, for input $x_{\ell-1}$ at layer $\ell$, Pre-Norm executes:

$$a_\ell = \mathrm{LN}(x_{\ell-1}) \ z_\ell = x_{\ell-1} + \mathrm{MSA}(a_\ell) \ b_\ell = \mathrm{LN}(z_\ell) \ x_\ell = z_\ell + \mathrm{MLP}(b_\ell) \ \text{output}_\ell = \mathrm{LN}(x_\ell)$$

In contrast, Residual-Post-Norm delays the application of LN until after each residual sum. If $F_1(x) = \mathrm{MSA}(x)$ and $F_2(x) = \mathrm{MLP}(x)$, then:

$$z_1 = x_{\ell-1} + F_1(x_{\ell-1}) \ x' = \mathrm{LN}(z_1) \ z_2 = x' + F_2(x') \ x_\ell = \mathrm{LN}(z_2)$$

There is no normalization preceding either sub-layer; rather, normalization re-centers and re-scales activations after each addition, countering the accumulation of unnormalized outputs. This structure empirically prevents the "runaway" growth of activation magnitudes observed when training very deep or wide models (Liu et al., 2021).

2. Scaled Cosine Attention Mechanism

Traditional attention in Transformers employs scaled dot-product logits, computed as $\mathrm{Softmax}(QK^\top/\sqrt{d} + B)$, where $B$ denotes relative positional bias terms. Excessive scaling of input vectors ($\|Q\|, \|K\|$) leads to highly peaked or flat attention, driving instabilities especially when activations drift in magnitude.

Scaled Cosine Attention replaces the dot-product with bounded cosine similarity, eliminating dependence on activation norm. Each attention logit is given by:

$$\mathrm{sim}(q_i, k_j) = \frac{q_i \cdot k_j}{\|q_i\| \|k_j\|}\cdot \frac{1}{\tau} + B_{ij}$$

where $\ell$0, $\ell$1 is a learnable temperature parameter (distinct per-head and per-layer, lower-bounded by $\ell$2), and $\ell$3 is the positional bias. The attention weights are then normalized via softmax:

$\ell$4

and used to compute the output:

$\ell$5

This approach ensures that attention logits remain bounded and focus exclusively on angular similarity, yielding more robust gradients and preventing logit explosions (Liu et al., 2021).

3. Stabilization Rationale and Empirical Characteristics

Layer normalization immediately after each residual addition ("Post-Norm") re-centers and scales the amplitude of activations within each layer, arresting the progressive magnitude drift observed in Pre-Norm networks. This stabilization is essential for extremely deep or wide Transformer models, where uncorrected accumulation of residuals can provoke gradient collapse or explosion. Figure 1 of (Liu et al., 2021) demonstrates that Pre-Norm leads to activation growth across four orders of magnitude, while Post-Norm constrains activations to a tightly bounded range.

Likewise, Scaled Cosine Attention removes the contribution of vector magnitude to attention logits. Dot-product attention can result in excessively sharp or nearly uniform distributions when $\ell$6 and $\ell$7 norms drift, indicating less informative attention maps and deteriorated gradient flow. Cosine attention enforces bounded, direction-only similarity, empirically yielding gentler gradients and improved robustness (Liu et al., 2021).

4. Implementation within Swin Transformer V2 Blocks

The integration of these techniques in Swin Transformer V2 modifies each encoder block as follows. Standard Window-based Multi-Head Self-Attention and MLP sub-layers are composed, with residual sums immediately followed by LN. Pseudocode for a block is:

$\ell$9 An additional LayerNorm is inserted on the main branch after every six such blocks to further damp persistent magnitude drift when models scale to hundreds of layers (Liu et al., 2021).

5. Quantitative Ablations and Performance Outcomes

Empirical results on ImageNet-1K demonstrate consistent top-1 accuracy improvements when moving from (a) Pre-Norm + dot-product, through (b) Residual-Post-Norm + dot-product, to (c) Residual-Post-Norm + Scaled Cosine Attention:

Model	Pre-Norm + Dot-prod	Post-Norm + Dot-prod	Post-Norm + Cosine
Swin-T	81.5%	81.6%	81.7%
Swin-S	83.2%	83.3%	83.6%
Swin-B	83.6%	83.8%	84.1%

Furthermore, in a 658M-parameter setting, Pre-Norm training under SimMIM self-supervised pre-training diverged, whereas Residual-Post-Norm with Scaled Cosine Attention enabled smooth convergence. Comparative analyses with alternative LN configurations (including Sandwich Norm and original Post-Norm) identify the specific Residual-Post-Norm schedule in Swin V2 as optimal for stability–capacity trade-off without compromising representational expressiveness (Liu et al., 2021).

6. Implications and Significance for Large-Scale Vision Models

The adaptation of Residual-Post-Norm and Scaled Cosine Attention constitutes the foundation for Swin Transformer V2's ability to handle billion-scale parameters and high-resolution imagery (up to 1,536$\ell$81,536). These techniques address fundamental instabilities that otherwise preclude successful scaling of vision models. A plausible implication is that further advances in vision Transformer scale and deployment will require continued progression along the axes of normalization scheduling and attention mechanism design (Liu et al., 2021). The efficiency gains reported—orders of magnitude less labeled data and training time compared to prior billion-level visual models—are directly attributable to these stabilization strategies.