Neural Discrimination-Prompted Attention

• Neural Discrimination-Prompted Attention (NDPA) is a technique that injects explicit high-resolution priors into Transformer attention to guide feature refinement.
• NDPA integrates a two-stage cross-attention approach where low-resolution queries first attend to NDP-derived keys and values and then re-attend over refined low-resolution features.
• Empirical results demonstrate that NDPA improves restoration performance and efficiency by preserving critical high-resolution details in UHD image enhancement tasks.

Neural Discrimination-Prompted Attention (NDPA) is the attention mechanism inside the UHDPromer architecture, introduced for Ultra-High-Definition image restoration and enhancement in "Neural Discrimination-Prompted Transformers for Efficient UHD Image Restoration and Enhancement" (Wang et al., 1 Mar 2026). It injects an explicit prior—Neural Discrimination Priors (NDP)—into attention computation so that low-resolution features selectively attend to regions where high-resolution and low-resolution representations disagree most strongly. Within the overall design, NDPA operates as a two-stage, NDP-guided attention mechanism in a low-resolution Transformer branch, while complementary modules use the same prior to gate feed-forward processing and support super-resolution-guided reconstruction.

1. Neural Discrimination Priors as the basis of NDPA

UHDPromer is built on the observation that there implicitly exist neural differences between high-resolution and low-resolution features, and that exploring such differences can facilitate low-resolution feature representation. The architecture therefore maintains a high-resolution branch, HRFR, that extracts multi-scale features

$$\{\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\} \in \mathbb{R}^{H \times W \times C},$$

and a low-resolution branch, NDPT, operating on shuffle-down features

$$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$

where $s=8$ is the shuffle-down factor.

The prior used by NDPA is the Neural Discrimination Prior, defined for the $i$-th NDPT block and low-resolution location $x$ as

$\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$

Here $\mathcal{H}_i[\cdot,\cdot,\cdot]$ concatenates multi-scale high-resolution features along the channel dimension and applies a stride convolution with kernel size $s \times s$ and stride $s \times s$, producing an aligned high-resolution-derived low-resolution feature

$$\tilde{\mathbf{X}}_i(x)=\mathcal{H}_i[\mathbf{X}_1,\mathbf{X}_2,\mathbf{X}_3](x)\in\mathbb{R}^C.$$

The element-wise absolute value, exponential, and inverse square root compress the difference magnitude into $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$0.

The associated interpretation is that NDP is a spatial discriminativeness map. Ignoring the slight notation abuse, it can be viewed as decreasing with the magnitude of the discrepancy

$$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$1

This indicates that NDP encodes, per pixel, how structurally significant or misrepresented a location is in the low-resolution domain, as inferred from multi-scale high-resolution features. The text also states that when $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$2 approaches 1, the feature at position $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$3 notably diverges from low-resolution features and indicates greater discriminative potential; this is slightly inconsistent with the literal formula as written. The central conceptual point remains that NDP is used as an explicit prior about high-resolution/low-resolution mismatch.

2. Two-stage attention formulation

NDPA re-formulates attention by incorporating NDP to globally perceive useful discrimination information. Within a Neural Discrimination-Prompted Transformer Block, normalized low-resolution features $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$4 are projected into queries, keys, and values, while NDP features $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$5 are projected into their own keys and values: $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$6 The operators are $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$7, a $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$8 pointwise convolution; $$\mathbf{Y}_i \in \mathbb{R}^{H/s \times W/s \times C},$$9, a $s=8$0 depthwise convolution; and $s=8$1, splitting into multi-head tensors. The attention kernel is

$s=8$2

where $s=8$3 is a learnable scale.

Operationally, NDPA proceeds in two stages. First, it performs cross-attention between low-resolution queries and NDP-derived keys and values, producing a global discrimination-aware intermediate representation: $s=8$4 Second, it reuses this output as a prompt when computing attention with low-resolution keys and values. The paper gives this second step schematically rather than as a fully expanded algebraic expression, but the intended design is explicit: cross-attend to NDP, then re-attend over low-resolution features conditioned on the discrimination prompt.

This distinguishes NDPA from both standard self-attention and standard cross-attention. Standard self-attention uses $s=8$5, $s=8$6, and $s=8$7 from the same source. Standard cross-attention typically performs a single pass from one source to another. NDPA instead uses a two-stage procedure in which NDP functions as an external prompt that biases subsequent low-resolution attention patterns.

3. Placement within UHDPromer

The full UHDPromer pipeline begins with shallow embedding of the input UHD image $s=8$8: $s=8$9 Three ConvNeXt-v2 blocks in HRFR then compute

$i$0

A shuffle-down operation maps these features to

$i$1

A stack of $i$2 Neural Discrimination-Prompted Transformer Blocks operates on this low-resolution representation, guided by NDP (Wang et al., 1 Mar 2026).

Within the $i$3-th block, NDPA is coupled to a Neural Discrimination-Prompted Network (NDPN) through residual connections: $i$4 Here $i$5 is the low-resolution block input, $i$6 is LayerNorm, and $i$7 is the block-specific NDP feature map computed from current low-resolution features and multi-scale high-resolution features.

The architectural significance of NDPA is therefore twofold. First, it is the mechanism by which high-resolution/low-resolution discrepancy is converted into globally propagated discrimination information inside the low-resolution Transformer stream. Second, it is not an isolated attention replacement: it is paired with NDP-guided feed-forward gating in NDPN, so both attention and channel-spatial transformation are conditioned on the same prior.

4. Coupling with NDPN and reconstruction pathway

NDPN explores a continuous gating mechanism guided by NDP to selectively permit the passage of beneficial content. Its formulation is

$i$8

where $i$9 is channel concatenation, $x$0 is GELU, and $x$1 is element-wise multiplication. In effect, NDPN acts as an NDP-aware MLP: it fuses NDP with low-resolution content and uses gated transformations so that positions identified as discriminative can pass more information.

After the low-resolution NDPT stage, UHDPromer applies FeaSR, a feature super-resolution module that upsamples the NDPT output and produces a super-resolved image

$x$2

A subsequent super-resolution-guided reconstruction module, SRG-Recon, fuses HRFR outputs and FeaSR outputs via three ConvNeXt-v2 blocks, two $x$3 convolutions, and one $x$4 convolution, producing a residual $x$5 and the final restored image

$x$6

Training uses the joint loss

$x$7

where $x$8 is a spatial+frequency loss and $x$9. The relation between reconstruction and NDPA is direct but modular: NDPA and NDPN first produce low-resolution features already optimized to correct high-resolution/low-resolution differences via NDP; FeaSR then expands those NDP-refined features back to full resolution; SRG-Recon uses the super-resolved prediction together with high-resolution features as guidance for the final UHD output.

5. Empirical behavior and efficiency

UHDPromer is evaluated on three UHD image restoration and enhancement tasks: low-light image enhancement, image dehazing, and image deblurring, under two settings: training on general-resolution datasets and testing on UHD, and training and testing on UHD. In Setting 2, the reported results are 27.159 dB PSNR and 0.9285 SSIM on UHD-LL; 22.725 dB PSNR and 0.9432 SSIM on UHD-Haze; and 29.527 dB PSNR, 0.8584 SSIM, and LPIPS 0.2163 on UHD-Blur. The same section reports that UHDPromer achieves the best computational efficiency while still maintaining state-of-the-art performance on the three tasks (Wang et al., 1 Mar 2026).

The computational profile is central to the design rationale. UHDPromer has 0.7430M parameters. Reported $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$0 FLOPs are 32.56G, and runtime on the same GPU is 0.12 s. The low-resolution branch enabled by shuffle-down to $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$1 is the principal reason attention remains tractable at UHD scale, while NDPA and NDPN are intended to preserve the missing structural information that a naive low-resolution Transformer would discard.

The ablation entitled "Effect on Neural Discrimination-Prompted Transformers" isolates the role of NDPA and NDP. The full model, with NDP in both NDPA and NDPN, achieves 27.159 dB PSNR and 0.9285 SSIM on the main branch. Without NDP in both NDPA and NDPN, performance becomes 26.811 dB and 0.9282 SSIM. Without NDP in NDPA only, it drops to 26.183 dB and 0.9252 SSIM. Without NDP in NDPN only, it is 27.026 dB and 0.9283 SSIM. Replacing NDP with direct features yields 26.138 dB and 0.9272 SSIM, and using NDP only before the block rather than inside attention and feed-forward integration yields 26.161 dB and 0.9263 SSIM. These results indicate that removing NDP from NDPA causes the largest degradation and that the specific definition of NDP as a high-resolution/low-resolution difference is crucial.

A complementary finding is that enabling SR-guided reconstruction and supervising $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$2 yields a clear boost, especially in SSIM, compared to naive cascaded reconstruction without SR guidance. This situates NDPA not as an isolated mechanism but as part of a tightly coupled low-resolution-to-high-resolution restoration strategy.

6. Relation to other attention and prompting mechanisms

NDPA differs from standard self-attention because it is not limited to pairwise similarity within a single feature space. It differs from ordinary cross-attention because its first cross-attention stage is not the terminal operation; instead, the output of that stage becomes a prompt for subsequent attention over low-resolution keys and values. The mechanism is therefore discrimination-aware rather than merely multi-source.

The paper also contrasts NDPA with prompt-based vision Transformers in which learnable tokens or prompts are inserted into the input sequence to adapt pretrained models or condition on degradation. In NDPA, NDP is not a learned parameter. It is a data-driven prior derived from the actual discrepancy between high-resolution and low-resolution features at each block. In addition, NDP is used as an external attention source with its own keys and values, not simply as extra tokens appended to the low-resolution sequence.

Relative to other attention modifications for high-resolution and UHD restoration, NDPA is presented as an alternative to axis-based or windowed attention schemes and to heavy correlation-based transformations between high-resolution and low-resolution features. Instead of constructing large cross-correlation matrices, it computes a simple local difference via stride convolutions and exponentiation, then integrates that difference into attention as a prompt. This suggests an architectural trade-off in which explicit high-resolution/low-resolution discrepancy acts as a low-cost surrogate for more expensive correlation modeling.

The paper further implies broader applicability beyond the evaluated tasks. A plausible implication is that the same pattern—using differences between a richer reference representation and a compressed representation as a discrimination prior to prompt attention—could be generalized to other multi-resolution or cross-modal settings. The explicit examples given are super-resolution, HDR, pan-sharpening, and multimodal cases such as depth versus RGB or infrared versus RGB. The paper itself, however, focuses on UHD restoration and enhancement rather than experimentally validating those extensions.

7. Implementation characteristics and interpretive issues

Several implementation details shape the practical behavior of NDPA. The NDPT stack has depth $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$3. The base channel dimension is $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$4. Each NDPA uses 8 attention heads, and NDP is computed per NDPT block, making the priors layer-wise and dynamic. The stride convolution used in NDP computation has kernel and stride $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$5, matching the shuffle-down factor $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$6. Training uses AdamW with initial learning rate $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$7 decayed to $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$8 by cosine decay, and a training patch size of $\begin{split} NDP_{i}(x) = \frac{1}{\sqrt{e^{\textrm{abs}\left| \mathcal{H}_{i}\left[\mathbf{X}_1, \mathbf{X}_2, \mathbf{X}_3\right]\left(x\right)- \mathbf{Y}_{i}\left(x\right) \right|}}. \end{split} \tag{1}$9.

The underlying attention formulation is essentially Restormer’s multi-Dconv head transposed attention, but with NDP embedded into the computation. This matters because NDPA is not a wholly separate attention family; rather, it is a prior-conditioned modification of an existing efficient image-restoration Transformer pattern. Likewise, NDPN reuses depthwise and pointwise convolutional structure to preserve low cost while injecting NDP-guided gating.

An interpretive caution concerns the relationship between the literal NDP formula and the textual explanation of discriminative potential. The formula suggests that larger high-resolution/low-resolution discrepancies produce smaller NDP values, whereas the text states that values approaching 1 indicate greater divergence and greater discriminative potential. This inconsistency does not alter the operational role of the module inside the architecture: in all descriptions, NDP functions as a spatial prior indicating where the low-resolution representation is weak, ambiguous, or structurally underrepresented, and NDPA uses that prior to bias attention toward beneficial global information.

In summary, NDPA is a two-stage, NDP-conditioned attention mechanism embedded in an efficient low-resolution Transformer for UHD image restoration. Its defining features are the explicit construction of a high-resolution/low-resolution discrimination prior, the use of that prior as an external attention source, and the integration of discrimination-aware attention with NDP-guided feed-forward gating and super-resolution-guided reconstruction.