Stop-Gradient Attention (SGA)

• Stop-Gradient Attention (SGA) is a modified dot-product attention mechanism that eliminates gradient conflicts by halting back-propagation in the Q/K branch.
• SGA enhances training stability and output quality in reference-based line-art colorization, achieving lower FID scores and higher SSIM metrics.
• The method employs stop-gradient operations, double normalization, and residual connections to deliver significant performance improvements over traditional attention mechanisms.

Stop-Gradient Attention (SGA) is a modified attention mechanism designed to address gradient conflict in neural networks utilizing dot-product attention, particularly within the context of reference-based line-art colorization. By eliminating the gradient flow through the attention map computation for specific branches, SGA prevents destabilizing negative interactions during back-propagation, leading to improved training stability, more robust convergence, and enhanced perceptual and structural quality in generated outputs (Li et al., 2022).

1. Motivation and Theoretical Foundation

Dot-product attention, foundational in modern neural architectures, computes multiple branches of gradients: the skip (residual) connection branch, the Q/K branch (resulting from attention map computation), and the V branch (value aggregation). When applied to reference-based line-art colorization with complex objectives such as self-supervised reconstruction and GAN-based losses, empirical observations reveal that gradients from the skip and V branches align with the overall optimization direction. In contrast, the Q/K branch frequently exhibits negative cosine similarity—termed "conflict gradients"—which can destabilize or retard learning. Consequences include color bleeding, semantic mismatches, and, in severe cases, mode collapse.

SGA is introduced as a solution by preserving only dominant, optimally aligned gradient flows while removing the conflict-inducing Q/K branch gradients. This targeted intervention directly addresses multi-branch gradient interference endemic to conventional attention mechanisms under adversarial and reconstruction-based training.

2. Mathematical Formulation and Mechanistic Detail

The standard dot-product attention mechanism for feature maps $X \in \mathbb{R}^{n \times d}$ and $Y \in \mathbb{R}^{n \times d}$ involves linear projections:

• $Q = X W_q$, $K = Y W_k$, $V = Y W_v$ for $W_q, W_k, W_v \in \mathbb{R}^{d \times d}$
• Attention map: $A = \text{softmax}(Q K^\top / \sqrt{d})$
• Output: $O = A V + X$

In SGA, a stop-gradient operator $sg(\cdot)$ (implements identity on forward pass, zero gradient on backward pass) is applied to the attention map:

• $A_\text{fwd} = \text{softmax}(Q K^\top / \sqrt{d})$
• $Y \in \mathbb{R}^{n \times d}$0
• Output: $Y \in \mathbb{R}^{n \times d}$1

Backward differentiation operates as follows:

• $Y \in \mathbb{R}^{n \times d}$2 receives gradient through $Y \in \mathbb{R}^{n \times d}$3
• $Y \in \mathbb{R}^{n \times d}$4 via the skip branch
• No gradients propagate through $Y \in \mathbb{R}^{n \times d}$5 and $Y \in \mathbb{R}^{n \times d}$6: $Y \in \mathbb{R}^{n \times d}$7

This setup ensures that only gradients empirically aligned with the true descent direction—those from the skip and V branches—are utilized for parameter updates.

3. Implementation and Practical Considerations

SGA is operationalized in practice by employing standard tensor libraries' stop-gradient facilities (e.g., torch.no_grad() or .detach() in PyTorch) during the attention score computation. Double normalization (row-wise softmax followed by column-wise normalization) replaces vanilla softmax to further stabilize the attention map against scaling variations in feature representations.

Feature embeddings are transformed with LeakyReLU activations, and SGA blocks are wrapped with BatchNorm and additional skip connections. Optionally, unused $Y \in \mathbb{R}^{n \times d}$8 and $Y \in \mathbb{R}^{n \times d}$9 layers may be pruned once their gradients are irrevocably blocked, marginally reducing computational footprint. Empirically, forward-pass expressivity remains unaffected after detaching these branches.

4. Network Architecture and Training Protocol

The SGA framework processes line-art sketches ($Q = X W_q$0) and warped references ($Q = X W_q$1) through respective encoders, yielding multi-scale feature maps ($Q = X W_q$2, $Q = X W_q$3), flattened to $Q = X W_q$4 for cross-modal aggregation. Two SGA block variants are alternated:

• cross-SGA: fuses $Q = X W_q$5 and $Q = X W_q$6 via stop-gradient attention,
• self-SGA: applies self-attention with gradient blocking to refine features.

This block arrangement is followed by a U-Net-style decoder with residual connections to reconstruct the colored output at $Q = X W_q$7 resolution. Single-headed attention is used, with double normalization for stability.

The overall loss combines:

1. $Q = X W_q$8 (reconstruction, $Q = X W_q$9; $K = Y W_k$0)
2. $K = Y W_k$1 (GAN least-squares; $K = Y W_k$2)
3. $K = Y W_k$3 (VGG19 perceptual; $K = Y W_k$4)
4. $K = Y W_k$5 (Gram matrix-based style loss; $K = Y W_k$6).

Training utilizes Adam optimizer ($K = Y W_k$7), with generator learning rate $K = Y W_k$8, discriminator $K = Y W_k$9, and a 40-epoch schedule.

5. Empirical Results and Comparative Analysis

Comprehensive quantitative analysis demonstrates that SGA delivers significant improvements over competing attention mechanisms. Compared to SCFT on the anime dataset:

• FID: SCFT = 44.65, SGA = 29.65 ($V = Y W_v$034%; $V = Y W_v$1 = 27.21%)
• SSIM: SCFT = 0.788, SGA = 0.912 ($V = Y W_v$216%; $V = Y W_v$3 = 15.70%)

On multi-domain benchmarks (anime, cat, dog, wildlife), SGA consistently achieves the lowest FID and highest SSIM in three of four datasets. Representative results include:

• Cat: FID = 34.35 (SGA) vs 36.33 (SCFT), SSIM = 0.843 (SGA) vs 0.636 (SCFT)
• Dog: FID = 54.76 vs 79.08, SSIM = 0.841 vs 0.683
• Wildlife: FID = 15.19 vs 24.93, SSIM = 0.831 vs 0.633

Ablation studies confirm that the primary gains originate from the stop-gradient operation, with further robustness conferred by double normalization and the inclusion of self-SGA.

6. Key Implementation Recommendations

Critical implementation details for robust deployment include:

• Strict application of stop-gradient facilities (torch.no_grad(), .detach()) to attention score computation to preempt any Q/K branch back-propagation.
• Adoption of double normalization (row and column-wise) to maintain attention stability across heterogeneous feature scales.
• Integration of LeakyReLU activations in feature embedding, with post-SGA BatchNorm and skip-connections.
• Optionally removing unused weight layers (e.g., $V = Y W_v$4, $V = Y W_v$5) post-detachment if parameters remain static.

These recommendations yield not just theoretical but empirically validated improvements in stability and output quality.

7. Impact and Significance in Line-Art Colorization

SGA constitutes an effective modification to dot-product attention, directly addressing the destabilizing effect of gradient conflict in reference-based line-art colorization tasks involving adversarial loss landscapes. By pruning the Q/K branch gradients, SGA stabilizes training (notably reducing mode collapse and parameter oscillation in L_rec), enhances perceptual fidelity, and improves outline preservation as measured by FID and SSIM (Li et al., 2022).

The empirical superiority and methodological simplicity of SGA suggest its broader applicability to attention-based architectures confronted with similar multi-branch gradient conflicts. The official implementation is accessible at https://github.com/kunkun0w0/SGA.