---
title: Multi-size Swin Transformer Block (MSTB)
url: https://www.emergentmind.com/topics/multi-size-swin-transformer-block-mstb
type: topic
---

# Multi-size Swin Transformer Block (MSTB)

The Multi-size Swin Transformer Block (MSTB) is a specialized architecture developed to enhance the representational efficiency and computational performance of vision Transformers, with principal applications in single image super-resolution and wide-angle portrait correction. MSTB leverages multi-window, parallel, and multi-scale self-attention mechanisms based on Swin Transformer principles to capture information at diverse spatial scales within lightweight models. Two prominent MSTB variants are introduced in the contexts of MSwinSR [2210.11019] (image super-resolution) and MS-UNet [2109.08024] (portraits correction), each exhibiting distinct configurations tailored to their respective domains.

## 1. Multi-size Swin Transformer Block: Core Architectural Principles

The MSTB, as instantiated in MSwinSR [2210.11019], operates on an input tensor $F_0 \in \mathbb{R}^{H \times W \times C}$. It simultaneously applies four parallel Swin-style Multi-Head Self-Attention (MSA) modules, each differing in window size and shift configuration:
- **W-MSA**: regular window, size $M \times M$, no shift.
- **SW-MSA**: shifted window (by $\lfloor M/2 \rfloor$ pixels), same size $M \times M$.
- **W-MSA-½**: regular window, size $(M/2) \times (M/2)$.
- **SW-MSA-½**: shifted window (by $\lfloor (M/2)/2 \rfloor$), size $(M/2) \times (M/2)$.

Each path begins with LayerNorm, then computes MSA, and concludes with a residual addition of $F_0$. Their outputs $F_1, F_2, F_3, F_4$ are concatenated along the channel dimension, followed by an additional LayerNorm and a compact MLP. The MLP reduces channel dimension in two steps ($4C \to 2C \to C$), each time applying GELU activation. The MLP output is finally added back to $F_0$ as a residual. This architecture enhances width-wise expressivity by fusing multi-scale contexts via parallel processing.

## 2. Self-Attention Mechanism and Multi-size Operation

Within each parallel MSA branch, the feature tensor $X$ (pre-LN) is linearly projected into $Q, K, V$ using weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ (where $d = C$). For $h$ heads, each head has $d_k = d / h$-dimensional $Q^\ell, K^\ell, V^\ell$ with $\ell=1, \ldots, h$. Attention in each window is computed as:
$$
\text{Attention}(Q^\ell, K^\ell, V^\ell) = \operatorname{softmax}\left( \frac{Q^\ell (K^\ell)^T}{\sqrt{d_k}} + B^\ell \right) V^\ell
$$
where $B^\ell$ is a learnable relative position bias per head and branch. The attended representations from all heads are concatenated and projected via $W_O$. The primary means of multi-size operation is realized by using different window partitions (i.e., $M \times M$ and $(M/2) \times (M/2)$, with/without shifts), enabling the block to concurrently process information at varied receptive field scales.

## 3. Positional Encoding, Normalization, and MLP Design

Positional encoding in MSTB strictly follows the Swin Transformer scheme. Each branch maintains a bias table of size $(2M-1) \times (2M-1)$ per head for the larger windows, and $(M-1) \times (M-1)$ for the smaller windows. During attention, the appropriate slice of $B^\ell$ is indexed by the relative coordinates within the window.

Layer normalization is applied before every MSA in each branch and again post-concatenation before the shared MLP. The MLP employs two linear layers with GELU activation, following a $4C \to 2C \to C$ channel reduction. These design choices ensure each output maintains the original spatial and channel dimensions for subsequent residual addition.

## 4. Integration into Network Architectures and Computational Impact

In the MSwinSR framework [2210.11019], MSTBs are organized into stages that replace the Residual Swin Transformer Blocks (RSTBs) of SwinIR. Specifically, each stage contains $L$ MSTBs, followed by a $3 \times 3$ convolution and a long-skip residual connection. The default model utilizes three stages with depth vector $[2,2,2]$, totaling six MSTBs. This arrangement yields a total of 24 MSA calls—comparable to SwinIR's six Swin layers per RSTB across four RSTBs.

Crucially, placing four parallel MSAs before a single, shared MLP results in a $4\times$ reduction in MLP count, leading to substantial parameter and FLOP reductions. On the CelebA $4\times$ super-resolution benchmark, MSwinSR with MSTB achieves a parameter reduction of $30.7\%$ (897.2K to 621.9K), FLOP reduction of $9.94\%$ (4.187G to 3.771G), and a $0.07$ dB gain in PSNR compared to SwinIR (see Table 1).

| Model Variant           | Parameter Count | Multi-adds | PSNR (dB) |
|------------------------|-----------------|------------|-----------|
| SwinIR (4×6 layers)    | 897,200         | 4.187G     | 29.47     |
| MSwinSR [2,2,2] (MSTB) | 621,900         | 3.771G     | 29.54     |

## 5. Multi-Scale Swin Transformer Block in Portrait Correction Networks

In wide-angle portrait correction [2109.08024], MSTB denotes a related but distinct block: the Multi-Scale Swin Transformer Block. Here, each MSTB comprises two serial “sub-blocks” that alternate between normal and shifted window partitioning, window-based MSA, residual plus LN, and an MLP with GELU activation. Importantly, each sub-block fuses two query/key/value branches:
- **Global window branch** derives $Q$ directly from the windowed (possibly shifted) input.
- **Local Dense Connection Module (DCM) branch** extracts $K$ and $V$ through a small residual module employing three depthwise separable $3 \times 3$ convolutions (with dilations $\{1,2,3\}$), then windows and flattens the output.

The attention function fuses these scales by computing attention globally over windows while incorporating local structural priors through DCM-based keys and values. Alternating shifted and regular windows ensures cross-window information percolation and alleviates blocking artifacts. The MLP maintains a two-layer, channel-expanding (commonly $r=4$) structure. Parameter and FLOP complexity scale quadratically with channel dimension per block, $\mathcal{O}(C^2)$.

## 6. Hyperparameters and Ablation Findings

In MSwinSR MSTB [2210.11019], default settings are:
- Embedding dimension $C = 60$
- Number of heads $h = 6$, so $d_k = 10$
- Window sizes: $M=8$ for “full” branches; $M/2=4$ for “half” branches
- MLP path: $4C \to 2C \to C$

For depth ablation with constant total attention calls (24 per model), $[2,2,2]$ staging yields superior PSNR/parameter tradeoffs over $[1,1,1,1,1,1]$. In portrait correction MSTB [2109.08024], architectural choices are less explicit; values for key hyperparameters (window size, channel count, head number) are not specified, but design follows typical Swin-T-sized backbones. Each MSTB is preceded/followed by standard Swin Transformer network elements.

## 7. Comparative Analysis and Applicability

MSTBs, across both MSwinSR and MS-UNet, advance Transformer-based vision models by concurrently mixing multi-scale (and, in MS-UNet, local-global) feature extraction within single blocks. By enabling parallel windowed self-attention computations over multiple partitionings, they grant the network enhanced feature aggregation without the parameter or computational penalties of stacking many deep blocks. The MSTB design results in more parameter- and computation-efficient models with either improved or equivalent reconstruction metrics relative to conventional Swin Transformer architectures [2210.11019], [2109.08024]. These advantages suggest MSTBs are well suited for real-time or resource-constrained vision tasks requiring rich hierarchical representation synthesis at modest computational cost.

Source: https://www.emergentmind.com/topics/multi-size-swin-transformer-block-mstb