Unified Multi-Weather Image Restoration with Intra-Task Difficulty and Inter-Task Contribution

Abstract

Recent studies have witnessed significant advances in unified multi-weather image restoration, which aims to handle diverse weather degradations within a single model. In this work, we observe that rain, haze, and snow restoration exhibit substantial differences in both degradation characteristics and learning dynamics, making straightforward joint optimization prone to performance imbalance and ineffective knowledge transfer. To this end, we propose UMWIR-Net, a unified multi-weather image restoration network equipped with an Asymmetric Task Collaborative Learning strategy. ATCL consists of Intra-Task Difficulty Optimization and Inter-Task Contribution Scheduling. Specifically, Intra-Task Difficulty Optimization jointly models the remaining restoration error and recent learning progress to dynamically estimate the optimization difficulty of each weather task, thereby assigning larger weights to slowly converging and under-optimized tasks. Inter-Task Contribution Scheduling measures the directional influence of a source-task update on the validation objective of a target task, constructs an asymmetric task-contribution matrix, and accordingly promotes tasks that provide stronger transferable knowledge while compensating those that benefit less from collaborative learning. In this manner, different weather restoration tasks collaborate selectively and asymmetrically, allowing the model to exploit complementary knowledge across tasks and improve overall restoration performance. Furthermore, UMWIR-Net adopts a wavelet-based Transformer backbone to capture low- and high-frequency information, enabling effective modeling of both global structures and local details for diverse weather restoration. Extensive experiments on multi-weather image restoration datasets show that UMWIR-Net achieves state-of-the-art performance and delivers more balanced restoration quality across rain, haze, and snow removal.

1. Introduction

Images captured under adverse weather conditions often suffer from severe visibility degradation and structural corruption, which can substantially affect the reliability of visual perception systems in practical applications such as autonomous driving, intelligent transportation, and outdoor surveillance [1,2,3]. Different weather phenomena produce distinct visual degradation patterns. Rain streaks may occlude local structures and introduce directional high-frequency artifacts, haze reduces scene visibility through atmospheric scattering and global contrast attenuation, while snow particles often cause irregular occlusion and texture corruption. Therefore, effectively restoring images degraded by diverse weather conditions is important not only for improving visual quality, but also for supporting robust high-level vision systems.

Existing task-specific restoration methods have achieved remarkable progress in individual tasks, including image deraining [4,5,6], dehazing [7,8,9], and desnowing [10]. However, these methods are usually designed and trained for a particular degradation type. A model optimized for rain removal may fail to process hazy or snowy images, which requires multiple task-specific models to be separately stored, selected, and executed in practical systems. To overcome this limitation, unified multi-weather image restoration aims to handle multiple weather degradations within a single model. Recent unified restoration approaches [11,12,13,14] have demonstrated the potential of sharing restoration knowledge across degradation types and provide a more compact solution than maintaining separate models for individual tasks.

Most existing unified restoration methods, such as TransWeather [11], AirNet [15], PromptIR [16], DiffUIR [17], and AdaIR [18], mainly focus on improving degradation representation, feature adaptation, or network architecture, while multiple degradation tasks are jointly optimized within a shared restoration model. Although such a paradigm enables a single network to acquire multiple restoration capabilities, it implicitly assumes that different weather tasks can be learned at similar rates and benefit equally from joint training. In practice, rain, haze, and snow restoration exhibit substantially different degradation characteristics and learning dynamics. Some tasks can rapidly reach satisfactory performance, while others may converge slowly and remain under-optimized throughout joint training. Consequently, simple tasks can dominate the shared optimization process, whereas difficult tasks receive insufficient learning attention, resulting in imbalanced restoration quality across different weather conditions. Existing strategies based on fixed data proportions, equal task weights, or instantaneous training losses cannot reliably reflect such heterogeneous optimization states [19,20]. A large task loss may originate from its numerical scale rather than its actual learning difficulty, while a slowly decreasing loss does not necessarily indicate whether the task still has substantial room for improvement. Therefore, effective unified restoration requires a task-aware optimization mechanism that jointly considers the remaining restoration deficiency and the recent learning progress of each weather task.

In addition to unequal learning difficulty, previous multi-task learning studies have shown that knowledge transfer and task interactions can vary substantially across different task pairs [21,22,23]. In unified multi-weather restoration, we further observe that such knowledge transfer is inherently directional. Although different tasks share natural image statistics, clean-scene structures, edge priors, and texture regularities, the knowledge learned from one task does not necessarily benefit another task to the same extent. For example, global structure and color recovery learned from haze removal may facilitate background reconstruction in rainy and snowy images. In contrast, the local high-frequency suppression capability learned from rain or snow removal may provide less benefit to haze restoration. More importantly, optimizing one task may improve another task, whereas the reverse optimization direction may contribute weakly or even introduce negative interference. Such a relationship cannot be sufficiently characterized by symmetric task similarity or uniform feature sharing. Nevertheless, existing unified restoration methods typically do not explicitly model such directional cross-task interactions during joint optimization and generally optimize multiple task objectives together. As a result, useful transferable knowledge may not be adequately exploited, while conflicting task updates can impair the shared representation. This motivates us to explicitly characterize the directional contribution between weather tasks and schedule their collaboration according to their asymmetric relationships. These two observations, i.e., heterogeneous intra-task learning dynamics and asymmetric inter-task knowledge transfer, are summarized in Figure 1. In left figure of Figure 1, the convergence curves are obtained from a separate diagnostic run over a fixed 50-epoch training window solely to visualize the relative convergence behaviors of different weather restoration tasks; no early stopping is applied, and this diagnostic setting is independent of the training schedule of the final model.

Another challenge arises from the heterogeneous reconstruction requirements of different weather degradations. Haze mainly affects low-frequency appearance, including global illumination, contrast, and color distribution, and thus requires effective modeling of long-range scene structures. Rain streaks and snow particles, by contrast, usually introduce localized high-frequency corruption and require accurate reconstruction of edges, textures, and fine details. Although Transformer-based restoration models [24,25] provide strong global modeling capability, directly processing all frequency components in a homogeneous feature space may limit their ability to simultaneously recover global structures and local details under diverse weather conditions. A unified restoration backbone should therefore preserve complementary frequency information and provide suitable representations for both large-scale appearance degradation and fine-grained structural corruption.

To address these issues, we propose UMWIR-Net, a unified multi-weather image restoration network equipped with an Asymmetric Task Collaborative Learning strategy, termed ATCL. ATCL consists of two complementary components, namely Intra-Task Difficulty Optimization and Inter-Task Contribution Scheduling. Intra-Task Difficulty Optimization jointly models the remaining restoration error and recent learning progress to dynamically estimate the optimization difficulty of each weather task. It consequently assigns greater optimization weights to tasks that converge slowly and remain insufficiently learned, preventing easy tasks from dominating the joint training process. Inter-Task Contribution Scheduling explicitly measures the directional influence of a source-task update on the validation objective of a target task. Based on these directional effects, we construct an asymmetric task-contribution matrix to characterize the six transfer relationships among rain, haze, and snow restoration. The resulting scheduler promotes tasks that provide stronger transferable knowledge to the other tasks, while compensating tasks that receive limited benefit from collaborative learning. By jointly modeling intra-task difficulty and inter-task contribution, ATCL accounts for both the optimization state of each task and its directional collaborative relationships with the others, thereby generating adaptive task weights for selective and asymmetric joint optimization. In this manner, UMWIR-Net can exploit complementary knowledge across weather tasks, alleviate task imbalance and ineffective transfer, and improve the overall restoration performance. Furthermore, we adopt a wavelet-based Transformer backbone to accommodate the heterogeneous frequency characteristics of diverse weather degradations. The wavelet representation preserves complementary low- and high-frequency information, allowing the network to effectively model global structures affected by haze and local details corrupted by rain streaks or snow particles. The proposed ATCL is applied only during training and introduces neither additional model parameters nor computational overhead during inference. Extensive experiments on multi-weather image restoration datasets demonstrate that UMWIR-Net achieves state-of-the-art performance and delivers more balanced restoration quality across rain, haze, and snow removal.

The main contributions of this work are summarized as follows:

• We propose UMWIR-Net, a unified multi-weather image restoration network that jointly addresses the heterogeneous learning behaviors and reconstruction requirements of rain, haze, and snow restoration.
• We introduce an intra-task difficulty optimization that jointly considers the remaining restoration error and recent learning progress to dynamically identify slowly converging and under-optimized tasks, enabling balanced allocation of optimization resources.
• We develop an inter-task contribution scheduling that models directional knowledge transfer through an asymmetric task-contribution matrix, promoting tasks that provide transferable restoration knowledge while compensating tasks that receive limited collaborative support.
• Extensive experiments demonstrate that UMWIR-Net achieves state-of-the-art and more balanced performance across multiple weather restoration tasks.

2. Related Work

2.1. Single-Weather Image Restoration

Single-weather image restoration focuses on recovering clean visual content from images affected by a specific adverse weather condition [26,27,28,29]. Since rain, haze, and snow exhibit substantially different degradation mechanisms, task-specific methods usually incorporate dedicated architectures or priors to capture their characteristic distributions. Haze mainly causes depth-dependent contrast attenuation and color distortion, whereas rain streaks and snow particles generally introduce spatially non-uniform, high-frequency corruption. These differences have led to the development of specialized methods for image dehazing, deraining, and desnowing [30,31,32,33]. For image dehazing, early approaches commonly rely on atmospheric scattering models and handcrafted priors to estimate transmission maps and atmospheric light. Recent learning-based methods instead recover clear images through end-to-end feature modeling. DehazeFormer [34] introduces a Transformer-oriented architecture tailored to haze removal, enabling effective modeling of non-local degradation and global scene information. DCMPNet [7] further exploits the relationship between scene depth and haze distribution through collaborative depth estimation and image dehazing. These methods demonstrate the importance of global context and depth-related information for restoring contrast, color, and distant scene structures. However, these methods are specifically designed for image dehazing and do not address the joint restoration and collaborative optimization of heterogeneous weather degradations within a unified model. Image deraining methods mainly aim to separate rain-induced components from underlying scene details while avoiding the removal of genuine textures. PReNet [35] progressively refines intermediate restoration results through recurrent processing, while RCDNet [36] incorporates model-driven rain convolutional dictionaries to characterize rain formation. With the development of Transformer architectures, IDT [37] and DRSformer [6] improve the modeling of long-range dependencies and spatially varying rain patterns. RLP [38] introduces rain-related priors to enhance degradation representation, whereas NeRD-Rain [5] combines multi-scale Transformer modeling with implicit representations to capture both scale-specific and shared rain characteristics. MSDT [39] and URIR [40] further investigate effective feature interaction and unified rain representation for complex rainy scenes. Despite these advances, existing deraining methods remain specialized for rain removal and still face challenges in distinguishing rain streaks from fine background textures, particularly under spatially varying and nighttime conditions, while lacking the ability to exploit transferable restoration knowledge from other weather restoration tasks. Compared with rain streaks, snow particles exhibit greater variations in shape, scale, and opacity, and large particles may severely occlude scene content. SnowFormer [41] employs Transformer-based feature interaction to model multi-scale snow distributions and restore snow-covered regions. More recent studies also explore large-scale real-world data and perceptual supervision to improve the robustness of desnowing models under complex snowfall conditions [10]. These developments highlight the importance of reconstructing partially occluded structures, yet existing methods remain specialized for snow removal and do not leverage complementary knowledge from other weather restoration tasks.

Although single-weather methods achieve strong restoration performance by exploiting degradation-specific characteristics, each model is generally applicable to only one predefined weather condition. Practical systems must therefore maintain and select multiple independent models when encountering rain, haze, or snow. Moreover, separately trained networks cannot fully exploit the structural, textural, and natural-image knowledge that may be shared across different weather restoration tasks. These limitations motivate the development of unified frameworks capable of handling heterogeneous weather degradations within a single model.

2.2. Unified Multi-Weather Image Restoration

Unified multi-weather image restoration aims to recover images degraded by different adverse weather conditions within a single model. Compared with maintaining separate networks for rain, haze, and snow removal, this paradigm reduces model redundancy and facilitates the sharing of common restoration knowledge across tasks. TransWeather [11] introduces learnable weather-type queries into a Transformer-based architecture to adapt the decoding process to different degradations. Patil et al. [12] formulate multi-weather restoration from a domain-translation perspective, enabling a shared model to accommodate multiple weather domains. Although these approaches improve deployment efficiency, directly optimizing heterogeneous degradations within a common feature space may weaken the representation of task-specific characteristics and does not explicitly account for their different optimization states. Subsequent studies enhance degradation awareness through explicit degradation representations, prompt learning, and adaptive feature modulation. AirNet [15] employs contrastive learning to extract degradation representations that guide a unified restoration network. PromptIR [16] introduces degradation-specific prompts to dynamically regulate restoration features, while AdaIR [18] adapts a shared backbone to different restoration tasks through parameter-efficient conditioning. WGWS-Net [42] and Histoformer [43] further investigate weather-aware feature interaction, multi-scale representation learning, and distribution-oriented modeling. These methods improve the capability of a unified model to distinguish different degradation patterns, but their adaptation mechanisms mainly operate at the feature or parameter level without explicitly coordinating the heterogeneous optimization dynamics among different restoration tasks. More recent approaches explore generative and sequence-modeling paradigms to strengthen unified restoration. DiffUIR [17] leverages diffusion priors to recover images under multiple degradation types, whereas MambaIRv2 [44] employs state-space modeling to capture long-range dependencies with favorable computational efficiency. Other methods introduce expert-based architectures or specialized branches to balance shared and degradation-specific representations [13,14]. These developments substantially improve the representation capacity and flexibility of unified restoration networks. However, they mainly strengthen representation capacity and do not explicitly model dynamic task difficulty and directional cross-task contribution during joint optimization.

Nevertheless, most existing methods primarily focus on designing stronger backbones, degradation-aware prompts, or adaptive feature modules, while their training procedures commonly rely on mixed datasets, uniform sampling, or fixed task weights. Such strategies implicitly assume that different weather tasks exhibit comparable learning dynamics and benefit equally from joint optimization. In practice, rain, haze, and snow restoration may differ significantly in convergence speed, remaining restoration deficiency, and transferable knowledge. Moreover, the benefit transferred from one task to another is generally directional rather than symmetric. Consequently, difficult tasks may remain insufficiently optimized, while ineffective task interactions may impair the shared representation. Motivated by these observations, UMWIR-Net addresses these optimization-level limitations by dynamically estimating the learning difficulty of each weather task and its directional contribution to other tasks, thereby adaptively adjusting their collaboration during joint training.

2.3. Multi-Task Learning and Task Balancing

Multi-task learning jointly optimizes several related tasks through shared representations, allowing knowledge learned from one task to benefit others [21]. Conventional approaches usually combine task-specific losses using equal or manually defined weights. However, such strategies assume comparable loss scales, convergence rates, and optimization demands across tasks, and therefore may lead to imbalanced learning.

Adaptive weighting methods estimate task importance from training statistics. Uncertainty-based weighting adjusts task weights according to estimated noise levels [19], while Dynamic Weight Averaging and related methods emphasize tasks with slower loss reduction or larger improvable gaps [20]. Nevertheless, instantaneous loss values or recent decline rates alone may not reliably reflect whether a task remains under-optimized, because they are also affected by numerical scale and training stage. Gradient-based methods, such as MGDA, PCGrad, CAGrad, and Nash-MTL, mitigate negative transfer by searching for compatible update directions or resolving gradient conflicts [45,46,47,48]. However, gradient compatibility does not necessarily indicate the actual benefit that one task provides to another over the entire training process. Task-relation methods further analyze task affinity, grouping, and transferability. TAG estimates how optimizing one task influences another task and uses these relationships for task grouping [22]. Adaptive Task Balancing further considers inter-task contribution and intra-task difficulty for task weighting [23]. Although Ref. [23] also considers inter-task contribution and intra-task difficulty, these factors are estimated in a preparatory stage and then used for task weighting. In contrast, ATCL estimates task difficulty and directional contribution online from the current optimization state, allowing task relationships to evolve dynamically during training. However, many existing approaches rely on symmetric affinity or static task statistics, whereas knowledge transfer is often directional and task difficulty may change during training.

These limitations are particularly relevant to unified multi-weather restoration, where rain, haze, and snow exhibit different convergence behaviors and asymmetric transfer effects. Existing methods rarely consider both the dynamic optimization state of each task and its directional contribution to others. This motivates a task-balancing strategy that jointly models task-specific difficulty and asymmetric inter-task collaboration.

3. Proposed Method

We consider unified multi-weather image restoration involving rain removal, haze removal, and snow removal. The task set is denoted by

$$\mathcal{T}=\left\{{\mathcal{T}}_r,{\mathcal{T}}_h,{\mathcal{T}}_s\right\},$$

(1)

where the subscripts r, h, and s correspond to rain, haze, and snow, respectively, and ${\mathcal{T}}_i$ denotes a generic restoration task with $i\in \{r,h,s\}$. For each task ${\mathcal{T}}_i$, its corresponding paired training set is denoted by ${\mathcal{D}}_i$ and written as:

$${\mathcal{D}}_i={\left\{\left(x_i^n,y_i^n\right)\right\}}_{n=1}^{N_i},$$

(2)

where $x_i^n$ denotes a degraded observation and $y_i^n$ is the corresponding clean image. Our goal is to learn a unified restoration model $f_{\theta }$ that handles different weather degradations without requiring task-specific networks or degradation labels during inference. As illustrated in Figure 2, UMWIR-Net consists of a shared wavelet-based Transformer backbone and an Asymmetric Task Collaborative Learning strategy, termed ATCL. The restoration backbone captures complementary low- and high-frequency information to model both global structures and local details under diverse weather degradations. Specifically, the restoration backbone adopts a three-stage encoder–decoder architecture with a base feature dimension of 32. At each encoder stage, we employ the Haar discrete wavelet transform (DWT) to decompose the feature representation into one low-frequency subband (LL) and three high-frequency subbands (HL, LH, and HH). The LL component is fused with the corresponding multi-scale input feature and processed by Transformer blocks, whereas the three high-frequency components are adaptively fused and refined by high-frequency enhancement blocks. The numbers of low-frequency Transformer blocks across the three stages are set to $[1,2,4]$, while those of the high-frequency enhancement blocks are set to $[1,1,2]$. Both branches employ single-head attention, and the feed-forward expansion factor of the low-frequency Transformer branch is set to 2.0. During decoding, the low- and high-frequency features are progressively reconstructed through the corresponding inverse wavelet transform (IWT). We adopt the Haar wavelet due to its compact support, orthogonality, and low computational complexity, enabling efficient frequency decomposition and reconstruction while preserving complementary structural and detail information. ATCL coordinates how rain, haze, and snow restoration are jointly optimized within the shared network. For a mini-batch ${\mathcal{B}}_i$ sampled from task ${\mathcal{T}}_i$, the task-specific restoration risk is defined as

$${\mathcal{R}}_i\left(\theta ;{\mathcal{B}}_i\right)={\displaystyle \frac{1}{\left|{\mathcal{B}}_i\right|}}\sum _{\left(x,y\right)\in {\mathcal{B}}_i}\mathsf{\ell }\left(f_{\theta }\left(x\right),y\right),$$

(3)

where ${\mathcal{R}}_i(\theta ;{\mathcal{B}}_i)$ denotes the empirical restoration risk of task ${\mathcal{T}}_i$ over mini-batch ${\mathcal{B}}_i$ under the current model parameters $\theta$, and $\mathsf{\ell }(·,·)$ denotes the basic restoration loss. In our implementation, it is instantiated as:

$$\mathsf{\ell }\left(\widehat{y},y\right)=\sqrt{{∥\widehat{y}-y∥}_2^2+{ϵ}^2}+\xi \left[1-SSIM\left(\widehat{y},y\right)\right],$$

(4)

where the first term is the Charbonnier penalty and the second term promotes structural consistency.

A straightforward unified training paradigm minimizes the average restoration risk over all weather tasks. However, this symmetric formulation overlooks two important properties. First, rain, haze, and snow restoration exhibit substantially different learning dynamics. Some tasks converge rapidly, whereas others remain under-optimized during joint training. Second, knowledge transfer among weather tasks is inherently directional. Optimizing one task may benefit another task, while the reverse transfer can be considerably weaker or even detrimental. To address these issues, ATCL consists of two complementary components, namely intra-task difficulty optimization and inter-task contribution scheduling. Intra-task optimization estimates the optimization difficulty of each task according to its remaining restoration error and recent learning progress. Inter-task scheduling measures the directional influence of one task on another and constructs an asymmetric task-contribution matrix. The two components are jointly converted into adaptive task weights for selective and asymmetric optimization. Specifically, the intra-task difficulty estimated in Figure 2b and the inter-task contribution obtained in Figure 2c are jointly integrated to determine the adaptive task weights, which are subsequently applied to the joint optimization of the shared restoration network shown in Figure 2a.

3.1. Intra-Task Difficulty Optimization

Different weather degradations impose heterogeneous reconstruction demands. Haze mainly causes global contrast attenuation, color distortion, and visibility reduction, whereas rain streaks and snow particles usually introduce localized structural corruption and high-frequency interference. Consequently, their restoration objectives follow different convergence trajectories even when the same backbone and loss function are employed. Directly comparing task losses cannot reliably reflect such differences, because the absolute loss scales may vary across degradation types. A large loss does not necessarily indicate that a task is difficult to optimize, while a small loss does not guarantee that the task has been sufficiently learned. We therefore estimate intra-task difficulty by jointly considering the remaining restoration error and recent learning progress.

• Probe-based Task State Estimation. For each task ${\mathcal{T}}_i$, we construct a small and fixed probe set ${\mathcal{P}}_i$ that is excluded from parameter optimization. At training stage t, the probe restoration error is evaluated as

  $$v_i^t={\mathcal{R}}_i\left({\theta }^t;{\mathcal{P}}_i\right).$$

  (5)

The same restoration objective is used for both training and probe evaluation, ensuring that the estimated task state is directly related to the final restoration quality. Since the probe error can fluctuate due to data variation, we maintain an exponential moving average

$${\overline{v}}_i^t={\mu }_v{\overline{v}}_i^{t-1}+\left(1-{\mu }_v\right)v_i^t,$$

(6)

where ${\mu }_v\in [0,1)$ is the momentum coefficient. The smoothed error ${\overline{v}}_i^t$ provides a stable description of the current optimization state of task ${\mathcal{T}}_i$.

• Remaining Restoration Error. To eliminate the scale discrepancy among weather tasks, the current probe error is normalized using the error measured at the end of the equal-weight warm-up stage

  $$e_i^t={\displaystyle \frac{{\overline{v}}_i^t}{{\overline{v}}_i^w+ϵ}},$$

  (7)

  where w denotes the end of warm-up and $ϵ$ prevents numerical instability. The normalized residual $e_i^t$ measures how much restoration error remains relative to the initial optimization state of the same task. A larger value indicates that the task still has considerable room for improvement.

• Recent Learning Progress. We further characterize the short-term learning dynamics of task ${\mathcal{T}}_i$ over a temporal window of length $\Delta$. Its relative learning progress is defined as

  $$p_i^t={\displaystyle \frac{{\left[{\overline{v}}_i^{t-\Delta }-{\overline{v}}_i^t\right]}_{+}}{{\overline{v}}_i^{t-\Delta }+ϵ}},$$

  (8)

  where ${\left[z\right]}_{+}=max(0,z)$. A large $p_i^t$ indicates that the task is still improving rapidly, whereas a small value suggests slow convergence or an optimization plateau. When the probe error increases, the progress is set to zero because performance degradation should not be interpreted as effective learning.

• Optimization Difficulty Estimation. Based on the remaining restoration error and recent learning progress, the dynamic optimization difficulty of task ${\mathcal{T}}_i$ is formulated as

  $$d_i^t=e_i^{t}exp\left(-{\displaystyle \frac{p_i^t}{{\tau }_p}}\right),$$

  (9)

  where ${\tau }_p>0$ controls the sensitivity of the difficulty estimate to recent learning progress. The multiplicative formulation treats recent progress as a dynamic modulator of the remaining restoration error. When a task retains a large restoration error but exhibits limited recent improvement, $d_i^t$ becomes large, indicating that the task remains under-optimized. When its error is still relatively large but decreases rapidly, the progress term suppresses the difficulty score, preventing excessive optimization resources from being prematurely assigned to a naturally converging task. Therefore, $d_i^t$ captures the practical optimization difficulty of a weather task under the current shared model rather than merely reflecting its absolute loss magnitude.

3.2. Inter-Task Contribution Scheduling

Although rain, haze, and snow have different physical degradation mechanisms, their restoration processes share natural image statistics, clean-scene structures, edge priors, and texture regularities. Such shared knowledge enables one weather restoration task to benefit another during joint training. Nevertheless, the direction and strength of this transfer are generally different. For example, global structure and color recovery learned from haze removal may facilitate background reconstruction in rainy and snowy images. In contrast, local high-frequency suppression learned from rain or snow removal may provide limited benefit to haze restoration. Therefore, symmetric task similarity is insufficient for characterizing cross-task collaboration. We explicitly estimate the directional contribution from a source task ${\mathcal{T}}_i$ to a target task ${\mathcal{T}}_j$.

• Directional Cross-task Influence. Let ${\theta }_s\subseteq \theta$ denote the parameters shared by all weather tasks. Since UMWIR-Net adopts a fully shared restoration backbone, ${\theta }_s$ can include all trainable backbone parameters. At training stage t, the source-task training gradient is given by

  $$g_i^t={\nabla }_{{\theta }_s}{\mathcal{R}}_i\left({\theta }^t;{\mathcal{B}}_i\right),$$

  (10)

  while the probe gradient of target task ${\mathcal{T}}_j$ is computed as

  $$h_j^t={\nabla }_{{\theta }_s}{\mathcal{R}}_j\left({\theta }^t;{\mathcal{P}}_j\right).$$

  (11)

The training gradient describes the update direction induced by the source task, whereas the probe gradient represents the direction associated with the validation objective of the target task.

Consider a virtual update of the shared parameters using source task ${\mathcal{T}}_i$

$${\theta }_s^{\prime }={\theta }_s^t-\eta g_i^t,$$

(12)

where $\eta$ is a sufficiently small virtual step size. Using a first-order Taylor expansion around ${\theta }_s^t$, the resulting change in the target probe objective is approximated as

$${\mathcal{R}}_j\left({\theta }_s^{\prime };{\mathcal{P}}_j\right)-{\mathcal{R}}_j\left({\theta }_s^t;{\mathcal{P}}_j\right)\approx -\eta 〈g_i^t,h_j^t〉.$$

(13)

A positive inner product indicates that the source-task update is aligned with the direction that reduces the target-task probe error. Task ${\mathcal{T}}_i$ therefore provides a positive contribution to task ${\mathcal{T}}_j$. A negative inner product instead reveals conflicting optimization directions and potential negative transfer.

• Scale-normalized Task Contribution. The raw gradient inner product is sensitive to gradient magnitude. We therefore compute a scale-invariant cross-task alignment

  $$s_{i\to j}^t={\displaystyle \frac{〈g_i^t,h_j^t〉}{{∥g_i^t∥}_2{∥h_j^t∥}_2+ϵ}}.$$

  (14)

The value $s_{i\to j}^t$ measures whether the optimization direction of source task ${\mathcal{T}}_i$ is consistent with the validation objective of target task ${\mathcal{T}}_j$. However, self-optimization effectiveness can vary substantially across target tasks. The cross-task alignment should therefore be interpreted relative to the effect of the target task learning from its own training data. We calculate the self-alignment of task ${\mathcal{T}}_j$ as

$$s_{j\to j}^t={\displaystyle \frac{〈g_j^t,h_j^t〉}{{∥g_j^t∥}_2{∥h_j^t∥}_2+ϵ}}.$$

(15)

The normalized directional contribution is then defined as

$$c_{i\to j}^t=tanh\left({\displaystyle \frac{s_{i\to j}^t}{{\tau }_c\left(\left|s_{j\to j}^t\right|+ϵ\right)}}\right),i\ne j,$$

(16)

where ${\tau }_c>0$ controls the sensitivity of contribution estimation. The hyperbolic tangent bounds the contribution within $(-1,1)$ and suppresses extreme gradient observations.

A positive $c_{i\to j}^t$ indicates that learning source task ${\mathcal{T}}_i$ provides useful transferable knowledge to target task ${\mathcal{T}}_j$. A value close to zero represents weak interaction, whereas a negative value indicates harmful interference. Since the reverse transfer is evaluated using different source gradients, target probe gradients, and normalization references, the task relationship generally satisfies

$$c_{i\to j}^t\ne c_{j\to i}^t.$$

(17)

• Asymmetric Task-contribution Matrix. For rain, haze, and snow restoration, the directional task contributions form the matrix

  $${\mathbf{C}}^t=\left[\begin{array}{ccc}0& c_{r\to h}^t& c_{r\to s}^t\\ c_{h\to r}^t& 0& c_{h\to s}^t\\ c_{s\to r}^t& c_{s\to h}^t& 0\end{array}\right].$$

  (18)

The off-diagonal entries represent the six directional transfer relationships among rain, haze, and snow restoration. Since an estimate obtained from a single pair of mini-batches may be noisy, we maintain an exponential moving average

$${\overline{c}}_{i\to j}^t={\mu }_c{\overline{c}}_{i\to j}^{t-1}+\left(1-{\mu }_c\right)c_{i\to j}^t,$$

(19)

where ${\mu }_c$ is the contribution momentum coefficient. The contribution matrix is updated every K iterations rather than at every optimization step, thereby limiting the additional gradient computation.

• Contribution-aware Task Scheduling. The task-contribution matrix describes not only how much knowledge a task provides to the others, but also how much benefit it receives from collaborative learning. We summarize these two directional properties using outgoing contribution and received benefit. The outgoing contribution of task ${\mathcal{T}}_i$ is defined as

  $$o_i^t={\displaystyle \frac{1}{N-1}}\sum _{j\ne i}{\overline{c}}_{i\to j}^t.$$

  (20)

A large $o_i^t$ indicates that optimizing task ${\mathcal{T}}_i$ tends to improve the other weather tasks. Equation (20) retains both positive and negative values. Consequently, a task that frequently interferes with other tasks receives a lower outgoing-contribution score.

The benefit received by task ${\mathcal{T}}_i$ from the other tasks is defined as

$$r_i^t={\displaystyle \frac{1}{N-1}}\sum _{j\ne i}{\left[{\overline{c}}_{j\to i}^t\right]}_{+}.$$

(21)

A large $r_i^t$ indicates that task ${\mathcal{T}}_i$ receives substantial positive transfer from collaborative learning. A small value means that it benefits less from the other weather tasks and therefore requires additional optimization compensation. Both $o_i^t$ and $r_i^t$ are derived from the same asymmetric task-contribution matrix and jointly characterize the inter-task contribution scheduling requirement.

3.3. Asymmetric Collaborative Weighting

ATCL jointly models intra-task difficulty and inter-task contribution. Intra-task optimization determines whether an individual task remains difficult and under-optimized, while inter-task contribution scheduling determines whether the task provides transferable knowledge to the others or receives limited benefit from them. These two components are integrated to generate adaptive task weights for selective and asymmetric joint optimization.

• Inter-task Contribution Scheduling Score. Since outgoing contribution $o_i^t$ and received benefit $r_i^t$ are measured on different numerical ranges, we first standardize them across the task dimension. More generally, for a task statistic ${\mathbf{z}}^t$, we define

  $${\widehat{z}}_i^t={\displaystyle \frac{z_i^t-mean\left({\mathbf{z}}^t\right)}{std\left({\mathbf{z}}^t\right)+ϵ}}.$$

  (22)

The standardization preserves the relative ordering among tasks while preventing one statistic from dominating merely because of its scale.

The inter-task contribution scheduling score of task ${\mathcal{T}}_i$ is formulated as

$${\varphi }_i^t=\alpha {\widehat{o}}_i^t-\beta {\widehat{r}}_i^t,$$

(23)

where $\alpha$ and $\beta$ control the influence of outgoing contribution and received benefit, respectively. A task with a large outgoing contribution receives a larger scheduling score because its optimization can improve the shared restoration capability. A task with a small received benefit is also compensated because it cannot sufficiently rely on knowledge transferred from the other tasks.

• Asymmetric Collaboration Priority. We standardize the intra-task difficulty ${\mathbf{d}}^t$ using Equation (22) and combine it with the inter-task contribution scheduling score

  $$a_i^t={\widehat{d}}_i^t+{\varphi }_i^t.$$

  (24)

Equation (24) explicitly preserves the two-component design of ATCL. The first term originates from intra-task difficulty optimization and emphasizes tasks that converge slowly and remain under-optimized. The second term originates from inter-task contribution scheduling and coordinates each task according to its directional contribution relationship with the others. Thus, the collaboration priority simultaneously reflects the optimization state of an individual task and its role in cross-task knowledge transfer.

The raw task weight is obtained through temperature-controlled normalization

$$q_i^t=N{\displaystyle \frac{exp\left(a_i^t/\tau \right)}{{\sum }_{k=1}^{N}exp\left(a_k^t/\tau \right)}},$$

(25)

where $\tau$ controls the degree of asymmetry among weather tasks. A smaller temperature produces a sharper weight distribution, whereas a larger value gradually approaches equal-weight training. The average value of the raw task weights is one, which maintains a stable global optimization scale.

Short-term fluctuations in the probe statistics may occasionally produce excessively large or small task weights. We therefore apply bounded projection

$${\tilde{q}}_i^t=clip\left(q_i^t,{\lambda }_{min},{\lambda }_{max}\right),$$

(26)

where ${\lambda }_{min}$ and ${\lambda }_{max}$ prevent task starvation and over-dominance, respectively. The projected weights are subsequently renormalized as

$${\lambda }_i^t=N{\displaystyle \frac{{\tilde{q}}_i^t}{{\sum }_{k=1}^N{\tilde{q}}_k^t}}.$$

(27)

• Joint Training and Inference. During each optimization step, equally sized mini-batches are sampled from the rain, haze, and snow restoration datasets. This sampling strategy avoids implicit task bias caused by different dataset sizes. The resulting ATCL objective is formulated as

  $${\mathcal{L}}_{\mathrm{ATCL}}^t={\displaystyle \frac{{\sum }_{i=1}^N{\lambda }_i^t{\mathcal{R}}_i\left(\theta ;{\mathcal{B}}_i\right)}{{\sum }_{i=1}^N{\lambda }_i^t}}.$$

  (28)

The adaptive weights are treated as detached training statistics and do not participate in gradient back-propagation. All weather tasks share the same restoration parameters, while their optimization strengths are dynamically adjusted according to their current learning states and directional collaborative relationships. Training begins with an equal-weight warm-up stage, during which ${\lambda }_i^t=1$ for all tasks. This stage allows the shared model to acquire basic restoration capability and provides reliable initial probe statistics. After warm-up, the intra-task difficulty and task-contribution matrix are updated periodically. The task weights remain fixed between two consecutive updates, which stabilizes joint optimization and limits additional computation. The probe-based difficulty estimation and gradient-based inter-task scheduling are activated only during training. During inference, UMWIR-Net performs a standard forward pass without probe sets, task labels, auxiliary branches, or dynamic weight estimation. Therefore, ATCL introduces neither additional model parameters nor computational overhead during inference. By jointly considering intra-task difficulty and inter-task contribution, ATCL enables different weather restoration tasks to collaborate selectively and asymmetrically, thereby improving overall restoration performance while alleviating task imbalance and ineffective knowledge transfer.

4. Experiments

4.1. Experiments Setup

4.1.1. Datasets and Metrics

To comprehensively evaluate the performance of our method on unified multi-weather image restoration, we conduct experiments on three datasets: WeatherBench [49], RainRAG [50], and FoundIR-Weather [51]. WeatherBench is a real-world dataset covering both daytime and nighttime scenes under three representative adverse weather conditions, namely rain, snow, and haze. It contains a total of 42,002 images. For the WeatherBench experiments, the rain, haze, and snow training subsets are jointly used to train a unified UMWIR-Net, which is subsequently evaluated on the corresponding WeatherBench test subsets. The training set of RainRAG consists of 52,869 paired images, while its test set contains 400 images divided into four representative subsets: daytime rain streaks (DRS), daytime raindrops (DRD), nighttime rain streaks (NRS), and nighttime raindrops (NRD). For evaluation on RainRAG, we separately retrain UMWIR-Net using its 52,869 training images and evaluate the resulting model on the 400 test images. FoundIR-Weather is a large-scale, high-quality real-world dataset covering a wide variety of scenes and weather conditions. We employ this dataset for cross-dataset generalization testing to evaluate the robustness and generalization capability of our model in diverse real-world scenarios.

We use PSNR [52] and SSIM [53] to evaluate restoration performance. PSNR measures pixel-level reconstruction accuracy, while SSIM assesses structural similarity in terms of luminance, contrast, and structure. Higher PSNR and SSIM values indicate better performance.

4.1.2. Comparison Methods

To provide a comprehensive evaluation, we compare our method with 22 representative image restoration approaches covering different task settings and architectural paradigms. For single-degradation restoration, we include DehazeFormer [34] and DCMPNet [7] for image dehazing; PReNet [35], RCDNet [36], IDT [37], DRSformer [6], RLP [38], MSDT [39], NeRD-Rain [5], and URIR [40] for image deraining; and SnowFormer [41] for image desnowing. We further compare with the general-purpose restoration backbones MPRNet [54], Restormer [24], and MambaIRv2 [44], as well as representative all-in-one approaches, including AirNet [15], TransWeather [11], PromptIR [16], WGWS-Net [42], DiffUIR [17], Histoformer [43], and AdaIR [18]. These baselines encompass CNN-based recurrent, progressive, and model-driven networks, Transformer-based architectures, diffusion models, and Mamba-based state-space models, enabling a comprehensive comparison across both task scope and network design. For a fair comparison, all competing methods are retrained using their officially released implementations under the corresponding experimental settings.

4.1.3. Implementation Details

All experiments are implemented using PyTorch 2.0.0 and conducted on a single NVIDIA GeForce RTX 3090 GPU. During training, the input and ground-truth images are cropped into $512\times 512$ patches, with a batch size of 4. The proposed network is optimized using AdamW with ${\beta }_1=0.9$, ${\beta }_2=0.99$, and a weight decay of $1\times {10}^{-3}$. The initial learning rate is set to $5\times {10}^{-4}$ and progressively reduced to $1\times {10}^{-7}$ using a cosine annealing restart scheduler. The entire training process contains 300,000 iterations. Validation is performed every 5000 iterations, and the checkpoint achieving the highest PSNR is selected for evaluation.

4.2. Evaluation on the WeatherBench Dataset

Table 1. Quantitative comparison on the WeatherBench dataset. The best and second-best values are bold and underlined, respectively. ↑ indicates that higher values correspond to better performance.

Methods	Dehaze		Derain		Desnow		Average
	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
DehazeFormer [34]	24.12	0.7447	36.05	0.9539	28.88	0.8492	29.68	0.8493
DCMPNet [7]	21.18	0.5056	32.04	0.8756	24.81	0.6135	26.01	0.6649
DRSformer [6]	19.95	0.6944	33.98	0.9426	28.00	0.8358	27.31	0.8243
NeRD-Rain [5]	21.52	0.7184	35.74	0.9503	28.87	0.8507	28.71	0.8398
SnowFormer [41]	22.71	0.7362	35.18	0.9514	29.30	0.8678	29.06	0.8518
MPRNet [54]	23.27	0.7393	36.14	0.9537	29.18	0.8598	29.53	0.8509
Restormer [24]	19.30	0.6866	34.49	0.9448	27.95	0.8360	27.25	0.8225
AirNet [15]	20.94	0.7054	33.59	0.9418	22.06	0.7799	25.53	0.8090
TransWeather [11]	19.79	0.6800	29.34	0.9034	24.96	0.7959	24.70	0.7931
PromptIR [16]	21.11	0.7128	34.54	0.9443	27.93	0.8363	27.86	0.8311
WGWS-Net [42]	13.79	0.6027	37.08	0.9608	20.81	0.7799	23.89	0.7811
DiffUIR [17]	22.74	0.7439	35.93	0.9547	29.50	0.8696	29.39	0.8561
Histoformer [43]	17.69	0.6691	30.70	0.9160	25.39	0.8076	24.59	0.7976
AdaIR [18]	23.08	0.7311	34.87	0.9458	28.44	0.8372	28.80	0.8380
Ours	25.09	0.7605	37.17	0.9583	29.90	0.8643	30.72	0.8610

reports the quantitative results on WeatherBench for dehazing, deraining, and desnowing. Our method achieves the best average performance, reaching 30.72 dB PSNR and 0.8610 SSIM. Compared with the second-best results, it provides improvements of 1.04 dB in PSNR over DehazeFormer and 0.0049 in SSIM over DiffUIR. For dehazing, UMWIR-Net obtains the best PSNR and SSIM values of 25.09 dB and 0.7605, outperforming DehazeFormer by 0.97 dB and 0.0158, respectively. For deraining, our method achieves the highest PSNR of 37.17 dB and the second-best SSIM of 0.9583, demonstrating effective rain removal and structural preservation. On desnowing, it obtains the highest PSNR of 29.90 dB and a competitive SSIM of 0.8643. Although several task-specific methods perform well on individual degradations, their performance varies considerably across other weather conditions. In contrast, UMWIR-Net maintains consistently strong results across all three tasks, validating the effectiveness of ATCL in balancing heterogeneous optimization requirements.

The qualitative comparisons are presented in Figure 3, Figure 4 and Figure 5. For hazy images, several competing methods leave residual haze or introduce noticeable color shifts, whereas our method restores clearer scene structures and more natural contrast. In rainy scenes, UMWIR-Net removes dense rain streaks while preserving edges and fine background textures, with fewer residual artifacts and less over-smoothing. For snowy images, our method effectively suppresses snow particles and reconstructs partially occluded regions with improved structural consistency. Overall, the quantitative and qualitative results demonstrate that UMWIR-Net provides robust and balanced restoration across diverse weather degradations.

4.3. Evaluation on the RainRAG Dataset

Table 2. Quantitative evaluations on the RainRAG dataset, where DRS, DRD, NRS, and NRD denote daytime rain streaks, daytime raindrops, nighttime rain streaks, and nighttime raindrops, respectively. ↑ indicates that higher values correspond to better performance.

Methods	RainRAG Dataset
	DRS		DRD		NRS		NRD		Average
	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
PReNet [35]	26.73	0.8233	21.45	0.7358	32.85	0.9520	24.29	0.8058	26.33	0.8292
RCDNet [36]	25.27	0.7949	21.46	0.7295	28.44	0.9070	20.62	0.5832	23.95	0.7537
MPRNet [54]	27.99	0.8242	22.50	0.7264	33.75	0.9541	24.46	0.8118	27.17	0.8291
Restormer [24]	28.45	0.8298	23.36	0.7547	33.92	0.9541	25.85	0.8234	27.89	0.8405
IDT [37]	26.97	0.8279	23.17	0.7648	33.05	0.9470	25.20	0.7963	27.10	0.8340
DRSformer [6]	28.08	0.8253	22.93	0.7499	32.86	0.9457	24.99	0.7953	27.22	0.8291
RLP [38]	28.01	0.8290	22.42	0.7353	32.06	0.9438	23.48	0.7638	26.49	0.8180
MSDT [39]	28.60	0.8414	23.31	0.7487	34.56	0.9591	25.28	0.8148	27.94	0.8410
NeRD-Rain [5]	28.11	0.8309	23.30	0.7513	33.88	0.9511	25.31	0.8027	27.65	0.8340
URIR [40]	28.29	0.8385	23.19	0.7483	34.32	0.9570	25.82	0.8263	27.91	0.8425
Ours	30.38	0.8386	23.60	0.7652	37.29	0.9686	24.06	0.8193	28.83	0.8479

Note: The best and second-best values are bold and underlined, respectively.

reports the quantitative results on the RainRAG dataset, which contains four representative rainy conditions: daytime rain streaks (DRS), daytime raindrops (DRD), nighttime rain streaks (NRS), and nighttime raindrops (NRD). These subsets exhibit substantial differences in illumination, rain appearance, and degradation severity, providing a comprehensive evaluation of deraining robustness. Our method achieves PSNR/SSIM values of 30.38/0.8386, 23.60/0.7652, 37.29/0.9686, and 24.06/0.8193 on DRS, DRD, NRS, and NRD, respectively. It obtains the best PSNR on DRS, outperforming the second-best MSDT by 1.78 dB, while achieving the second-best SSIM. On DRD, UMWIR-Net ranks first in both metrics, surpassing Restormer by 0.24 dB in PSNR and IDT by 0.0004 in SSIM. The advantage is more pronounced on NRS, where our method exceeds the second-best result by 2.73 dB in PSNR and 0.0095 in SSIM, demonstrating strong robustness under challenging nighttime illumination. Although the PSNR on NRD remains lower than several competing methods, our method achieves the third-best SSIM, indicating favorable structural preservation. Overall, UMWIR-Net reaches the highest average performance of 28.83 dB PSNR and 0.8479 SSIM, improving upon the second-best results by 0.89 dB and 0.0054, respectively. These results demonstrate its strong adaptability to different rain patterns and illumination conditions.

4.4. Evaluation on the FoundIR-Weather Dataset

For cross-dataset generalization evaluation, we directly apply the UMWIR-Net model trained on WeatherBench to the rain and haze subsets of FoundIR-Weather without any fine-tuning. As reported in

Table 3. Quantitative evaluations on the FoundIR-Weather [51] dataset. ↑ indicates that higher values correspond to better performance.

Datasets	Metrics	Restormer [24]	TransWeather [11]	DRSformer [6]	NeRD-Rain [5]	MambaIRv2 [44]	Ours
Rain	PSNR ↑	22.87	22.94	23.70	24.11	23.55	26.42
	SSIM ↑	0.7195	0.6814	0.7026	0.7169	0.7234	0.8060
Haze	PSNR ↑	12.88	17.43	15.01	17.19	14.33	17.67
	SSIM ↑	0.4096	0.6545	0.4787	0.6561	0.5036	0.6262
Average	PSNR ↑	17.88	20.19	19.36	20.65	18.94	22.05
	SSIM ↑	0.5646	0.6680	0.5907	0.6865	0.6135	0.7161

Note: The best and second-best values are bold and underlined, respectively.

, our method achieves 26.42 dB PSNR and 0.8060 SSIM on the rain subset, outperforming the second-best results by 2.31 dB and 0.0826, respectively. These substantial improvements demonstrate its effectiveness in removing spatially varying rain while preserving structural details. On the haze subset, UMWIR-Net obtains the highest PSNR of 17.67 dB, exceeding TransWeather by 0.24 dB. Although its SSIM of 0.6262 is lower than those of NeRD-Rain and TransWeather, it remains competitive in structural reconstruction. Averaged over the two degradation types, our method achieves the best PSNR/SSIM of 22.05 dB/0.7161, surpassing the second-best NeRD-Rain by 1.40 dB and 0.0296. These results confirm that UMWIR-Net generalizes effectively to real-world weather degradations and provides strong overall restoration performance across heterogeneous rain and haze conditions.

Visual comparisons on representative rainy and hazy scenes are provided in Figure 6. For rainy images, some competing methods leave visible rain streaks or remove fine background textures together with the degradation. Our method suppresses rain artifacts more thoroughly while preserving object boundaries and local structural details. For hazy scenes, existing methods may retain residual haze or produce unnatural color and contrast variations, particularly in distant regions. UMWIR-Net restores clearer scene structures, more natural colors, and improved global visibility with fewer artifacts. The enlarged patches further demonstrate its advantage in preserving fine textures and edge information, confirming its robustness under diverse real-world weather conditions.

4.5. Ablation Studies

4.5.1. Effectiveness of the Main Components

Table 4. Ablation study of the main components on WeatherBench, where Wavelet denotes the wavelet-based Transformer backbone, ITDO denotes intra-task difficulty optimization, and ITCS denotes inter-task contribution scheduling. ↑ indicates that higher values correspond to better performance, while ✔ and ✗ indicate that the corresponding component is enabled and disabled, respectively.

Variant	Wavelet	ITDO	ITCS	Dehaze		Derain		Desnow		Average
				PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
(a)	✗	✗	✗	22.26	0.7214	34.83	0.9488	28.48	0.8358	28.52	0.8353
(b)	✔	✗	✗	24.12	0.7497	36.39	0.9545	29.09	0.8429	29.87	0.8490
(c)	✔	✔	✗	24.88	0.7569	36.95	0.9575	29.72	0.8615	30.52	0.8586
(d)	✔	✗	✔	24.84	0.7573	36.83	0.9569	29.49	0.8585	30.39	0.8576
Ours	✔	✔	✔	25.09	0.7605	37.17	0.9583	29.90	0.8643	30.72	0.8610

evaluates the contributions of the main components. Wavelet denotes the wavelet-based Transformer backbone, while ITDO and ITCS refer to Intra-Task Difficulty Optimization and Inter-Task Contribution Scheduling, respectively. Compared with the spatial-domain baseline, introducing the wavelet backbone consistently improves dehazing, deraining, and desnowing, demonstrating the benefit of jointly modeling low- and high-frequency information for heterogeneous weather degradations. Based on the wavelet backbone, ITDO improves the average performance to 30.52 dB/0.8586 by assigning greater optimization emphasis to slowly converging and under-optimized tasks. ITCS also provides clear gains, reaching 30.39 dB/0.8576, which confirms that directional task-contribution modeling facilitates more effective cross-task knowledge transfer. Although the two modules yield different improvements across individual tasks, both outperform the wavelet-only variant. Combining all three components achieves the best performance on every restoration task, with an average PSNR/SSIM of 30.72 dB/0.8610. The complete model surpasses the ITDO-only and ITCS-only variants by 0.20 dB and 0.33 dB in average PSNR, respectively. These results demonstrate that frequency-aware representation, intra-task difficulty optimization, and inter-task contribution scheduling are complementary and jointly improve restoration accuracy and task balance.

4.5.2. Comparison with Multi-Task Optimization Strategies

To further evaluate the effectiveness of ATCL, we conduct a controlled comparison with representative multi-task optimization strategies, including Equal Weighting (Variant b), DWA [20], GradNorm [55], and PCGrad [56]. All methods employ the same wavelet-based Transformer backbone and identical training settings, with only the multi-task optimization strategy being changed. As shown in

Table 5. Comparison of multi-task optimization strategies on WeatherBench using the same backbone and training settings. Average PSNR and SSIM are reported over dehazing, deraining, and desnowing. ↑ indicates that higher values correspond to better performance.

Variant	Equal Weighting	DWA [20]	GradNorm [55]	PCGrad [56]	ATCL (Ours)
PSNR ↑	29.87	30.05	30.57	30.31	30.72
SSIM ↑	0.8490	0.8532	0.8575	0.8554	0.8610

, Equal Weighting achieves an average PSNR/SSIM of 29.87 dB/0.8490, while DWA, GradNorm, and PCGrad obtain 30.05 dB/0.8532, 30.57 dB/0.8575, and 30.31 dB/0.8554, respectively. ATCL achieves the best performance of 30.72 dB/0.8610, outperforming the strongest competing strategy, GradNorm, by 0.15 dB in PSNR and 0.0035 in SSIM. These results demonstrate that ATCL provides more effective task coordination than representative task-weighting and gradient-balancing strategies under the same restoration backbone and training configuration.

4.5.3. Analysis of Intra-Task Difficulty Optimization

Table 6. Analysis of intra-task difficulty optimization on WeatherBench. Inter-task contribution scheduling is enabled for all variants, while remaining restoration error and recent learning progress denote the two components used for dynamic optimization difficulty estimation. ↑ indicates that higher values correspond to better performance, while ✔ and ✗ indicate that the corresponding component is enabled and disabled, respectively.

Variant	Remaining Restoration Error	Recent Learning Progress	Dehaze		Derain		Desnow		Average
			PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
(d)	✗	✗	24.84	0.7573	36.83	0.9569	29.49	0.8585	30.39	0.8576
(e)	✔	✗	24.95	0.7571	37.04	0.9576	29.73	0.8614	30.57	0.8587
(f)	✗	✔	25.03	0.7595	37.09	0.9581	29.88	0.8630	30.67	0.8602
Ours	✔	✔	25.09	0.7605	37.17	0.9583	29.90	0.8643	30.72	0.8610

analyzes the two components of Intra-Task Difficulty Optimization, while Inter-Task Contribution Scheduling remains enabled for all variants. Without ITDO, variant (d) obtains an average PSNR/SSIM of 30.39 dB/0.8576. Introducing only the remaining restoration error in variant (e) improves the result to 30.57 dB/0.8587, showing that emphasizing tasks with larger unresolved errors helps alleviate under-optimization. Using only recent learning progress, variant (f) further achieves 30.67 dB/0.8602 and performs better than variant (e) on most metrics. This indicates that recent convergence behavior provides a more responsive estimate of the current task state. However, progress alone cannot determine whether a slowly changing task has already converged or still retains substantial restoration errors. Combining both components produces the best average performance of 30.72 dB/0.8610 and yields the strongest results across all three weather tasks. These results confirm that remaining restoration error and recent learning progress provide complementary information, enabling more reliable difficulty estimation and more balanced multi-task optimization.

4.5.4. Analysis of Inter-Task Contribution Scheduling

Table 7. Analysis of inter-task contribution scheduling on WeatherBench. Intra-task difficulty optimization is enabled for all variants. Outgoing contribution and received benefit denote the two task-level statistics derived from the task-contribution matrix. ↑ indicates that higher values correspond to better performance, while ✔ and ✗ indicate that the corresponding component is enabled and disabled, respectively.

Variant	Outgoing Contribution	Received Benefit	Asymmetric	Dehaze		Derain		Desnow		Average
				PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
(c)	✗	✗	✗	24.88	0.7569	36.95	0.9575	29.72	0.8615	30.52	0.8586
(g)	✔	✗	✔	25.03	0.7600	37.13	0.9581	29.88	0.8640	30.68	0.8607
(h)	✗	✔	✔	24.92	0.7593	37.10	0.9578	29.85	0.8632	30.62	0.8601
(i)	✔	✔	✗	25.06	0.7603	37.15	0.9582	29.90	0.8642	30.70	0.8609
Ours	✔	✔	✔	25.09	0.7605	37.17	0.9583	29.90	0.8643	30.72	0.8610

analyzes the components of Inter-Task Contribution Scheduling, with Intra-Task Difficulty Optimization enabled for all variants. Without contribution scheduling, variant (c) achieves an average PSNR/SSIM of 30.52 dB/0.8586. Introducing only outgoing contribution in variant (g) improves the result to 30.68 dB/0.8607, indicating that emphasizing tasks that provide useful knowledge to others enhances shared optimization. Variant (h), which uses only received benefit, also improves the average performance to 30.62 dB/0.8601, confirming the effectiveness of compensating tasks that obtain limited support from collaborative learning. When both statistics are used with symmetric task relationships, variant (i) reaches 30.70 dB/0.8609. The complete asymmetric formulation further improves the result to 30.72 dB/0.8610 and achieves the best or comparable performance across all three weather tasks. Compared with variant (c), the full ITCS provides gains of 0.20 dB and 0.0024 in average PSNR and SSIM, respectively. These results demonstrate that outgoing contribution and received benefit are complementary, while preserving transfer direction provides a modest advantage over symmetric interaction modeling. To further assess the robustness of this improvement, we repeat both the symmetric and asymmetric variants using three random seeds under identical training settings. As shown in

Table 8. Repeated comparison of symmetric and asymmetric task collaboration on WeatherBench using three random seeds. ↑ indicates that higher values correspond to better performance.

Strategy	Seed 0		Seed 1		Seed 2		Mean ± Std.
	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑	PSNR ↑	SSIM ↑
Symmetric	30.70	0.8609	30.80	0.8615	30.95	0.8627	30.82 ± 0.10	0.8617 ± 0.0007
Asymmetric (Ours)	30.72	0.8610	30.88	0.8620	30.98	0.8635	30.86 ± 0.11	0.8622 ± 0.0010

Note: The best values are bold.

, the asymmetric formulation consistently outperforms the symmetric variant across all three seeds. Averaged over the three runs, the asymmetric strategy achieves 30.86 ± 0.11 dB/0.8622 ± 0.0010, compared with 30.82 ± 0.10 dB/0.8617 ± 0.0007 for the symmetric strategy. The consistent improvements across different random seeds support that the observed advantage is not tied to a particular random initialization.

5. Conclusions

We have presented UMWIR-Net, a unified multi-weather image restoration network for rain, haze, and snow removal. To address task imbalance and ineffective knowledge transfer, we introduce an asymmetric task collaborative learning strategy consisting of intra-task difficulty optimization and inter-task contribution scheduling. The former emphasizes slowly converging and under-optimized tasks, while the latter models directional task contributions to enable selective and asymmetric collaboration. In addition, a wavelet-based Transformer backbone captures complementary low- and high-frequency information for effective restoration of global structures and local details. Extensive experiments demonstrate that UMWIR-Net achieves state-of-the-art and more balanced performance across multiple weather restoration tasks.