SoftAdapt: Adaptive Loss Weighting

• SoftAdapt is an adaptive technique that dynamically weights multi-part loss functions using recent loss trends to improve optimization stability.
• It employs a softmax-based rule with a focus parameter to adjust weights in real time, mitigating imbalances caused by fixed or heuristic settings.
• Empirical studies demonstrate its effectiveness in PINNs, graph matching, and deep learning, achieving faster convergence and higher accuracy than traditional methods.

SoftAdapt is a family of adaptive techniques for weighting multi-part loss functions in neural network optimization. It dynamically adjusts the scalar weights assigned to each loss component based on recent performance trends, aiming to redress imbalances in composite objectives where fixed or heuristically tuned weights can result in suboptimal convergence, instability, or the dominance of large-scale terms. SoftAdapt has been applied to deep learning architectures, physics-informed neural networks (PINNs), and combinatorial optimization problems such as graph matching, with extensions enabling efficient and stable computation via Hadamard-Equipped Sinkhorn operations.

1. Mathematical Foundations of SoftAdapt

SoftAdapt addresses optimization scenarios featuring a composite objective,

$L(\theta) = \sum_{i=1}^K w_i\,\ell_i(\theta),$

where $\ell_i(\theta)$ are component losses (e.g., reconstruction, regularization, adversarial penalties), and $w_i$ are loss weights. Traditional practice either sets all $w_i$ equal or relies on data-specific heuristics or grid search. However, this often leads to inefficient learning dynamics; certain loss terms (due to differing scales) can prematurely dominate gradient updates, impairing convergence and generalization (Heydari et al., 2019).

SoftAdapt replaces fixed $w_i$ with adaptive weights $w_i^{(t)}$, determined on each iteration by the recent trends of the component losses. Specifically, for each loss component $i$ at iteration $t$, SoftAdapt defines a finite-difference rate

$s_i^{(t)} = \bar\ell^{(t)}_i - \bar\ell^{(t-1)}_i,$

where $\bar\ell^{(t)}_i$ is a moving average over the previous $n$ epochs. Intuitively, $s_i^{(t)} > 0$ indicates underperformance (loss increasing), prompting SoftAdapt to increase that term's weight.

The core SoftAdapt rule forms a softmax across these rates: $$\alpha_i^{(t)} = \frac{\exp(\beta\,s_i^{(t)})}{\sum_{j=1}^K \exp(\beta\,s_j^{(t)})}.$$ The focus parameter $\beta > 0$ controls the sharpness; larger $\beta$ narrows focus on the steepest-rising losses, while $\beta \to 0$ yields uniform weights.

Variants include:

• Loss-weighted SoftAdapt: $\alpha_i$ is further scaled by the magnitude of the current loss, favoring larger absolute losses.
• Normalized SoftAdapt: Rates are normalized to prevent single terms from dominating exponentiation, i.e., $s'_i = s_i / (\sum_j |s_j| + \varepsilon)$.

The final weighted loss passed to the optimizer is

$$\mathrm{WLoss}(\theta) = \sum_{i=1}^K \alpha^{(t)}_i\,\ell_i(\theta).$$

In PINN contexts, the adaptive weights are constructed directly from single-step loss changes: for $L(\theta) = \sum_{i=1}^M \lambda_i L_i(\theta),$

$$\Delta L_i(t) = L_i(t) - L_i(t-1), \quad \lambda_i(t) = \frac{ \exp(\Delta L_i(t)/T) }{ \sum_{j=1}^M \exp( \Delta L_j(t)/T ) }.$$

The temperature $T$ governs the redistribution sharpness (Bischof et al., 2021).

2. Algorithmic Implementation and Variants

SoftAdapt is computationally lightweight; it maintains a rolling buffer of length $n$ for each component loss, computes per-iteration finite differences, and applies softmax-based normalization and potential scaling. The algorithmic steps are as follows (Heydari et al., 2019):

1. Update buffers and compute $f_i \leftarrow$ average of last $n$ values of $\ell_i$.
2. Compute rate $s_i = f_i - ($previous$\,f_i )$.
3. Optionally normalize $s_i$.
4. Compute raw softmax weights $\alpha_i$ with focus parameter $\beta$.
5. Optionally apply loss-weighted scaling $\alpha_i \leftarrow f_i \alpha_i$.
6. Normalize $\alpha_i$ to sum to 1.
7. Compute weighted loss WLoss and update model via optimizer.

In PINNs, SoftAdapt adapts the weighting for collocation, boundary, and data loss terms directly at each iteration, without requiring extra backward passes or complex gradient computations. Memory and runtime overheads are negligible compared with model forward/backward passes (Bischof et al., 2021).

Algorithmic summary for PINNs:

Step	Operation	Purpose
1	Evaluate $L_i(t)$	Current loss for each term
2	Retrieve $L_i(t-1)$	Previous-epoch loss
3	Compute $\Delta L_i(t)$	Difference (trend)
4	Softmax $\lambda_i(t)$ over $\Delta L_i$	Adaptive weighting
5	Form weighted joint loss	For optimizer update
6	Update network/parameters	Gradient step
7	Store loss for next iteration	For future steps

The temperature/focus parameter ($\beta$ or $T$) is user-controlled; low values diffuse weight equally, while high values sharply redistribute toward the rapidly increasing losses.

3. Extensions: Adaptive Softassign and Hadamard-Equipped Sinkhorn

An important generalization of SoftAdapt arises in combinatorial optimization, notably in adaptive softassign approaches for graph matching (Shen et al., 2023). Here, the softassign operator is defined as: $$S_\tau(X) = \mathrm{Sinkhorn}\left(\exp\left(\frac{X}{\tau}\right)\right) \in \Sigma_{n \times n},$$ where $X$ is the score or gradient matrix, $\tau$ is the temperature ($\beta = 1/\tau$), and $\Sigma_{n \times n}$ is the Birkhoff polytope (all doubly stochastic matrices). Sinkhorn balancing guarantees row/column normalization.

The principal challenge is to choose the minimal $\tau$ such that the assignment score's deviation from the exact combinatorial optimum is below a prescribed error $\varepsilon$. The method incrementally increases $\beta$ by a stepsize $\Delta\beta$, using the exponential decrease in error established by

$$\| S_\tau(X) - S_\infty(X) \| \leq (c/\mu)\, e^{-\mu/\tau},$$

for suitable $c, \mu$ depending on $X$. The Hadamard-Equipped Sinkhorn formulas enable efficient computation for variable $\beta$. Once $S^{\beta_1}(X)$ is available, $S^{\beta_2}(X)$ is computed via: $$S^{\beta_2}(X) = \mathcal{P}_{sk}( [S^{\beta_1}(X)]^{\circ(\beta_2/\beta_1)} ),$$ where $\mathcal{P}_{sk}$ denotes Sinkhorn-balancing and $^{\circ}$ denotes Hadamard power. This avoids repeated exponentiation and improves numerical stability and runtime.

The adaptive softassign algorithm proceeds by increasing $\beta$ (reducing $\tau$), monitoring the $L_1$ difference of outputs, and stopping when the change falls below $\varepsilon$.

4. Empirical Results and Comparative Analysis

SoftAdapt has demonstrated several empirical advantages across domains:

• Sparse Autoencoder (MNIST): SoftAdapt consistently outperforms fixed $\lambda$ in hybrid losses, achieving 88% classification accuracy of latent codes (vs. 82% for fixed optimal $\lambda=10^{-4}$), with improved convergence and validation loss.
• IntroVAE (CELEBA): SoftAdapt, without pre-tuning, achieves SSIM $\approx 0.8473$ (vs. $0.7838$) and PSNR $\approx 23.93$ (vs. $22.24$) in face reconstruction at epoch 250, matching or exceeding the best fixed-weight variants.
• Classical Test Functions: On the Rosenbrock and Beale functions, SoftAdapt accelerates convergence by up to 43.3% over vanilla gradient descent with equal weighting (Heydari et al., 2019).
• PINN Benchmarks: On Burgers' equation (forward, inverse), Kirchhoff plate bending, and Helmholtz problems, SoftAdapt achieves best or competitive PDE losses and final errors among adaptive loss methods (GradNorm, learning rate annealing, ReLoBRaLo), generally dominating equal/manual weighting on convergence speed and accuracy. For example, on Burgers' forward, achieved training PDE loss is approximately $2\times 10^{-4}$ and validation error $8.1\times 10^{-4}$ (Bischof et al., 2021).

In graph matching, adaptive softassign (SoftAdapt with Hadamard-Equipped Sinkhorn) yields:

• PPI Network (n ≈ 1000, 25% noise): 75.1% accuracy in 22.6 s, surpassing large-graph baselines by 20% in accuracy and an order of magnitude in speed.
• Facebook Networks (n ≈ 4000): 85.7% accuracy in 393 s versus 60.8% in 4453 s for previous spectral methods (Shen et al., 2023).

SoftAdapt's overhead is minimal: typically 1–2% per iteration for deep learning and 0.4 s per 1,000 Adam steps in PINNs, compared to much larger overheads for gradient-norm or annealing-based schemes.

5. Hyperparameter Selection and Practical Guidance

Key hyperparameters include:

• Focus parameter $\beta$ (or temperature $T=1/\beta$): Defaults such as $\beta=0.1$ or $T=0.1$ are effective in diverse settings. Larger $\beta$ (lower $T$) sharpens focus but increases sensitivity to noise in loss trends.
• Lookback window $n$: Small values ($n=3$–$5$) suffice for rate estimation. Excessive $n$ reduces the adaptivity of weights. Initialization typically sets $\alpha_i = 1/K$ until buffers are populated.
• Stabilization $\varepsilon$: Small constant ($10^{-8}$) prevents division instabilities.

For adaptive softassign, the increment $\Delta\beta \approx \ln n$ is recommended for scaling with problem size, and monitoring of assignment change per iteration provides a principled stopping criterion (Shen et al., 2023).

Across all use cases, SoftAdapt is architecturally and optimizer agnostic and does not require gradient-norm computation, auxiliary networks, or manual search over loss weights.

6. Limitations and Directions for Extension

SoftAdapt has limitations in settings with highly non-differentiable, high-variance, or discrete loss terms, potentially requiring additional smoothing or lower $\beta$ for stability (Heydari et al., 2019). In stiff multi-task settings with strongly conflicting gradients, gradient-alignment-based reweighting (e.g., GradNorm) may outperform scalar-rate based methods. In PINN and graph-matching applications, SoftAdapt may underweight initially small but important losses if temperature is not sufficiently low.

Potential extensions include:

• Learnable focus parameter or automated temperature/β scheduling.
• Incorporation of higher-order statistics (e.g., variance of loss change) into weight updates.
• Integration with gradient-norm normalization for joint magnitude and focus control.
• Adaptation to policy-gradient reinforcement learning and settings with non-scalar or non-differentiable objectives.
• Formal convergence analysis, particularly under stochastic gradients.
• Automated selection of moving average window size based on loss curvature.

SoftAdapt's principled yet simple approach to allocating attention across loss components enables more robust and efficient training in composite-objective scenarios, mitigating the need for heuristic or dataset-specific hyperparameter searches (Heydari et al., 2019, Bischof et al., 2021, Shen et al., 2023).