Omni Self-Attention in Lightweight Image SR

• Omni Self-Attention (OSA) is a two-stage mechanism that fuses spatial and channel interactions to capture comprehensive omni-axis dependencies in lightweight image SR.
• It addresses the limitations of uni-dimensional self-attention and homogeneous aggregation by sequentially computing spatial and channel covariances through tensor rotations.
• Empirical results demonstrate that OSA improves PSNR benchmarks like Urban100 while offering faster convergence and efficient computation.

Omni Self-Attention (OSA) is an attention primitive introduced in "Omni Aggregation Networks for Lightweight Image Super-Resolution" (Wang et al., 2023) for lightweight image super-resolution (SR). It is defined as a two-stage self-attention mechanism that simultaneously models pixel-interaction from both spatial and channel dimensions, mining the potential correlations across omni-axis (i.e., spatial and channel). In the paper’s formulation, OSA addresses two stated limitations of prior lightweight vision transformers for SR: uni-dimensional self-attention modeling and homogeneous aggregation schemes, both of which are argued to restrict the effective receptive field (ERF) and to contribute to premature saturation in shallow models (Wang et al., 2023).

1. Motivation and problem formulation

OSA was proposed in the context of lightweight SR, where the reconstruction of fine textures requires both long-range spatial dependencies and rich inter-channel correlations (Wang et al., 2023). The motivation is explicitly twofold. First, standard attention in lightweight SR models is described as uni-dimensional. Spatial self-attention models relationships between pixels or tokens across the image plane, but it does not explicitly propagate information between channels. Channel attention models relationships among channels, but usually through scalar reweighting, which the paper characterizes as weaker than true channel-to-channel interaction (Wang et al., 2023).

Second, the paper argues that repeatedly stacking only one kind of operator yields a homogeneous aggregation scheme that is not ideal for SR, because texture structures exist at multiple scales and lightweight models cannot simply become much deeper to compensate (Wang et al., 2023). The proposed response is a broader aggregation philosophy that combines local convolution, meso-scale OSA, and global-scale OSA within the full Omni-SR architecture (Wang et al., 2023).

Within this formulation, OSA is the key attention mechanism rather than the entire network. Its role is to provide the spatial-channel interaction needed at each attention stage, while the surrounding architecture supplies the local, meso-scale, and global pathways (Wang et al., 2023).

2. Core architectural definition

The OSA block is defined as a two-stage attention process: first spatial self-attention, then channel-wise self-attention, with a tensor rotation or permutation connecting the two stages (Wang et al., 2023). Given an input feature map $X \in \mathbb{R}^{HW \times C}$, the block first forms spatial query, key, and value projections: $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$

It then computes spatial attention as

$$Y_s = \mathcal{A}^s(Q^s, K^s, V^s) = \text{SoftMax}(Q^s K^{sT}) \cdot V^s$$

which yields an attention map over spatial positions of size $HW \times HW$ (Wang et al., 2023).

The representation is subsequently rotated or transposed so that channel relationships become the attention axis. The paper denotes this with a rotation operator $\mathcal{R}(\cdot)$ and writes

$$Q^c = \mathcal{R}(Q'),\quad K^c = \mathcal{R}(K'),\quad V^c = \mathcal{R}(V')$$

where $Q',K',V'$ are the inputs to channel attention, either taken from the spatial stage or directly copied from $Q^s, K^s, V^s$ (Wang et al., 2023).

Channel-wise attention is then computed as

$$Y_c = \mathcal{A}^c(Q^c, K^c, V^c) = \text{SoftMax}(K^c Q^{cT}) \cdot V^c$$

which yields a $C \times C$ channel interaction map (Wang et al., 2023). The inverse rotation returns the result to the original orientation. The defining property is therefore not parallel fusion but a cascade in which spatial covariance and channel covariance are computed sequentially through matrix operations (Wang et al., 2023).

This design is presented as dense interaction modeling. The paper states that OSA cascades spatial covariance and channel covariance rather than treating channel attention as a simple scalar gate, and it frames the resulting interaction space as “omni-axis” (Wang et al., 2023).

3. Mathematical interpretation and comparison with related attention forms

The paper’s central mathematical claim is that OSA does not merely concatenate spatial and channel attention outputs; instead, it sequentially computes them so that the output of one influences the other through tensor rearrangement (Wang et al., 2023). The essential computation is summarized as

$$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$0

followed by

$$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$1

after rotation into the channel-oriented representation (Wang et al., 2023).

The authors describe this as the “cascading computation of the cross-covariance matrices between spatial/channel dimensions,” enabling dense interactions (Wang et al., 2023). This suggests that OSA is intended to capture not only which spatial positions matter or which channels matter, but also how those two dependency structures interact within a single block.

The comparison to other lightweight attention modules is explicit. Relative to standard spatial self-attention, OSA adds explicit channel mixing; standard spatial attention can use channel features to compute spatial scores, but it does not truly model channel-to-channel relations in the output (Wang et al., 2023). Relative to SE-style channel attention, the distinction is that SE produces scalar weights per channel, whereas OSA computes a full channel affinity matrix so channels can exchange information rather than only be rescaled (Wang et al., 2023). Relative to CBAM or BAM-style hybrids, the paper states that those modules combine spatial and channel attention, but channel attention remains scalar or gating-based; OSA instead uses covariance-based self-attention in the channel branch (Wang et al., 2023).

The paper also notes that OSA can be used as a drop-in replacement for attention blocks in Swin-like architectures and can work with variants such as Swin or Halo-style windowing (Wang et al., 2023). In context, this is important because it locates OSA as a reusable attention operator rather than a mechanism tied to a single backbone.

4. Efficiency, window partitioning, and scale-specific variants

Although OSA models both spatial and channel dimensions, it is explicitly designed to remain lightweight (Wang et al., 2023). The principal efficiency mechanism is window partitioning. Instead of applying full attention globally over all $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$2 tokens, attention is computed within smaller groups or windows, which reduces the size of the attention matrix substantially (Wang et al., 2023).

The paper specifically states that OSA uses a smaller attention map size in the channel self-attention stage and is less computationally intensive than cascade shifted-window self-attention in Swin (Wang et al., 2023). Scale-specific variants are built by changing the window partition strategy.

For Meso-OSA, the feature map is partitioned into non-overlapping $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$3 windows, with the reshaping described as

$$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$4

(Wang et al., 2023).

For Global-OSA, the feature is partitioned into a uniform $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$5 grid, each grid cell covering $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$6, with the reshape

$$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$7

(Wang et al., 2023).

These reshaping rules are important because they show that meso-scale and global-scale OSA are not distinct attention formulas. They are instances of the same OSA mechanism applied under different partitioning schemes, allowing the receptive field to change without abandoning the shared spatial-channel interaction design (Wang et al., 2023).

5. Function inside Omni-SR and the Omni-Scale Aggregation Group

OSA is used inside the Omni-Scale Aggregation Group (OSAG), which is the repeated building block of Omni-SR (Wang et al., 2023). Each OSAG contains four components:

1. Local Convolution Block (LCB): a lightweight convolutional module using pointwise and depthwise convolutions plus channel attention, aimed at local detail extraction.
2. Meso-OSA block: OSA within local windows to capture mid-range interactions.
3. Global-OSA block: OSA on a sparse grid to capture global interactions efficiently.
4. ESA block: a final enhancement module from prior lightweight SR work (Wang et al., 2023).

The architecture is therefore organized as a hierarchical progression from local convolution for fine details, to meso-scale attention for mid-range patterns, to global attention for broader contextual dependencies (Wang et al., 2023). OSA is the operator that enables both attention stages in a unified way.

The paper states that OSAG enables information propagation between any pair of tokens in the feature map, but with a more balanced receptive field than using only one kind of operator (Wang et al., 2023). A plausible implication is that OSA should be understood less as an isolated attention formula and more as the attention core of a multi-scale aggregation design whose surrounding modules are intended to mitigate ERF saturation in shallow lightweight SR networks.

6. Empirical findings and ablation evidence

The empirical results reported for Omni-SR attribute part of the model’s performance to the omni-axis interactions provided by OSA (Wang et al., 2023). The paper reports record-high performance on lightweight SR benchmarks, including 26.95 dB@Urban100 $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$8 with only 792K parameters, and the more detailed benchmark results include 26.64 dB on Urban100 $$Q^s = X \mathcal{W}_q,\quad K^s = X \mathcal{W}_k,\quad V^s = X \mathcal{W}_v$$9 with DIV2K only, 26.95 dB with DF2K, 31.02 dB on Manga109 $$Y_s = \mathcal{A}^s(Q^s, K^s, V^s) = \text{SoftMax}(Q^s K^{sT}) \cdot V^s$$0 with DIV2K only, and 31.50 dB with DF2K (Wang et al., 2023).

The most direct evidence for OSA comes from the ablations on OSA variants. Removing channel-wise interaction produces the spatial-only variant $$Y_s = \mathcal{A}^s(Q^s, K^s, V^s) = \text{SoftMax}(Q^s K^{sT}) \cdot V^s$$1 and reduces performance by about 0.13 dB compared with the full model. Removing spatial self-attention yields the channel-only variant $$Y_s = \mathcal{A}^s(Q^s, K^s, V^s) = \text{SoftMax}(Q^s K^{sT}) \cdot V^s$$2 and also reduces performance (Wang et al., 2023).

Variant	Urban100 $$Y_s = \mathcal{A}^s(Q^s, K^s, V^s) = \text{SoftMax}(Q^s K^{sT}) \cdot V^s$$3
spatial-only	26.51
channel-only	26.45
full OSA	26.64

These results are presented in the paper as evidence that both axes are necessary and that the full spatial-plus-channel cascade performs best (Wang et al., 2023). The paper also states that replacing OSA components with SE or CBAM lowers PSNR, which is used to support the claim that covariance-based interaction is stronger than scalar gating in this setting (Wang et al., 2023).

The optimization analysis further reports that OSA converges faster than spatial-only or channel-only attention, that features produced by OSA have higher normalized entropy, and that OSA gives a higher diffusion index, indicating stronger long-range interaction (Wang et al., 2023). This supports the view that OSA affects both final reconstruction quality and training behavior.

7. Terminological scope and relation to other “omni” attention mechanisms

The term “Omni Self-Attention” in the strict sense refers to the spatial-plus-channel self-attention block introduced for lightweight image SR in Omni-SR (Wang et al., 2023). It should not be conflated with every attention mechanism that uses “omni,” “omnidirectional,” or “shared multimodal self-attention” terminology.

A directly related but distinct mechanism appears in "Omni-directional attention mechanism based on Mamba for speech separation" (Xue et al., 23 Jan 2026). That paper proposes an efficient omni-directional attention mechanism built upon unidirectional Mamba, modeling global dependencies from ten different directions on the spectrogram while preserving linear complexity (Xue et al., 23 Jan 2026). The conceptual similarity is that both methods seek more comprehensive dependency modeling than a single-axis or single-order operator. The mechanism, however, is different: OA uses Mamba-based multi-directional scanning over a 2D time-frequency spectrogram and channel scans, rather than the quadratic self-attention cascade of spatial and channel covariance used by OSA (Xue et al., 23 Jan 2026).

A different usage appears in "iFLYTEK-Embodied-Omni Technical Report" (Zhang et al., 24 Jun 2026), where image, language, video, and action tokens are concatenated and passed through shared Omni Multi-Modal Self-Attention layers under modality-aware masking (Zhang et al., 24 Jun 2026). Here the emphasis is cross-modal alignment in a unified embodied foundation model, not the spatial-channel fusion problem of low-level vision. Similarly, "OmniNet: Omnidirectional Representations from Transformers" (Tay et al., 2021) uses omnidirectional attention to let tokens attend across both width and depth of the network via a meta-learner over stacked hidden states (Tay et al., 2021). This is again distinct from OSA’s two-stage spatial and channel self-attention block.

The distinction is also important in the opposite direction. "Self-attention for Enhanced OAMP Detection in MIMO Systems" (Fuchs et al., 2023) introduces a self-attention-enhanced OAMP detector, but the paper explicitly does not describe a distinct “Omni Self-Attention (OSA)” architecture; it uses standard transformer-style self-attention adapted to MIMO tokenization (Fuchs et al., 2023).

Taken together, these papers indicate that “omni” attention terminology has become a broader label for mechanisms that enlarge dependency modeling beyond a single axis, layer, modality, or scan order. In the technical literature, however, OSA most specifically denotes the two-stage spatial-plus-channel self-attention block introduced for lightweight SR in Omni-SR (Wang et al., 2023).