Conceptio › Archive › arXiv CS

arXiv CSopen access

Conflict-Aware Harmonized Rotational Gradient for Multiscale Kinetic Regimes

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

CONFLICT-AWARE HARMONIZED ROTATIONAL GRADIENT FOR MULTISCALE KINETIC REGIMES ∗

arXiv:2604.24745v1 [cs.LG] 27 Apr 2026

ZHANGYONG LIANG† Abstract. In this paper, we propose a harmonized rotational gradient method, termed HRGrad, for simultaneously tackling multiscale time-dependent kinetic problems with varying small parameters. These parameters exhibit asymptotic transitions from microscopic to macroscopic physics, making it a challenging multi-task problem to solve over all ranges simultaneously. Solving tasks in different asymptotic regions often encounter gradient conflicts, which can lead to the failure of multi-task learning. To address this challenge, we explicitly encode a hidden representation of these parameters, ensuring that the corresponding solving tasks are serialized for simultaneous training. Furthermore, to mitigate gradient conflicts, we segment the prediction results to construct task losses and introduce a novel gradient alignment metric to ensure a positive dot product between the final update and each loss-specific gradient. This metric maintains consistent optimization rates for all task losses and dynamically adjusts gradient magnitudes based on conflict levels. Moreover, we provide a mathematical proof demonstrating the convergence of the HRGrad method, which is evaluated across a range of challenging asymptotic-preserving neural networks (APNNs) scenarios. We conduct an extensive set of experiments encompassing the Bhatnagar-Gross-Krook (BGK) equation and the linear transport equation in all ranges of Knudsen number. Our results indicate that HRGrad effectively overcomes the ‘failure modes’ of APNNs in these problems. Key words. Multi-task learning, Asymptotic transitions, Bhatnagar-Gross-Krook equation, Linear transport equation, Asymptotic-preserving neural networks AMS subject classifications. 65N22, 65N55, 68T07

1. Introduction. Kinetic equations have been widely used in many areas such as mechanics, rarefied gas, plasma physics, astrophysics, semiconductor device modeling, and social and biological sciences [42]. They describe the non-equilibrium dynamics of a system composed of a large number of particles and bridge atomistic and continuum models in the hierarchy of multiscale modeling. The Boltzmann-type equation, as one of the most representative models in kinetic theory, provides a powerful tool to describe molecular gas dynamics, radiative transfer, plasma physics, and polymer flow [1]. In particular, the linear semiconductor Boltzmann-BGK equation we will study in this article describes the evolution of electron density(in the phase space) within semiconductor devices, in the presence of an external electrical potential and scatterings among various particles [17]. Ignoring relativistic effects, quantum mechanics provides a sufficiently accurate framework for understanding the physical properties of matter. Within the broad category of multiscale modeling, kinetic equations play a central role as bridges between atomistic and continuum descriptions [4]. However, a fundamental challenge arises from their inherently multiscale nature: kinetic equations typically involve small or multiple spatial and/or temporal scales, characterized by the dimensionless Knudsen number, which represents the ratio of the mean free path (or time) to the macroscopic length (or time) scale. In multiscale computations, one often encounters the need to couple models across different scales, each requiring distinct numerical strategies [45]. When the Knudsen number is small, kinetic equations can be rigorously approximated by macroscopic hydrodynamic or diffusion equations [2]. Yet, direct numerical simulations in this regime become prohibitively expensive, since resolving the small scales explicitly is computationally intractable. To overcome this difficulty, ∗ Submitted to the editors DATE. † National Center for Applied Mathematics, Tianjin University, Tianjin, 300072, China.

Z.Y. LIANG

the concept of asymptotic-preserving (AP) schemes was introduced. An AP scheme is designed to seamlessly capture the asymptotic transition from kinetic or hyperbolic equations to their macroscopic hydrodynamic or diffusive limits within the discrete numerical framework [14]. This strategy avoids the explicit coupling of microscopic and macroscopic solvers; instead, the microscopic scheme naturally degenerates into a macroscopic solver in the vanishing Knudsen number limit. Over the past two decades, AP schemes have proven to be a powerful and robust framework for multiscale problems. Their main advantage lies in their efficiency and accuracy in the hydrodynamic or diffusive regime, since they eliminate the need to resolve small physical parameters while faithfully capturing the emergent macroscopic behavior. As such, AP schemes provide a unifying and effective methodology for handling the inherent complexity of multiscale kinetic and hyperbolic problems. The high dimensionality and multiscale nature of kinetic simulations call for efficient numerical methods. Classical Monte Carlo methods are widely used but suffer from low-order accuracy and increasing cost as the Knudsen number decreases [32, 8, 33]. To alleviate these difficulties, a series of approaches based on deep neural networks (DNNs) has recently been proposed. Training these DNN-based solvers requires minimizing a high-dimensional, non-convex loss and remains difficult in the small-parameter regimes of multiscale kinetic problems. For time-dependent PDEs, physics-informed neural networks (PINNs) minimize least-squares residual risks of the governing equations [35]. Yet, depending on how the loss function is designed, such formulations are only able to capture the leading-order or single-scale behavior of the solution. Such limitations arise from loss design, first-order optimization, and the Frequency Principle [47], which biases DNNs toward low frequencies and away from small-scale structures. To address these challenges, a new class of asymptoticpreserving neural networks (APNNs) [15, 16, 25, 5] has been developed, which integrates the asymptotic-preserving (AP) strategy with PINNs to tackle multiscale physical problems involving small-scale parameters. By embedding asymptotic information, APNNs overcome the shortcomings of standard PINNs and yield robust multiscale approximations. However, existing APNNs are typically limited to problems with a specific small-scale parameter. For many multiscale physical problems, the scale parameter often spans the entire spectrum from macroscopic regimes, through transitional regions, to microscopic regimes. For example, in the radiative transfer equation, the scale parameter (i.e., the Knudsen number) may vary significantly, covering regimes from kinetic to diffusive, thus exhibiting pronounced multiscale transitions. A natural question is whether APNNs can be pretrained across a range of scale parameters and directly predict the solution for downstream values, including those unseen during pretraining. Even with AP decomposition, however, heterogeneous scale-dependent behaviors can induce severe gradient conflicts across asymptotic regimes. This raises two intertwined challenges: pretraining AP solvers across the parameter spectrum and resolving the resulting multi-task gradient conflicts. Existing methods for parameter-dependent PDEs either rely on data-intensive neural operators or use collocation-based solvers that typically require retraining for new parameter values. To address this limitation, two directions have been explored: multi-task learning (MTL) [3, 51, 52, 38] over a parameter domain, and parameterization strategies for specific parameter values. Existing approaches to conflicting gradients in MTL mainly focus on dynamic loss re-weighting in classical machine learning [41, 49], with only limited extensions to PDE solvers, where residual, initial, and boundary losses may be imbalanced. In this setting, self-adaptive weighting schemes [37, 40, 44, 46, 48]

HRGRAD

are widely used, while self-paced learning (SPL) [13, 50, 53] further emphasizes difficult collocation-point losses; Cognitive Physics-Informed Neural Networks (CoPINN) [9] recast weighted-loss optimization as collocation-point selection. Another line of work addresses imbalance through optimizer preconditioning. Unlike self-adaptive weighting schemes that focus on gradient magnitudes, these methods modify gradient directions in MTL, where opposing gradients can yield inefficient updates. Representative methods include PCGrad [49], CAGrad [23], IMTL-G [24], ConFIG [26], and SOAP [43]. However, these approaches mainly address PDEs with fixed parameters. For simultaneous pretraining over varying parameter values, recent work has explored meta-learning [7, 12, 39], typically via a bi-level formulation in which the outer loop assigns weights across parameter values and the inner loop solves the associated PINN problems. Another approach directly parameterizes the scale parameter itself [34], feeding sampled parameter values to the network as an additional input while encoding collocation points and parameters separately before concatenation. Such parameterization, however, is less effective for strongly multiscale problems with drastic regime changes and severe gradient conflicts. To the best of our knowledge, no existing optimization framework resolves these conflicts for simultaneous training across asymptotically distinct parameter regimes. Motivated by these challenges, we develop a physics-informed continuous multitask optimization framework for asymptotic-preserving neural networks over a continuum of scale parameters. Our central observation, which is formalized in Section 3, is that multiscale APNN training suffers from two coupled failure modes: directional contradiction between microscopic and macroscopic task gradients, and magnitude disparity caused by stiff asymptotic scaling. These effects make standard projectionbased gradient manipulation intrinsically dissipative for multiscale kinetic learning, as it clips weak microscopic gradients precisely when non-equilibrium features are the most vulnerable. To overcome this difficulty, we propose the Harmonized Rotational Gradient (HRGrad) method, which replaces lossy Euclidean projection with harmonized, energy-preserving gradient rotation and fair aggregation. As illustrated in Figure 1, HRGrad proceeds through four coordinated stages: harmonized-cone construction with physical anchoring, isometric task-wise rotation, MER-based angle selection, and fair consensus aggregation. The first three stages resolve directional contradiction without dissipating gradient energy. The final aggregation stage mitigates magnitude disparity by restoring balanced multiscale contributions in the resulting update. Our main contributions can be summarized as follows: • HRGrad for multiscale APNNs. We propose the Harmonized Rotational Gradient (HRGrad) method as a physics-informed gradient preprocessing and aggregation framework for continuous multi-task APNN training over a continuum of Knudsen numbers, enabling one unified model to learn across kinetic, transitional, and macroscopic regimes. • Kinetic-aware gradient guidance. We design HRGrad to explicitly account for the multiscale structure of kinetic equations, so that weak but essential updates associated with microscopic non-equilibrium dynamics are protected during APNN optimization rather than being overwhelmed by dominant macroscopic gradients. • Harmonized rotational preprocessing. We introduce a harmonized cone together with a physically anchored reference direction, and use manifold isometric rotation to preprocess conflicting task gradients. This mechanism removes destructive gradient conflicts while preserving the norm information

Z.Y. LIANG

εk ∈ B

Cones and anchor

Ĝπ2 Kinetic-regime gradient spectrum

Ĝπ3 d∗

d∗ = U

ε→0 gmacro

Ĝπ1

k U (Ĝπk )

ε→1 gmicro

transition mixed gi

⊥ g1 α∗ 2

Isometric rotation εj -task wj d∗ rj (αj ) αj ḡj

g2rot

rj (εj ) = cos αj ḡj + sin αj wj d∗

gHRGrad

g1rot α∗ 1 ⊥ g2 g1

MER angle search L search

α∗ j (εj ) align prox MER

g⊤ 1 g2 < 0

Fair aggregation

α∗

α π/2

g2rot

{εi }m i=1

g1rot

h i † ⊤ gu (B) = U (MB ) 1

Fig. 1. Schematic of the proposed HRGrad framework. Starting from task gradients sampled over macro, transitional, and micro regimes, HRGrad first constructs a harmonized cone together with a physical anchor. Then, it rotates conflicting gradients by isometric operators and selects MER-optimal rotation angles. Finally, it performs fair aggregation to produce the update direction.

carried by the original gradients, especially in the kinetic regime. • Magnitude restoration and fair aggregation. We further equip HRGrad with magnitude restoration and fair consensus aggregation, so that the final update remains positively aligned with all rotated task gradients and balances optimization progress across different scales and physical objectives. • Theory and multiscale validation. We provide a convergence analysis for the proposed gradient preprocessing framework and validate HRGrad on the Boltzmann–BGK equation, the linear transport equation, the ES-BGK equation, and the linear semiconductor Boltzmann–Poisson equation, demonstrating stable and robust performance across diverse multiscale kinetic regimes. The remainder of the paper is organized as follows. Section 2 introduces the Boltzmann-BGK equation and the linear transport equation, commonly used asymptoticpreserving numerical methods, and the formulation of the APNNs loss. Section 3 analyzes the failure modes of multiscale APNN training and develops the HRGrad method, including its geometric motivation, theoretical properties, and implementation. Section 4 showcases numerical results for several kinetic equations solved by HRGrad and compares them with existing multi-task learning baselines and parameterized APNN strategies. Finally, Section 5 offers the conclusion and outlook. 2. Preliminaries. This section briefly introduces the kinetic equations of interest and the framework of asymptotic-preserving neural networks (APNNs). 2.1. Multiscale kinetic problems. We consider two prototypical multiscale kinetic equations that describe the evolution of a particle distribution function f (t, x, v) under the scaling of the Knudsen number ε.

HRGRAD

The Boltzmann-BGK equation. The Boltzmann-BGK equation reads (2.1)

∂t f + v∂x f =

1 (M (U ) − f ) , ε

where M (U ) is the local Maxwellian defined byRthe macroscopic moments U = ⟨mf ⟩ = (ρ, ρu, E)T , with m = (1, v, 21 v 2 )T and ⟨·⟩ = · dv. As ε → 0, we have f → M (U ), and the moments satisfies the local conservation laws: (2.2)

∂t U + ∂x ⟨vmf ⟩ = 0,

which reduces to the macroscopic compressible Euler equations [10]. The linear transport equation. The one-dimensional linear transport equation under diffusive scaling is Z 1 1 1 ′ (2.3) ε∂t f + v∂x f = f dv − f . ε 2 −1 By defining the even and odd parities r(t, x, v) = 21 (f (v) + f (−v)) and j(t, x, v) = 1 2ε (f (v) − f (−v)), the system transforms into: ( 2 ε ∂t r + ε2 v∂x j = ρ − r, (2.4) ε2 ∂t j + v∂x r = −j, R1 where ρ = ⟨r⟩ := 0 r dv. In the diffusive limit ε → 0, we obtain r → ρ and j → −v∂x ρ, and ρ satisfies the diffusion equation: (2.5)

∂t ρ −

1 ∂xx ρ = 0. 3

2.2. Asymptotic-preserving neural networks (APNNs). Standard PhysicsInformed Neural Networks (PINNs) approximate f by directly minimizing the PDE residuals. However, in the diffusive or fluid regime (ε → 0), the standard PINN loss fundamentally fails to capture the macroscopic evolution constraints, leading to incorrect predictions of the macroscopic dynamics [15, 16, 25]. To overcome this, Asymptotic-Preserving Neural Networks (APNNs) [15] modify the loss function by explicitly coupling the microscopic kinetic equations with their corresponding macroscopic conservation laws.

Fig. 2. Illustration of the APNNs framework.

As shown in Figure 2, the strategy involves embedding the macroscopic equations directly into the neural network training. For instance, when solving the BoltzmannBGK equation, APNNs approximate f and U by continuous neural networks fθ and

Z.Y. LIANG

Uθ , and define the loss functional as (2.6)

LεAPNNs = LεRes + LεClaw + LεCnt + LεBC + LεIC ,

where the PDE structure constraints are:  2  Lε = λ1 ε(∂t fθ + v ∂x fθ ) − (M (Uθ ) − fθ ) ,   Res 2 (2.7) LεClaw = λ2 ∂t Uθ + ∂x ⟨vmfθ ⟩ ,    ε 2 LCnt = λ3 Uθ − ⟨mfθ ⟩ . Similarly, for the linear transport equation, APNNs enforce the macroscopic density evolution: (2.8)

∂t ρθ + ⟨v∂x jθ ⟩ = 0.

By jointly minimizing the residuals of both the microscopic and macroscopic equations, APNNs preserve the correct asymptotic behavior as ε → 0, enabling uniformly accurate simulations across all scales of ε. 3. Methodology. To simultaneously solve the multiscale kinetic equations across the continuous Knudsen spectrum ε ∈ (0, 1], we formulate the training process as a continuous multi-task learning (MTL) problem. In each training iteration, we sample a mini-batch of m scale parameters B = {ε1 , ε2 , · · · , εm }. For each sampled parameter εi , we adopt the Asymptotic-Preserving Neural Network (APNN) architecture described in Section 2 as the unified base model. Consequently, the objective for the i-th task is precisely defined as the APNN loss functional evaluated at εi , i.e., i Li (θ) := LεAPNNs (θ) (as formulated in Eq. 2.7 for the Boltzmann-BGK equation or its equivalent for the linear transport equation). The optimization trajectory is thus driven by the task-specific gradients gi = ∇θ Li (θ). When the neural solver is faced with tasks across drastically different kinetic regimes simultaneously, we observe that existing MTL methods suffer from severe ‘failure modes’, making multi-task optimization exceptionally challenging. To clarify the rationale and innovation behind our proposed method, we first mathematically dissect these failure modes from the perspective of multiscale gradient conflicts. Based on this analysis, we develop the Harmonized Rotational Gradient (HRGrad) method, which abandons traditional dissipative projections in favor of manifold isometric rotations, ensuring that the preconditioning gradient consistently maintains physically valid, conflict-free, and energy-preserving updates. 3.1. Gradient conflict in multiscale kinetic regimes. Although APNNs employ macro-micro and even-odd decompositions [15, 16, 25] to enforce asymptotic correctness within a single regime, jointly training tasks across vastly different ε values with shared parameters θ inevitably triggers severe gradient conflicts (i.e., ⟨gi , gj ⟩ < 0). This geometric misalignment originates from the fundamentally opposing physical behaviors captured by the APNN sub-losses across different kinetic regimes. Specifically, let gmicro = ∇θ LεRes denote the gradient driven by microscopic transport mechanisms, and gmacro = ∇θ LεClaw (or macroscopic density evolution) denote the gradient driven by macroscopic relaxation and conservation. We analyze the dynamics across three typical regimes: • Micro-limit regime (ε → 1; kinetic side): The dynamics are highly nonequilibrium and transport-dominated. The optimization is overwhelmingly

HRGRAD

guided by the microscopic residual LεRes (e.g., the transport term ε(∂t fθ + v∂x fθ ) in Eq. 2.7). The corresponding gradients gmicro actively push the network to capture and preserve high-frequency, anisotropic structures in the phase space (such as Knudsen boundary layers), strongly resisting any premature relaxation toward the local Maxwellian M (Uθ ). • Macro-limit regime (ε → 0; fluid/diffusive side): The system reaches the hydrodynamic or diffusive limit. The stiff collision penalty 1ε (M (Uθ )−fθ ) strictly enforces the local relaxation fθ → M (Uθ ) (or rθ → ρθ for the linear transport equation). The optimization focus shifts entirely to the macroscopic conservation laws (e.g., LεClaw in Eq. 2.7 or the density evolution in Eq. 2.8). The gradients gmacro from this regime drive the neural network parameters to rapidly collapse onto a smooth, low-frequency macroscopic equilibrium manifold, heavily penalizing deviations from local equilibrium. • Intermediate regime (ε ∈ (εmin , 1)): Neither physical limit dominates; free transport and collisional relaxation contribute comparably. The network receives mixed gradient signals demanding both macroscopic consistency (moment evolution) and microscopic fidelity (non-equilibrium detail). Minibatches from this band yield gradients that fluctuate between macro-scale isotropic smoothing and micro-scale anisotropic retention, severely increasing the frequency of pairwise disagreements. From a mathematical optimization perspective, this multiscale physical transition introduces two fatal “failure modes” for conventional gradient aggregation: Failure Mode 1: Directional Contradiction (⟨gmicro , gmacro ⟩ < 0). Macroscopic objectives actively smooth out non-equilibrium details to enforce local equilibrium, while microscopic objectives strive to preserve them. A parameter update step ∆θ = −ηgmacro that strictly favors the macroscopic continuum limit will inevitably degrade the microscopic transport fidelity. Based on the first-order Taylor expansion, the variation in the microscopic loss is: (3.1)

Lmicro (θ + ∆θ) ≈ Lmicro (θ) − η⟨gmicro , gmacro ⟩.

Since the underlying physical mechanisms are fundamentally opposing (isotropic smoothing vs. anisotropic high-frequency preservation), we rigorously have ⟨gmicro , gmacro ⟩ < 0. This yields an undesirable strict increase in the micro-loss (Lmicro (θ + ∆θ) > Lmicro (θ)). Consequently, these gradients inherently point in geometrically opposing directions, yielding persistent obtuse angles ϕ ∈ (π/2, π] and triggering severe negative transfer across the continuous ε spectrum. Failure Mode 2: Magnitude Disparity and Energy Clipping. Beyond directional conflict, the explicit O(1/ε) (in Eq. 2.1) or O(1/ε2 ) (in Eq. 2.4) scaling introduces extreme numerical kinetic stiffness. As ε → 0, the macroscopic conservation and relaxation losses exponentially dominate the total loss landscape, creating an extreme spectral gap in gradient magnitudes: ∥gmacro ∥2 1 (3.2) lim =O → ∞, (p ≥ 1). ε→0 ∥gmicro ∥2 εp This mathematically guarantees that the macroscopic gradient overwhelmingly eclipses the microscopic gradient (∥gmacro ∥2 ≫ ∥gmicro ∥2 ). Under this extreme multiscale heterogeneity, standard gradient manipulation methods fundamentally fail. Traditional optimizers (e.g., PCGrad, ConFIG) rely on Euclidean orthogonal projections P to resolve conflicts, blindly projecting the conflicting components of the microscopic gradients onto the normal plane of the dominant

Z.Y. LIANG

macroscopic gradient gmacro : (3.3)

P(gmicro ) = gmicro −

⟨gmicro , gmacro ⟩ gmacro . ∥gmacro ∥22

By evaluating the remaining kinetic energy (L2 -norm) of the projected microscopic gradient, we reveal the inherent mathematical flaw of the projection operator: ⟨gmicro , gmacro ⟩2 2 2 ∥P(gmicro )∥2 = ∥gmicro ∥2 1 − ∥gmicro ∥22 ∥gmacro ∥22 (3.4) 2 = ∥gmicro ∥2 (1 − cos2 ϕ) = ∥gmicro ∥22 sin2 ϕ. This equation uncovers the catastrophic energy clipping mechanism: when the multiscale conflict is severe (ϕ → π), sin2 ϕ → 0. Because the original microscopic kinetic signal ∥gmicro ∥2 is already extremely weak due to the magnitude disparity, the projection mathematically acts as a numerical low-pass filter, irreversibly diminishing its magnitude to near zero (∥P(gmicro )∥2 → 0). As illustrated in Figure 3(a), traditional projection methods (e.g., PCGrad, ConFIG) mathematically act as an energy clipping mechanism (∥P(gmicro )∥2 ≪ ∥gmicro ∥2 ). As the conflict angle ϕ → 180◦ , the norm ∥P(gmicro )∥2 → 0, irreversibly dissipating high-frequency kinetic features. To address this, our proposed Harmonized Rotational Gradient (HRGrad) performs an isometric manifold rotation towards a physical anchor d∗ (Figure 3(b)). This approach rigorously guarantees non-conflict while flawlessly preserving 100% of the gradient rot magnitude (i.e., kinetic physical energy, ∥gmicro ∥2 ≡ ∥gmicro ∥2 ). N (d∗rot ) gmicro

N (gmacro ) R

gmicro

ϕ P(gmicro )

(a) Euclidean projection

gmicro gmacro

ϕ d∗

(b) Manifold isometric rotation

Fig. 3. Comparison of gradient aggregation strategies under multiscale conflicts. (a) Euclidean projection severely clips the magnitude of the microscopic gradient P(gmicro ). (b) Manifold isometric rot ∥2 ≡ ∥gmicro ∥2 . rotation avoids directional conflicts while preserving the gradient magnitude ∥gmicro

As a result, the high-frequency non-equilibrium physical features are permanently erased during parameter updates. The shared APNNs suffer from asymptotic manifold drift and degenerate into a naive macroscopic Euler/Diffusion solver, entirely losing its capacity to capture multiscale transition details. This profound vulnerability mathematically necessitates the development of the Harmonized Rotational Gradient (HRGrad) method, which abandons dissipative projections in favor of isometric manifold rotations. 3.2. Harmonized Rotational Gradient (HRGrad). Generally, we consider a multiscale optimization procedure with a set of m task losses, i.e., {L1 , L2 , · · · , Lm }, sampled from a continuous Knudsen spectrum ε ∈ (0, 1]. Let G = [g1 , g2 , · · · , gm ] ∈

HRGRAD

(a) ε = 1.0

(b) ε = 0.1

(c) ε = 0.01 Fig. 4. MAD-MTL predicts the macroscopic density moments ρ of the 1D Riemann problem. The comparison results are shown for (a) ε = 1.0, (b) ε = 0.1, and (c) ε = 0.01.

(a) ε = 1.0

(b) ε = 0.1

(c) ε = 0.01 Fig. 5. MAD-MTL predicts the macroscopic velocity moments u of the 1D Riemann problem. The comparison results are shown for (a) ε = 1.0, (b) ε = 0.1, and (c) ε = 0.01.

RD×m denote the primal gradient matrix corresponding to each loss function, where gi = ∇θ Li . Traditional gradient manipulation methods (e.g., PCGrad, ConFIG) rely on Euclidean orthogonal projections to resolve gradient conflicts (i.e., gi⊤ gj < 0). However, orthogonal projection P is inherently a non-isometric mapping, meaning that its operation strictly diminishes the gradient norm (∥P(gi )∥ < ∥gi ∥). In kinetic theory, microscopic non-equilibrium mechanisms (e.g., Knudsen boundary layers) produce gradients with significantly smaller magnitudes than macroscopic conservation laws. As intuitively illustrated in Figure 7(a), PCGrad resolves conflicts by summing orthogonal projections (O), which acts as a mathematical energy clipping oper-

Z.Y. LIANG

(a) ε = 1.0

(b) ε = 0.1

(c) ε = 0.01 Fig. 6. MAD-MTL predicts the macroscopic temperature moments T of the 1D Riemann problem. The comparison results are shown for (a) ε = 1.0, (b) ε = 0.1, and (c) ε = 0.01.

ator (∥O(gi , gj )∥ ≪ ∥gi ∥), inadvertently dissipating these high-frequency microscopic physical features. IMTL-G (Figure 7(b)) rescales gradients to be equal-length but may not fully guarantee manifold feasibility. ConFIG (Figure 7(c)) attempts to mitigate magnitude loss by summing unit vectors U of the orthogonal components, but the underlying direction still relies on a lossy projection basis. Furthermore, blindly seeking non-conflict directions in Euclidean space often forces the update trajectory to drift away from the physically feasible asymptotic-preserving (AP) manifold. gHRGrad gPCGrad

⊥ g1

g2O(g2 , g1 )

O(g1 , g2 ) g1

(a) PCGrad

α2 g2

g2 gIMTL-G

gConFIG

g2rot

Hd∗ ⊥ g2 g1rot

U(O(g2 , g1 )) U(O(g1 , g2 ))

α1 g1

(b) IMTL-G

(d) HRGrad (Ours)

Fig. 7. Geometric sketch comparison of gradient manipulation methods under conflicting tasks (g1 , g2 ). (a) PCGrad relies on energy-clipping orthogonal projections. (b) IMTL-G rescales gradients to be equal-length. (c) ConFIG sums unit vectors of orthogonal components. (d) HRGrad isometrically rotates gradients into a Harmonized Cone H, resolving conflicts while fully preserving kinetic energy.

To address these fundamental theoretical flaws, we propose the Harmonized Rotational Gradient (HRGrad). As visualized in Figure 7, HRGrad elevates conflict resolution from Euclidean projection to manifold isometric rotation. Specifically, PCGrad (Figure 7(a)) resolves conflicts by summing orthogonal projections (O), which inherently acts as an energy clipping operator (∥O(gi , gj )∥ ≪ ∥gi ∥), irreversibly dissipating microscopic physical features. IMTL-G (Figure 7(b)) rescales gradients to be equal-length but may not fully guarantee manifold feasibility. Con-

HRGRAD

FIG (Figure 7(c)) attempts to mitigate magnitude loss by summing unit vectors U of the orthogonal components, but the underlying direction still relies on a lossy projection basis. In contrast, HRGrad (Figure 7(d)) explicitly constructs the Harmonized Cone H = K ∩ K∗ and extracts a physically valid anchor d∗ . Conflicting gradients are deflected into H via manifold isometric rotation (orange arcs), mathematically ensuring absolute non-conflict while flawlessly preserving 100% of the kinetic energy (∥girot ∥ ≡ ∥gi ∥). This approach is strictly bounded within a physically valid harmonized cone, which operates in four rigorous stages. Stage 1: Harmonized Cone and Physical Anchoring. A physically valid multi-scale update direction must be both feasible (expressible as a non-negative linear combination of actual physical mechanisms) and non-conflicting across all ε scales. We define the Primal Gradient Cone K = {Gλ | λ ∈ Rm + } to guarantee physical feasibility, and the Dual Gradient Cone K∗ = {y ∈ RD | G⊤ y ≥ 0m } to guarantee absolute non-conflict. HRGrad restricts the optimization search space to their exact geometric intersection, the Harmonized Cone H := K ∩ K∗ . By leveraging the Double Description Method (DDM) to solve the algebraic constraint G⊤ Gλ ≥ 0m , we analytically extract the extreme rays Π = [π1 , · · · , πp ] of the feasible weight space. The physical extreme rays in the gradient space are given by rj = Gπj . We then determine the absolute physical target direction d∗ ∈ Int(H) as the normalized centroid of these rays:  (3.5)

d∗ = U 

p X

 U(rj ) ,

j=1

where U (x) = x/(∥x∥ + δ) is the normalization operator, and δ > 0 is a small constant introduced for numerical stability. This strictly locks the neural network’s evolution trajectory onto the AP manifold. Stage 2: Manifold Isometric Rotation Operator. For any normalized task gradient ḡi = U(gi ) that conflicts with the physical anchor (i.e., ḡi⊤ d∗ < 0), HRGrad completely abandons the projection operator. Instead, we construct a parameterized rotation operator on the 2D Riemannian tangent plane Ti = span(ḡi , d∗ ). Using the Gram-Schmidt process, we first extract the feasible orthogonal reference state: (3.6)

wi = U d∗ − (ḡi⊤ d∗ )ḡi .

The parameterized SO(2) isometric rotation operator ri (αi ) with a deflection angle αi ∈ [0, π/2] is formulated as: (3.7)

ri (αi ) = cos(αi )ḡi + sin(αi )wi .

This operator guarantees Isometric Fidelity (∥ri (αi )∥ ≡ 1). The directional deflection is completely decoupled from the vector norm, ensuring that 100% of the micro-scale kinetic energy is preserved during conflict resolution. Stage 3: Optimizing MER via MCAE. Blindly rotating the gradient to a completely orthogonal boundary constitutes a local hard-adjustment, which severely distorts task-specific physical features. To achieve a global Multi-task Equilibrium ∗ Relationship (MER), we seek the optimal rotation angles α∗ = [α1∗ , . . . , αm ] by mini-

Z.Y. LIANG

mizing a global energy functional: (3.8) X 1 min m L(α) = 1 − ri (αi )⊤ rj (αj ) +λ m(m − 1) i<j α∈[0,π/2] {z } | Global Alignment (Suppress UCAE)

1 X ∥ri (αi ) − ḡi ∥2 4m i=1 {z } |

Global Proximity (Suppress OCAE)

A second-order Taylor expansion of the Miss-Correction Angular Error (MCAE) demonstrates that any deviation from the optimal angle αi∗ introduces quadratic optimization penalties. Hard projections induce severe Over-Correction Angular Error (OCAE), obliterating microscopic specificities, while insufficient rotation leads to Under-Correction Angular Error (UCAE). Furthermore, to address the severe kinetic stiffness that varies drastically across ε, we introduce an adaptive search strategy for the inner optimization of Eq. 3.8. Let (t) (t) (t−1) (t−1) (t) (t) ∆i = (Li −Li )/(Li +δ) be the relative loss change, and st = std(∆1 , . . . , ∆m ) (t) be its standard deviation. The number of inner SGD steps αsteps is dynamically adjusted: st (t) (3.9) αsteps = αmin + (αmax − αmin ) . st + kstd This mechanism allocates more computational effort to accurately locate α∗ when the multi-scale system enters highly stiff transitional regimes (where st surges), strictly preventing OCAE distortion on high-frequency micro-features. Stage 4: Magnitude Restoration and Fair Aggregation. After achieving the perfect angular harmony, HRGrad flawlessly restores the intrinsic physical magnitudes to prevent macroscopic regimes from dominating the update. The physically restored gradient is given by: ( ∥gi ∥ · ri (αi∗ ), if ḡi⊤ d∗ < 0 rot (3.10) gi = . gi , otherwise To ensure scale-invariance and fair learning across the entire continuous spectrum, we adopt a pseudo-inverse equal-length aggregation. We construct the normalized rot orientation matrix M = [U(g1rot ), U(g2rot ), · · · , U(gm )]. Utilizing the Moore-Penrose † pseudoinverse M , we derive the globally fair consensus direction gu that maintains equal projection lengths across all rotated tasks: (3.11) gu = U (M † )⊤ 1m , where 1m is a vector of ones. The final HRGrad update gradient is formulated by projecting the physically restored gradients onto this fair direction: ! m X rot ⊤ (3.12) gHRGrad = GHRGrad (g1 , · · · , gm ) := (gi ) gu gu . i=1

3.3. Model architecture. To seamlessly integrate HRGrad into deep learning optimizers, we apply the GHRGrad operator directly to the pseudo-first-moments of the gradients within the Adam optimizer. This elegantly stabilizes the multiscale optimization trajectory. The complete process is detailed in Algorithm 3.1. 3.4. Pre-training.

HRGRAD

Algorithm 3.1 HRGrad for Multi-task Learning Require: Initial network weights θ0 , learning rate γ, Adam coefficients β1 , β2 , ϵ. Require: HRGrad hyper-parameters αmin , αmax , kstd , λ, δ. 1: Initialize first momentum [mg1 ,0 , mg2 ,0 , . . . , mgm ,0 ] ← [0, 0, . . . , 0] and second momentum v0 ← 0. 2: Initialize step counters [tg1 , tg2 , . . . , tgm ] ← [0, 0, . . . , 0]. 3: for t = 1, 2, . . . do 4: i ← (t − 1) mod m + 1 5: tgi ← tgi + 1 6: mgi ,tgi ← β1 mgi ,tgi −1 + (1 − β1 )∇θt−1 Li mg1 ,tg mgm ,tg 1 m 7: Compute bias-corrected momentums: [m̂g1 , . . . , m̂gm ] ← , . . . , tg m tg 1 1−β1

8: 9: 10: 11: 12: 13: 14: 15: 16: 17: 18: 19: 20: 21: 22: 23: 24: 25: 26: 27: 28: 29: 30:

1−β1

HRGrad Operator GHRGrad applied to Ĝ = [m̂g1 , . . . , m̂gm ]: Stage 1 (Harmonized cone and physical anchor). Solve the DDM constraint Ĝ⊤ Ĝ λ ≥ 0m to extract extreme rays Π = [π1 , . . . , πp ] of H = K ∩ K∗ . ComputePphysical extreme rays rk ← Ĝπk for k = 1, . . . , p and the physical anchor p d∗ ← U k=1 U(rk ) ∈ Int(H). Stage 2 (Per-task rotation references). Form the unit gradients m̄gj ← U(m̂gj ) for j = 1, . . . , m and identify the conflicting ∗ set C ← { j : m̄⊤ gj d < 0 }. for j ∈ C do ∗ wj ← U d∗ − (m̄⊤ gj d )m̄gj . end for Stage 3 (Joint MER optimization with adaptive steps). (t) Compute relative loss variance st and adaptive step count αsteps . (t) Fix αj ≡ 0 for j ∈ / C and run αsteps inner SGD iterations on the MER objective L(α) ∗ in Eq. 3.8, jointly over {αj }j∈C ⊂ [0, π/2], to obtain α∗ = [α1∗ , . . . , αm ]. Stage 4 (Magnitude restoration and fair aggregation). for j = 1 to m do if j ∈ C then rj ← cos(αj∗ )m̄gj + sin(αj∗ )wj , mrot ← ∥m̂gj ∥ rj . j else ← m̂gj . mrot j end if end for rot Form the normalized rotated matrix M ← [U(mrot 1 ), . . . , U(m m )]. † ⊤ Compute the fair aggregation direction gu ← U (M ) 1m . P m rot ⊤ m̂g ← GHRGrad (m̂g1 , . . . , m̂gm ) := j=1 (mj ) gu gu .

31: Estimate the unified gradient: gc ← [m̂g (1 − β1t ) − β1 mt−1 ]/(1 − β1 ). 32: mt ← β1 mt−1 + (1 − β1 )gc . 33: vt ← β2 vt−1 + (1 − β2 )gc2 . 34: v̂g ← vt /(1 − β2t ). p 35: Update network weights: θt ← θt−1 − γ m̂g /( v̂g + ϵ). 36: end for

3.5. Fast fine-tuning. 3.6. Convergence analysis. We establish the convergence guarantees of HRGrad for both convex and non-convex settings. We first derive the key geometric properties arising from the four-stage construction, and then exploit them to prove descent

Z.Y. LIANG

lemmas and convergence theorems. Recall from Stage 2 and Stage 4 that the output of the rotation and magnituderestoration process yields ( ∥gi ∥ ri (αi∗ ), if ḡi⊤ d∗ < 0, (3.13) girot = gi , otherwise, where ri (αi∗ ) = cos(αi∗ )ḡi + sin(αi∗ )wi is the SO(2) isometric rotation with αi∗ ∈ [0, π/2], ḡi = U(gi ), and wi = U d∗ − (ḡi⊤ d∗ )ḡi . Lemma 3.1 (Geometric Properties of HRGrad). Let {girot }m i=1 and gHRGrad be rot defined in Eqs. (3.13), (3.11)–(3.12), and let M = [U (g1rot ), . . . , U (gm )] ∈ RD×m have full column rank. Then: (i) (Isometric Fidelity) ∥girot ∥ = ∥gi ∥ for all i = 1, . . . , m. (ii) (Non-Conflict) (girot )⊤ gHRGrad ≥ 0 for all i, with strict inequality whenever ∥gHRGrad ∥ > 0. (iii) (Equal Cosine Similarity) There exists a uniform constant Sc > 0 such that (girot )⊤ gHRGrad = Sc , ∥girot ∥ ∥gHRGrad ∥

(3.14)

(iv) (Aggregate Product Identity)

m X

girot

∀ i = 1, . . . , m. ⊤

gHRGrad = ∥gHRGrad ∥2 .

i=1 ∗ Proof. (i) qFor conflicting tasks: ri (αi ) satisfies ḡi ⊥ wi and ∥ḡi ∥ = ∥wi ∥ = 1, so ∥ri (αi∗ )∥ = cos2 (αi∗ ) + sin2 (αi∗ ) = 1. By Stage 4 magnitude restoration girot = ∥gi ∥ ri (αi∗ ), hence ∥girot ∥ = ∥gi ∥. For non-conflicting tasks girot = gi by Eq. (3.13). (ii) By Stages 1–3, every girot is deflected into the Harmonized Cone H = K ∩ ∗ ∗

K , and the anchor d is the normalized centroid of the extreme rays of H, so the fair aggregation direction gu ∈ Int(H) ⊆ K∗ . The dual cone property of K∗ gives (girot )⊤ gu ≥ 0 for all i. From Eq. (3.12): (girot )⊤ gHRGrad =

(3.15)

m X (gjrot )⊤ gu · (girot )⊤ gu ≥ 0. | {z } j=1 ≥0 | {z } ≥0

Equality holds iff (girot )⊤ gu = 0 for all i, i.e., ∥gHRGrad ∥ = 0. † ⊤ ⊤ (iii) Since M has full column rank, (M ) M (M ⊤ M )−1 1m , and (3.16)

M ⊤ gu =

= M (M M )−1 .

M ⊤ M (M ⊤ M )−1 1m 1m = M (M ⊤ M )−1 1m M (M ⊤ M )−1 1m

Hence gu ∝

=: Sc 1m ,

so U (girot )⊤ gu = Sc > 0 uniformly in i. Because gHRGrad ∝ gu , the cosine similarity in Eq. (3.14) is constant for all i. (iv) Summing over all i and applying (iii): (3.17)

m X

girot

⊤

gHRGrad =

i=1

m X

∥girot ∥ ∥gHRGrad ∥ Sc = ∥gHRGrad ∥2 ,

i=1

rot ⊤ where ∥gHRGrad ∥ = i (gi ) gu ∥gu ∥ = rot ⊤ rot each (gi ) gu = ∥gi ∥ Sc .

rot i ∥gi ∥ Sc

follows from ∥gu ∥ = 1 and

HRGRAD

Remark 3.2. Isometric Fidelity P ensures ∥gHRGrad ∥ = 0 if and only if ∥gi ∥ = 0 for all i, since ∥gHRGrad ∥ = Sc i ∥gi ∥ with Sc > 0. This prevents the spurious premature stagnation caused by gradient cancellation in projection-based methods (where ∥P(gi )∥ → 0 even when ∥gi ∥ ≫ 0), guaranteeing that the optimizer halts only at genuine joint stationary points of all task losses. To formalize Pm the descent analysis, we introduce the proxy aggregate objective 1 L̂(θ) := m i=1 L̂i (θ), where each L̂i is the locally linear surrogate with ∇θ L̂i (θ) = rot gi (θ). By Lemma 3.1(i), L̂i and Li share the same critical points: ∇θ L̂i = 0 ⇔ ∇θ Li = 0. Theorem 3.3 (Convergence in the Convex Setting). Assume (a) the proxy m objectives Pm {L̂i }i=1 are convex and differentiable; (b) the aggregate proxy gradient ĝ := i=1 girot is L-Lipschitz continuous with constant L > 0. Then updating along gHRGrad with step size γ ≤ L2 guarantees Lγ (3.18) L̂(θ+ ) ≤ L̂(θ) − γ 1 − ∥gHRGrad ∥2 ≤ L̂(θ), 2 converging to either a location where ∥gHRGrad ∥ = 0 or the global optimum. Proof. The L-Lipschitz gradient condition yields the quadratic upper bound: (3.19)

L̂(θ+ ) ≤ L̂(θ) + ĝ ⊤ (θ+ − θ) +

L + ∥θ − θ∥2 . 2

Substituting θ+ = θ − γgHRGrad : (3.20)

L̂(θ+ ) ≤ L̂(θ) − γ ĝ ⊤ gHRGrad +

Lγ 2 ∥gHRGrad ∥2 . 2

Applying the Aggregate Product Identity (Lemma 3.1(iv)), ĝ ⊤ gHRGrad = ∥gHRGrad ∥2 . Inserting into Eq. (3.20): Lγ (3.21) L̂(θ+ ) ≤ L̂(θ) − γ 1 − ∥gHRGrad ∥2 . 2 When γ ≤ L2 the coefficient (1 − Lγ 2 ) ≥ 0, so Eq. (3.18) holds. Equality requires ∥gHRGrad ∥ = 0; by convexity this is the global optimality condition. Theorem 3.4 (Convergence in the Non-Convex Setting). Assume the proxy P(a) m rot objectives {L̂i }m are differentiable but possibly non-convex; (b) ĝ = g is Li=1 i=1 i 1 Lipschitz continuous. Then updating along gHRGrad with step size γ ≤ L satisfies the ergodic bound: 2 m X 2 L̂(θ0 ) − L̂∗ k (3.22) min gi ≤ , 1≤k≤K γ α2 K i=1 where gik = ∇θ Li (θk ), L̂∗ = inf θ L̂(θ), and α = min1≤k≤K Sck > 0 P is the minimum k equal cosine similarity across all iterations. Consequently, mink → 0 as i gi K → ∞. Proof. Applying the descent lemma with γ ≤ L1 and the Aggregate Product Identity (Lemma 3.1(iv)): (3.23)

L̂(θk+1 ) ≤ L̂(θk ) −

γ k ∥g ∥2 . 2 HRGrad

Z.Y. LIANG

By Isometric Fidelity (Lemma 3.1(i)) and Equal Cosine Similarity (iii): (3.24)

k ∥gHRGrad ∥=

m X

∥gik,rot ∥ Sck =

i=1

m X

(3.25)

L̂(θ

m X

i=1

where the final step uses the triangle inequality k+1

∥gik ∥ Sck ≥

i=1

i ∥gi ∥ ≥ ∥

m X γ gik ) ≤ L̂(θ ) − (Sck )2 2 i=1 k

gik Sck ,

i gi ∥. Therefore:

Telescoping from k = 1 to K and using L̂ ≥ L̂∗ : (3.26)

K X

(Sck )2 k=1

m X i=1

gik

≤

2 L̂(θ0 ) − L̂∗ . γ

2 Setting α = mink Sck > 0 and dividing P k by α K gives Eq. (3.22). As K → ∞ the righthand side vanishes, hence mink → 0, establishing convergence to a stationary i gi point.

4. Numerical results. In this section, we present the numerical results of the proposed HRGrad for several multiscale kinetic equations, including the Boltzmann– BGK equation, the linear transport equation, and two additional kinetic models with more challenging asymptotic structures. Besides the classical BGK and linear transport benchmarks, we further consider the ES-BGK equation and the linear semiconductor Boltzmann–Poisson equation. The former extends the standard BGK model by replacing the isotropic Maxwellian equilibrium with an ellipsoidal Gaussian depending on the stress tensor, thereby allowing the correct Prandtl number in the fluid limit. The latter extends the linear transport setting by incorporating a selfconsistent electric field, a velocity-space drift term, and a Poisson coupling, with the drift–diffusion–Poisson system as its asymptotic limit. These two additional examples are chosen to examine whether HRGrad remains effective when the multiscale conflict is caused not only by stiff relaxation, but also by high-order moment closure, anisotropy, self-consistent fields, and boundary layers. Extensive numerical results are presented for several problems chosen from kinetic regimes (ε ≈ O(1)) to hydrodynamic (diffusive) regimes (ε → 0), in order to compare our method and existing MTL methods. To verify and compare the performance of existing MTL methods and our method, we present both 1D and 2D numerical results for several problems chosen from rarefied regimes (ε ≈ O(1)) to hydrodynamic (diffusive) regimes (ε → 0). To validate the effectiveness of our method, we conduct all numerical experiments with the initial and boundary conditions kept consistent with those of APNNs [15, 16]. Furthermore, we consider additional solution tasks for Knudsen numbers ε ∈ (0, 1]. All experiments are conducted on an Nvidia RTX-3090-24GB GPU. Code and data for the following experiments are available at https://github.com/liangzhangyong/HRGrad. 4.1. MTL baselines. We compared the performance of the proposed HRGrad method with ten popular MTL baselines. Besides PCGrad and IMTL-G from before, we also compare to Linear scalarization baseline (LS) [51], Uncertainty Weighting (UW) [18], Dynamic Weight Average (DWA) [27], Gradient Sign Dropout (GradDrop) [6], Conflict-Averse Gradient Descent (CAGrad) [23], Random Loss Weighting (RLW) [21], Nash bargaining solution for MTL (Nash-MTL) [30], and Fast Adaptive

HRGRAD

Multitask Optimization (FAMO) [22]. We use two metrics to measure the performance of different methods: the mean rank metric (M R, lower ranks being better) [30, 22], which represents the average rank across all tasks. An M R of 1 means consistently outperforming all others. In addition, we consider the average F1 score (F1 , larger is better) to measure the average performance on all tasks. MGDA. In MGDA [36], the gradient vectors are first normalized, and the update step is taken in the direction of the minimum-norm point within the convex hull formed by their normalized convex combination. We employed the ℓ2 normalization, proposed in MGDA, which has the advantage of providing a stable definition of the update direction. CAGrad. In CAGrad [23], the hyperparameter c ∈ [0, 1) constrains the update direction to remain within a certain distance from the average of the loss gradients, and we set c = 0.5 in our experiments. ConFIG. In ConFIG [26], the update direction is originally obtained by solving a least-squares problem to find a vector whose inner product with each loss gradient equals 1. However, this procedure produced NaN values in all five independent runs on NS2d-C, Volterra1d, and Volterra system. Therefore, in our implementation, instead of using the unstable least-squares method, we computed the pseudo-inverse of the loss gradient matrix and multiplied it by 1m to obtain the update direction instead. 4.2. Experiment setting. Some settings for the numerical experiments are stated as follows. Our method is a gradient preconditioner that can be embedded into existing MTL methods. The Adam optimizer of the gradient descent algorithm is commonly used for MTL optimization of the task loss functions. The initial value of the parameters in all numerical experiments is generated by Xavier initialization. All the hyperparameters are chosen for the best performance after trying these experiments. We adopt the APNNs architecture and automatically satisfy the constraint settings in each task. Meanwhile, we adopt the meta-auto-decoder (MAD) method to construct the MTL framework for faster convergence speed without losing accuracy compared to other MTL methods across tasks. The activation function we used is σ(x) = x/(1 + exp(−x)) for BGK problems and σ(x) = tanh(x) for linear transport problems. In particular, the spatial domain of interest, denoted as D, includes the interval [0, 1] for the linear transport equation and [−0.5, 0.5] for the Boltzmann-BGK equation. Additionally, the parameter V is assigned a value of 10. The number of quadrature points is 30 for the linear transport equation and 64 for the Boltzmann-BGK equation. To train the networks, the Adam [19] version of the gradient descent method is used to solve the optimization problem with Xavier initialization. In practice, we need to tune the hyperparameters, such as neural network architecture, learning rate, and batch size, to obtain a good level of accuracy [28]. As a matter of experience, one can tune the weights of loss terms to equalize them, and a decreasing annealing schedule for the learning rate is used to achieve better numerical performance. We use an exponential decay strategy for an initial learning rate η0 = 10−3 with a decay rate of γ = 0.96 and a decay step of p = 200 iterations (4.1)

ηi = η0 · γ ⌊ p ⌋ ,

where, the variable i represents the current i−th iteration step, and the symbol ⌊·⌋ denotes the floor function. The reference solutions are obtained by standard finite difference methods. For most of the time we will check the relative ℓ2 error of the density ρ(x) between our

Z.Y. LIANG

method and reference solutions, e.g. for 1d case, sP 2 j |ρθ,j − ρj | P . (4.2) error := 2 j |ρj | Multiscale sampling. To enable effective pretraining across the full spectrum of Knudsen numbers ε ∈ [εmin , 1] with εmin > 0 sufficiently small (typically εmin = 10−6 or 10−4 ), a carefully designed sampling strategy for the scale parameter ε is crucial. Training difficulty is strongly heterogeneous across regimes. The small-ε regime is significantly harder: the equations become stiff, gradients are larger and more sensitive, and the network must accurately capture rapid relaxation while preserving macroscopic evolution and resolving thin layers. Naive uniform sampling in the linear scale, i.e., ε ∼ U(εmin , 1),

(4.3)

severely undersamples the difficult small-ε region. Such a sampling strategy biases training toward the transport-dominated regime, as the small-ε interval is effectively ignored in a linear scale. This imbalance severely impedes the convergence and triggers failure modes, particularly in the challenging continuum limit. To mitigate this, we adopt a log-uniform (log-scale) sampling strategy, which naturally assigns higher sampling density to small ε. Specifically, we introduce an auxiliary variable ξ uniformly distributed over the logarithmic range (4.4) ξ ∼ U log εmin , 0 , The scale parameter is then obtained via exponential mapping: (4.5)

ε = exp(ξ),

The resulting probability density function for ε is (4.6)

p(ε) =

1 1 = , ε (log 1 − log εmin ) ε log(1/εmin )

ε ∈ [εmin , 1].

Thus, the sampling density is proportional to 1/ε, ensuring exponentially higher sampling frequency in decades where stiffness and asymptotic complexity are greatest. In practice, during each training iteration, we independently sample a batch of b Nb values {εk }N k=1 according to this log-uniform distribution: (4.7) εk = exp ξk , ξk ∼ U log εmin , 0 , k = 1, . . . , Nb . For each sampled εk , we generate the corresponding task-specific collocation points in the domain as well as boundary and initial condition points, and compute the per-task residual losses using the APNNs formulation. The resulting task gradients {gk } are then aggregated and corrected via HRGrad to produce the final parameter update. This log-uniform strategy aligns with the multiplicative nature of multiscale transitions in kinetic theory and significantly improves training stability, convergence speed, and accuracy in the small-ε regime. It reduces the occurrence of failure modes observed under linear uniform sampling and enhances extrapolation capability to unseen ε → 0+ . For further refinement, one may incorporate adaptive importance resampling: periodically estimate the difficulty of logarithmically spaced ε-bins via moving-average

HRGRAD

residual loss or gradient norm, and adjust the sampling distribution to up-weight currently hard regimes. However, the baseline log-uniform sampling already provides substantial robustness when combined with HRGrad, and we employ it as the default strategy throughout all experiments. 4.3. Boltzmann-BGK problems. Consider the Boltzmann-BGK equation (4.8)

∂t f + v · ∇ x f =

1 (M (U ) − f ) , ε

v ∈ R.

and recall that the system of the BGK model for APNNs is   ε (∂t f + v∂x f ) = M (U ) − f, ∂t U + ∇x · ⟨vmf ⟩ = 0, (4.9)   U = ⟨mf ⟩ . Initial thermodynamic equilibrium is considered for all tests, as follows |v − u0 |2 ρ0 exp − := f0 (x, v). (4.10) f (0, x, v) = 1 2T0 (2πT0 ) 2 For these tests, a classical Riemann-like problem with Sod like initial velocity data is considered, as (4.11)

ρ0 (t, x) = ρL ,

u0 (t, x) = uL ,

T0 (t, x) = TL ,

x ≤ 0,

ρ0 (t, x) = ρR ,

u0 (t, x) = uR ,

T0 (t, x) = TR ,

x > 0.

Furthermore, HRGrad encodes the density distribution variable fθ and the macro variables ρθ , uθ , Tθ using two separate channels and decodes them concurrently in the base task. In this problem, we impose hard constraints on the predictions to ensure that fθ remains positive and that ρθ , uθ , Tθ automatically satisfy the boundary constraints. The formulation is as follows   fθ (t, x, v) := ln 1 + exp(f˜θ (t, x, v)) > 0,      x−xL xR −x    ρθ (t, x) := ρLxR −xL · ρRxR −xL · exp ((x − xL )(xR − x) · ρ̃θ (t, x)) , (4.12) p xR − x x − xL  uθ (t, x) := (x − xL )(xR − x) · ũθ (t, x) + uL + uR ,    xR − x L xR − x L    xR −x x−xL   Tθ (t, x) := T xR −xL · T xR −xL · exp (x − xL )(xR − x) · T̃θ (t, x) . L R 4.3.1. 1D smooth problems. In this section, we perform a test on the 1D smooth problem [20] for ε ∈ (0, 1]. The time interval for the test is set to t ∈ (0, 0.1]. The spatial domain is defined as x ∈ (−0.5, 0.5) with periodic boundary conditions applied. For the microscopic velocity space, the computational domain is chosen as v ∈ (−10, 10). The initial condition is specified by a Maxwellian distribution with particular macroscopic moments, as (4.13)

ρ0 (x) = 1 + 0.5 sin(2πx),

u0 (x) = 0,

T0 (x) = 1 + 0.5 sin(2πx + 0.2).

In this experiment, we employ the trapezoidal rule with 257 points for the numerical evaluation of macroscopic moments.

Z.Y. LIANG

4.3.2. 1D Riemann problems. In this section, we conduct a test on the onedimensional Sod tube problem [20] for ε ∈ (0, 1]. The computational settings for this test largely mirror those used in Subsection 4.3.1, with the exceptions being the boundary condition, the number of collocation points, and the initial condition. For the spatial domain, we implement the free-flow boundary condition. The initial condition is represented by a local Maxwellian distribution with macroscopic moments (ρ0 , u0 , T0 ). Specifically, the values on the interval (−0.5, 0) are set to (1, 0, 1), and on [0, 0.5), they are (0.125, 0, 0.8). Given the challenge of approximating jump functions with neural networks, we opt for a smoothed version of the macroscopic moments (4.14) ρ0 (x) = 1.5+(0.625−1.5)H(x),

u0 (x) = 0,

T0 (x) = 1.5+(0.75−1.5)H(x).

where, H(x) = (1 + tanh(20x)) /2 serves as a smoothed approximation of the Heaviside function. Prediction

1.1 1.0 0.9

0.2

0.8 0.04

x=-0.50

0.06

0.10

0.02

x=-0.30

0.04

x=-0.10

0.06

0.000875

1.4850 1.4800 0.00

0.02

0.04

0.06

0.08

0.10

0.00

0.02

0.04

t=0.00

0.06

0.08

1.25

0.10

(t, x)

0.04

0.06

0.04

t=0.02

0.06

0.08

0.10

0.625 0.00

0.02

0.04

t=0.04

0.06

0.08

0.10

4 2

0.626 0.02

0.00

0.02

0.04

t=0.06

0.06

0.08

0.10

0.00

1.2

0.8 0.6

0.4

0.2

0.0 x

0.2

0.6

0.4

0.2

0.0 x

0.2

0.6

0.4

1.0

0.4

0.2

0.0 x

0.2

0.6

0.4

1.0

0.4

0.2

0.0 x

0.2

0.6

0.4

0.06

0.08

0.10

1.0 0.8

0.8

(t, x)

1.4

(t, x)

1.4

(t, x)

1.4

(t, x)

1.4

1.0

0.04

t=0.10

1.4

1.0

0.02

t=0.08

1.4

1.0

0.10

x=0.50

0.628 0.627

0.65 0.00

0.08 1e 6+6.25e 1 8

0.629

0.75 0.70

1.30

1.4825

0.000850

1.35

1.45 1.30 1.15 1.00 0.85 0.70 0.55 0.40 0.25 0.10

0.630

0.80 (t, x)

(t, x)

1.4875

0.02

x=0.30

1.40

1.4900

0.000900

0.00

0.631

1.45

1.4925

0.000925

0.4

0.10

x=0.10

1.4950

0.000950

0.08

0.85

1.4975

0.000975

0.2

0.7

0.00

1.5000

0.001000

(t, x)

0.08

(t, x)

0.02

+1.499

(t, x)

0.4

0.7

0.00

(t, x)

0.4

0.0

0.8

(t, x)

0.9

0.2

1.2

0.0

1.0

0.2

1.3 (t, x)

1.1

0.0

1e 5

0.4

1.4

0.2

1.2

1.3

Absolute Error

1.5

0.4

1.4

0.2

error

Ground Truth 0.4

0.4

0.2

0.0 x

0.2

0.6

0.4

0.2

0.0 x

0.2

0.4

Fig. 8. HRGrad predicts the macroscopic density moments ρ of the 1D Riemann problem for ε = 0.1.

Ground Truth

Prediction

0.4

Absolute Error

0.2 0.4 0.04

x=-0.10

0.06

0.0001

0.01

0.0000

0.00 0.00

0.02

0.04

0.06

0.08

0.10

u(t, x)

0.1 0.0

t=0.00

0.02

0.04

0.06

0.08

0.10

0.00

0.02

0.04

t=0.02

0.30

0.04

0.06

0.08

0.00

0.0000

u(t, x)

0.3

0.15

0.00 0.4

0.2

0.0 x

0.2

0.4

0.02

0.04

0.06

0.08

0.10

0.2

0.4

0.2

0.0 x

0.2

0.4

error

0.08

0.10

0.00

0.3

0.4

0.2

0.0 x

0.2

0.4

0.06

0.08

0.10

t=0.10

0.2

0.0

0.04

0.5 0.4

0.2

0.02

t=0.08

0.3

0.1

0.4

0.06

0.4

0.0 0.2

0.3

0.0 0.0 x

0.0 0.04

0.4

0.1

0.2

0.5 0.02

0.5

0.1

0.4

1.0

0.00

0.4

0.05

0.04

1.5

t=0.06

0.10

0.02

2.0

0.0025

x=0.50

2.5

0.0075 0.0050

0.0

0.10

3.0 1e 5

0.0100

0.05 0.00

0.08

0.0125

0.10

t=0.04

0.20

0.00

0.0150

0.15

0.25

0.02

0.06

0.0175

0.20

0.10

0.04

x=0.30

0.25

0.00

1.2 0.02

x=0.10

0.30

0.2

0.02

0.0002

2.4

0.00

0.35 0.3

0.03

u(t, x)

0.0003

0.0

0.10

0.40

0.4

0.04

0.0004

0.08

u(t, x)

0.02

0.05

0.0005

3.6

0.4

0.00

x=-0.30

0.06

0.0006

0.0

0.10

u(t, x)

0.08

u(t, x)

x=-0.50

0.0007

0.06

u(t, x)

0.04

4.8

0.1

0.4 0.02

0.0 0.2

0.1

u(t, x)

0.0

0.2

u(t, x)

6.0

u(t, x)

0.3

0.2

0.00

7.2

0.2

0.3

9.6 8.4

0.4 0.2

0.0

1e 6

0.4

0.4 0.4

0.2

0.2 0.1 0.0

0.4

0.2

0.0 x

0.2

0.4

0.2

0.0 x

0.2

0.4

Fig. 9. HRGrad predicts the macroscopic velocity moments u of the 1D Riemann problem for ε = 0.1.

4.3.3. ES-BGK problems. To further examine the performance of HRGrad beyond the classical BGK relaxation model, we consider the ellipsoidal statistical BGK (ES-BGK) equation. Compared with the standard BGK model, the ES-BGK equation introduces an anisotropic Gaussian equilibrium depending on the stress tensor,

HRGRAD Ground Truth

Prediction 1.4

Absolute Error 1.4

0.4

1.2 1.0

0.0

0.8

0.0

0.6

0.2

0.6

0.2

0.4

0.04

x=-0.50

0.06

0.08

0.10

0.00

x=-0.30

0.04

x=-0.10

1.5

1.50

0.02

0.06

1.4 T(t, x)

1.35 1.30

1.3 1.2

1.3

1.1 1.0

1.2

1.25

T(t, x)

1.40

0.9

1.20

1.1

1.15 0.00

0.02

0.04

0.06

0.08

0.10

0.8 0.00

0.02

0.04

t=0.00

0.06

0.08

0.7

0.10

0.00

0.02

0.04

t=0.02

0.06

0.08

0.10

0.00

0.02

x=0.10

1.4

1.45

0.08

0.04

0.06

x=0.30

0.08

0.7

0.6

0.5

0.4

0.3

0.10

0.00

0.02

0.04

t=0.04

0.06

0.08

0.10

x=0.50

0.7

T(t, x)

0.02

1.08 0.96 0.84 0.72 0.60 0.48 0.36 0.24 0.12 0.00

error

T(t, x)

0.8

0.2

0.00

T(t, x)

0.2

1.0

T(t, x)

0.0

1e 5

0.4

1.2

0.2

T(t, x)

0.2

0.4

0.5 0.4 0.3

0.00

0.02

0.04

t=0.06

0.06

0.08

0.10

0.00

0.02

0.04

t=0.08

0.06

0.08

0.10

t=0.10

1.5 1.4

1.0

0.8

1.0 0.9

0.4

0.6

0.6 0.4

0.4

T(t, x)

0.6

T(t, x)

0.6

T(t, x)

1.1

T(t, x)

1.2

T(t, x)

1.3

0.8 0.4

0.2

0.0 x

0.2

0.4

0.2

0.4

0.2

0.0 x

0.2

0.4

0.2

0.4

0.2

0.0 x

0.2

0.4

0.2

0.4

0.2

0.0 x

0.2

0.4

0.2

0.4

0.2

0.0 x

0.2

0.4

0.2

0.4

0.2

0.0 x

0.2

0.4

Fig. 10. HRGrad predicts the macroscopic velocity moments T of the 1D Riemann problem for ε = 0.1.

which enables the model to recover the correct Prandtl number in the Navier–Stokes– Fourier asymptotic regime. This makes the ES-BGK problem a natural higher-level benchmark for testing multiscale optimization methods, since the neural solver must simultaneously resolve the kinetic distribution, the macroscopic conservation laws, and the anisotropic moment structure. We consider the one-dimensional-in-space and two-dimensional-in-velocity ESBGK equation (4.15)

∂t f + v1 ∂x f =

τ (ρ, T ) (G[f ] − f ) , ε

x ∈ Ω ⊂ R,

v = (v1 , v2 ) ∈ R2 ,

where f = f (t, x, v) is the particle distribution function, ε is the Knudsen number, and G[f ] is the ellipsoidal Gaussian equilibrium. The macroscopic quantities are defined by Z Z Z 1 (4.16) ρ= f dv, ρu = vf dv, E= |v|2 f dv. 2 R2 R2 R2 The temperature tensor is given by Z 1 (v − u) ⊗ (v − u)f dv, (4.17) Θ= ρ R2 and the corrected tensor temperature is (4.18)

T = (1 − ν)T I + νΘ,

−

1 ≤ ν < 1. 2

The ellipsoidal Gaussian is then defined as ρ 1 ⊤ −1 (4.19) G[f ] = p exp − (v − u) T (v − u) . 2 det(2πT ) The parameter ν controls the Prandtl number through (4.20)

Pr =

1 . 1−ν

Z.Y. LIANG

In the hydrodynamic limit ε → 0, the solution approaches a local equilibrium and the leading-order macroscopic dynamics are governed by the compressible Euler system. When the first-order Chapman–Enskog correction is taken into account, the ES-BGK model is consistent with the Navier–Stokes–Fourier equations with an adjustable Prandtl number. Therefore, this example allows us to test whether HRGrad can handle not only the kinetic-to-fluid transition, but also the optimization conflict induced by anisotropic stress and heat-flux corrections. We adopt a Riemann-type benchmark with left and right macroscopic states √ (ρR , uR , TR ) = (1, 0, 1.05), (4.21) (ρL , uL , TL ) = (1, M 2, 1), where M = 2.5. The initial distribution is taken as the corresponding local Maxwellian generated from these macroscopic states. The computational domain is chosen as (4.22)

x ∈ [−0.5, 0.5],

v ∈ [−V, V ]2 ,

with V = 8 unless otherwise specified. The representative Knudsen numbers are selected as (4.23) ε ∈ 5 × 10−1 , 10−1 , 10−2 , 10−3 , and the multiscale training version samples ε from the same log-uniform distribution used in the other experiments. In the neural approximation, we use separate network channels for the distribution function fθ and the macroscopic variables Uθ = (ρθ , uθ , Tθ ). To improve stability, the positivity of fθ , ρθ , and Tθ is enforced through softplus or exponential parameterizations. Since the ES-BGK equilibrium depends on the tensor T , we additionally impose a stress-consistency loss between the stress tensor computed from fθ and the auxiliary macroscopic stress representation. The loss contains the kinetic residual, the conservation-law residual, the moment-consistency residual, the ES equilibrium residual, and the initial and boundary losses. These loss components are treated as different tasks and are aggregated by HRGrad. The main diagnostic quantities are the density ρ, velocity u, temperature T , heat flux Z 1 |v − u|2 (4.24) Q1 (t, x) = (v1 − u1 )f dv, ε R2 2 and the anisotropy indicator (4.25)

A(t, x) =

∥Θ − T I∥F . T

For large ε, the HRGrad prediction is compared with a high-resolution discretevelocity ES-BGK solver. For small ε, we also compare the macroscopic quantities with the corresponding fluid-limit reference solution. This setup provides a more stringent test than the standard BGK equation, because the method must preserve both the stiff relaxation structure and the anisotropic moment information across different kinetic regimes. 4.4. Linear transport problems. We consider representative slab-geometry transport benchmarks from the rarefied regime (ε ≈ O(1)) to the diffusive regime (ε → 0). To keep the presentation focused, we organize them into three groups: a

HRGRAD

smooth periodic benchmark, a family of one-dimensional non-periodic stress tests, and a two-dimensional smooth benchmark. Together they cover different initial-boundary conditions, heterogeneous coefficients, boundary-layer effects, and dimensional extension. Consider the linear transport equation Z 1 1 1 f dv ′ − f , (4.26) ε∂t f + v∂x f = ε 2 −1 and recall that the even-odd decomposition of this equation is  2 ε ∂t r + ε2 v∂x j = ρ − r,    ε2 ∂ j + v∂ r = −j, t x (4.27)  ∂ ρ + ⟨v∂  t x j⟩ = 0,   ρ = ⟨r⟩ . 4.4.1. Smooth initial data with periodic BC. We start from the rarefied regimes where ε = 1 to the diffusive regime where ε → 0, and consider periodic boundary conditions with a smooth initial data as follows ρ(x) v2 f0 (x, v) = √ e− 2 , 2π

(4.28) where (4.29)

ρ(x) = 1 + cos(4πx).

The source term, scattering, and absorbing coefficients are set as (4.30)

σS = 1,

σA = 0,

Q = 0,

ε ∈ (0, 1].

To enhance numerical performance, exact periodic boundary conditions are enforced. The approach relies on a Fourier basis, where the transform (4.31)

T : x → {sin(2πjx), cos(2πjx)}kj=1

is applied prior to the first layer of the DNNs, as outlined in [11, 29]. In this paper, we set k = 8 to strike a balance between computational efficiency and the accuracy of the Fourier representation. 4.4.2. 1D non-periodic transport problems. We next collect three representative one-dimensional non-periodic tests, which stress HRGrad under boundarydriven transport, spatially varying scattering, and kinetic boundary layers. Isotropic in-flow boundary condition. We first consider isotropic in-flow boundary conditions given by (4.32)

x ∈ [0, 1],

FL (v) = 1,

FR (v) = 0.

The initial condition is set as f0 (x, v) = 0. The source term, scattering, and absorbing coefficients are defined as (4.33)

Q = 0,

σS = 1,

σA = 0,

ε ∈ (0, 1],

The results are illustrated in Figure, where the exact value of ⟨g⟩ = 0 is achieved. To enhance numerical performance, ρθ can be further designed to inherently satisfy the initial condition, which is given by (4.34)

ρθ (t, x) := t · exp (−ρ̃θ (t, x)) ≈ ρ(t, x).

Z.Y. LIANG

Variable scattering coefficient. Let (4.35)

x ∈ [0, 1],

FL (v) = 1,

FR (v) = 0,

and initial condition f0 (x, v) = 0. The source term, scattering, and absorbing coefficients are set as (4.36)

Q = 1,

σS = 1 + (10x)2 ,

σA = 0,

ε ∈ (0, 1].

In Figure, we report the numerical solution at time t = 0.0, 0.1, 0.2. In this problem, we have a source term and the scattering cross section that depend on x, so the scaling term σS /ε ranges from 1/ε → O(1), a problem with mixing scales. The numerical results show reasonably good performance of the proposed HRGrad method. Boundary-layer case. Let (4.37)

x ∈ [0, 1],

FL (v) = 5 sin(v),

FR (v) = 0,

and initial condition f0 (x, v) = 0. The source term, scattering, and absorbing coefficients are set as (4.38)

Q = 0,

σS = 1,

ε ∈ (0, 1].

σA = 0,

where since FL depends on v, there is a boundary layer near x = 0. 4.4.3. 2D smooth problems. Consider a 2D smooth problem from the Rarefied regime to the diffusive regime with (4.39)

Γ = [0, 1] × [0, 1],

FB (x, v) = 0,

n · v < 0,

x ∈ ∂Γ,

and initial condition f0 (x, v) = 0. The source term, scattering, and absorbing coefficients are set as (4.40)

Q = 1,

σS = 1,

ε ∈ (0, 1].

σA = 0,

were n denotes the exterior unit normal vector on ∂Γ. 4.4.4. Linear semiconductor Boltzmann–Poisson problems. We next consider a linear semiconductor Boltzmann–Poisson problem, which can be viewed as a natural extension of the linear transport equation with a self-consistent electric field. Compared with the pure linear transport benchmarks, this problem contains an additional velocity-space drift term, a Poisson equation for the electrostatic potential, contact boundary layers, and a drift–diffusion asymptotic limit. It therefore provides a more challenging multiscale test for HRGrad. We consider the one-dimensional-in-space and one-dimensional-in-velocity model

(4.41)

ε∂t f + v∂x f − E(t, x)∂v f =

1 Q(f ), ε

x ∈ (0, 1),

v ∈ R,

where f = f (t, x, v) is the electron distribution function, E is the electric field, and Q(f ) is a linear relaxation operator. In the numerical experiments, we use the relaxation-time approximation Z (4.42) Q(f ) = ρM − f, ρ(t, x) = f (t, x, v) dv, R

HRGRAD

where M (v) is a normalized Maxwellian. The electric potential ϕ and the electric field E are determined by the Poisson equation (4.43)

−β∂xx ϕ = ρ − c(x),

E = −∂x ϕ,

where β denotes the scaled Debye length and c(x) is the doping profile. In the diffusive scaling, as ε → 0, the distribution approaches local equilibrium, (4.44)

f = ρM + O(ε),

and the limiting macroscopic model is the drift–diffusion–Poisson system (4.45)

∂t ρ = ∂x (D∂x ρ + ηρE) ,

−β∂xx ϕ = ρ − c(x),

E = −∂x ϕ.

Here D is the diffusion coefficient and η is the mobility coefficient. This asymptotic structure is analogous to the diffusion limit of the linear transport equation, but the additional electric-field coupling introduces stronger task interaction and more severe optimization stiffness. We use a smooth n+ − n − n+ diode-type doping profile, C+ − C n x1 − x c(x) = Cn + 1 + tanh 2 s (4.46) C+ − C n x − x2 + 1 + tanh , 2 s where x1 = 0.3, x2 = 0.7, and s = 0.02. Unless otherwise stated, we set (4.47)

β = 2 × 10−3 ,

Vb = 5,

and impose the Dirichlet boundary condition for the electric potential (4.48)

ϕ(t, 0) = 0,

ϕ(t, 1) = Vb .

For the kinetic distribution, the inflow boundary conditions are prescribed as (4.49)

f (t, 0, v) = FL (v),

v > 0,

f (t, 1, v) = FR (v),

v < 0,

where FL and FR are Maxwellian inflow data associated with the contact densities. The initial condition is chosen as (4.50)

f (0, x, v) = c(x)M (v).

The representative Knudsen numbers are (4.51)

ε ∈ {1, 10−1 , 10−2 , 10−3 }.

For the multiscale version, ε is sampled from the log-uniform distribution on [εmin , 1], as described in the multiscale sampling paragraph. This is particularly important for the Boltzmann–Poisson problem, since the small-ε regime is dominated by the drift–diffusion balance and the Poisson coupling, whereas the kinetic regime requires resolving the velocity-space non-equilibrium induced by the electric field. For the neural representation, we use a micro–macro form (4.52)

fθ (t, x, v) = ρθ (t, x)M (v) + εgθ (t, x, v),

⟨gθ ⟩ = 0,

Z.Y. LIANG

together with an additional potential head ϕθ (t, x). The electric field is then computed by (4.53)

Eθ (t, x) = −∂x ϕθ (t, x).

The total loss contains the kinetic residual, the micro–macro consistency residual, the Poisson residual, the drift–diffusion limiting residual, and the initial and boundary losses. The Poisson residual is essential in this example, because an inaccurate electric field can lead to a wrong drift direction even if the density error is small. HRGrad is then applied to aggregate the gradients associated with the different ε-tasks and different physical residuals. The reference solutions are generated by a deterministic asymptotic-preserving solver based on parity or micro–macro decomposition, with Gauss–Hermite quadrature in velocity and finite differences in space. For ε = O(1), we compare the predicted density ρ, current J, electric field E, and potential ϕ with the kinetic reference solution. For ε ≪ 1, we additionally compare with the drift–diffusion–Poisson limit. The current is computed by Z (4.54) J(t, x) = vf (t, x, v) dv. R

The main error metrics are the relative ℓ2 errors of ρ, E, and ϕ, the Poisson residual, the mass conservation error, and the asymptotic-preserving error (4.55)

eAP (t) =

∥fθ − ρθ M ∥L1x,v ∥fθ ∥L1x,v

This example directly tests whether HRGrad can maintain stable training when the kinetic equation, the macroscopic limit equation, and the self-consistent field equation impose competing gradient directions. For the additional ES-BGK and semiconductor Boltzmann–Poisson experiments, the same multiscale training protocol is adopted. In particular, ε is treated as a task parameter and sampled by the log-uniform strategy. The loss functions are decomposed into kinetic residuals, macroscopic limiting residuals, moment or field consistency residuals, and initial-boundary residuals, so that HRGrad can explicitly correct the gradient conflict among different scales and physical constraints. 5. Conclusions. In this study, we have developed a harmonized rotational gradient method (HRGrad) for training the model to simultaneously solve both the Boltzmann-BGK equation and the linear transport equation across all ranges of Knudsen numbers. This is also a multi-scale problem where the Knudsen number gradually decreases, reflecting a transition from large-scale to small-scale phenomena, leading to a shift from free molecular flow to a continuum flow state. The computational challenge inherent in this task arises from the multi-scale nature of the kinetic equation and the simultaneous handling of different scale-solving issues. Existing methods for addressing such problems often experience gradient conflicts, making effective training difficult. To address this, we propose the harmonized rotational gradient method (HRGrad), which not only corrects the gradient magnitude but also adjusts the gradient direction, balancing the task losses for different Knudsen numbers by ensuring each task’s loss decreases at approximately the same rate. Additionally, we employ APNNs to solve these equations using macro-micro and even-odd decomposition for

HRGRAD

each task, addressing the small-scale challenges for PINNs. Through a series of numerical tests, we demonstrate that HRGrad is capable of predicting solutions for different Knudsen numbers of the Boltzmann-BGK equation and the linear transport equation, even for unseen Knudsen numbers. In addition, there is still room for improvement on the complexity of the proposed HRGrad method for solving high-dimensional kinetic problems, along with the computational cost of handling complex collisions and high-dimensional integrals, which we leave for future work. In future work, we aim to integrate a neural sparse representation method [20] to efficiently approximate the Maxwellian distribution within the BGK and quadratic collision models, thereby enhancing HRGrad’s computational efficiency for each task. Furthermore, we propose incorporating a separable physicsinformed neural network (SPINN)-based method [31] to address high-dimensional Boltzmann-BGK problems, enabling HRGrad to handle even more intricate highdimensional challenges effectively. Acknowledgments. We would like to thank the referees for their insightful comments that greatly improved the paper. The computing for this work was supported by the High Performance Computing Platform at Eastern Institute of Technology, Ningbo. Appendix A. Convergence analysis. In this appendix, we provide the formal proof of global convergence for both convex and non-convex landscapes for the proposed HRGrad method, together with two supplementary analyses: per-task gradient descent compatibility (how the isometric rotation preserves individual task optimization) and MCAE angular stability (how the MER objective achieves the minimal necessary rotation for conflict resolution). All four HRGrad stages contribute explicitly to the results below. Notation. Throughout this appendix, gi = ∇θ Li (θ) denotes the i-th task gradient, ḡi = U(gi ) its unit vector, αi∗ ∈ [0, π/2] the optimal rotation angle from the MER objective (Eq. (3.8)), and girot the isometrically rotated counterpart defined in Eq. (3.13). We write K = {Gλ | λ ∈ Rm + } for the Primal Gradient Cone, K∗ = {y ∈ RD | G⊤ y ≥ 0m } for the Dual Gradient Cone, and H = K ∩ K∗ for the Harmonized Cone. The convergence analysis rests on two geometric properties of HRGrad that follow directly from Lemma A.3 and one regularity condition on the alignment between the original task gradients and the fair consensus direction. (P1) Equal cosine similarity (Lemma A.3(iii)) There exists a constant Sc > 0 such that U (girot )⊤ gu = Sc for all i = 1, . . . , m. rot (P2) Magnitude conservation (Lemma Pm A.3(i),(iii)) ∥gi ∥ = ∥gi ∥ for all i, and consequently ∥gHRGrad ∥ = Sc i=1 ∥gi ∥. Assumption A.1 (Strict Common Descent). At every iterate θ, the original task gradients satisfy gi⊤ gu ≥ ρi ∥gi ∥ for some constants ρi > 0, i = 1, . . . , m. Set ρmin = min1≤i≤m ρi > 0. ∗ Remark A.2 (Sufficient condition for Assumption A.1). Let αmax = maxi αi∗ be rot the largest MER rotation angle. The chord-length identity ∥gi −gi ∥ = 2∥gi ∥ sin(αi∗ /2) and the Cauchy–Schwarz inequality give m m m m X X X X α∗ ⊤ rot ⊤ rot ⊤ max ∥gi ∥, (A.1) gi gu = (gi ) gu − (gi − gi ) gu ≥ Sc − 2 sin 2 i=1

i=1

so Assumption A.1 holds with ρmin =

∗ Sc − 2 sin(αmax /2)

∗ > 0 whenever αmax <

Z.Y. LIANG

2 arcsin(Sc /2). In the conflict-free case αi∗ = 0, this gives ρmin = Sc and κ = 1 (see Corollary A.5). Under Assumption A.1, we define the manifold mapping coefficient (A.2)

κ :=

ρmin ∈ (0, 1], Sc

which equals 1 precisely when all task gradients are non-conflicting. A.1. Complete Proof of the Geometric Lemma. Lemma A.3 (Geometric Properties of HRGrad). Let {girot }m i=1 and gHRGrad be rot defined in Eqs. (3.13), (3.11)–(3.12), and let M = [U (g1rot ), . . . , U (gm )] ∈ RD×m have full column rank. Then: (i) (Isometric Fidelity) ∥girot ∥ = ∥gi ∥ for all i. (ii) (Non-Conflict) (girot )⊤ gHRGrad ≥ 0 for all i, with strict inequality when ∥gHRGrad ∥ > 0. (girot )⊤ gHRGrad = Sc for all i. (iii) (Equal Cosine Similarity) ∃ Sc > 0 s.t. rot ∥g ∥ ∥gHRGrad ∥ Pm rot ⊤i (iv) (Aggregate Product Identity) gHRGrad = ∥gHRGrad ∥2 . i=1 gi Proof. (i) [Stage 2 isometry] For a conflicting task (ḡi⊤ d∗ < 0): the GramSchmidt step ensures ḡi ⊥ wi and ∥ḡi ∥ = ∥wi ∥ = 1, so the rotation preserves unit length: (A.3)

∥ri (αi∗ )∥2 = cos2 (αi∗ )∥ḡi ∥2 + sin2 (αi∗ )∥wi ∥2 = cos2 (αi∗ ) + sin2 (αi∗ ) = 1.

The Stage 4 magnitude restoration gives girot = ∥gi ∥ ri (αi∗ ), so ∥girot ∥ = ∥gi ∥. For a non-conflicting task (ḡi⊤ d∗ ≥ 0): girot = gi by definition (3.13), so the identity is trivially satisfied. (ii) [Harmonized Cone dual property] By Stages 1–3, every rotated gradient girot is deflected into H = K ∩ K∗ . Since gu is constructed as the normalized centroid of the extreme rays of H (Stage 4), we have gu ∈ Int(H) ⊆ K∗ . By the definition of the Dual Cone K∗ : gu ∈ K∗ =⇒ (girot )⊤ gu ≥ 0 ∀ i (since girot ∈ K). P rot ⊤ From Eq. (3.12), gHRGrad = j (gj ) gu gu : (A.4)

(girot )⊤ gHRGrad =

(A.5)

m X (gjrot )⊤ gu · (girot )⊤ gu ≥ 0. | {z } j=1 ≥0 {z } | ≥0

Equality holds iff j (gjrot )⊤ gu = 0, i.e., ∥gHRGrad ∥ = 0. (iii) [Pseudo-inverse fair aggregation] Since M has full column rank, the MoorePenrose pseudoinverse satisfies (M † )⊤ = M (M ⊤ M )−1 , so gu ∝ M (M ⊤ M )−1 1m . Computing the projected product: P

(A.6)

M ⊤ gu =

M ⊤ M (M ⊤ M )−1 1m 1m = M (M ⊤ M )−1 1m M (M ⊤ M )−1 1m

=: Sc 1m ,

Sc > 0.

Hence U(girot )⊤ gu = Sc uniformly for all i. Because gHRGrad ∝ gu , the cosine similarity (girot )⊤ gHRGrad /(∥girot ∥∥gHRGrad ∥) = Sc is constant across all tasks.

HRGRAD

(iv) [Summation identity] Using (iii), each cross-product is (girot )⊤ gHRGrad P = rot ⊤ rot ∥g ∥ ∥g ∥ S . Summing and using the magnitude formula ∥g ∥ = HRGrad c HRGrad i (gi ) gu = Pi rot i ∥gi ∥Sc (from ∥gu ∥ = 1 and (iii)): (A.7) m m m X ⊤ X X girot gHRGrad = ∥girot ∥ ∥gHRGrad ∥ Sc = ∥gHRGrad ∥· ∥girot ∥Sc = ∥gHRGrad ∥2 . i=1

i=1

}

= ∥gHRGrad ∥

A.2. Convergence for Convex Losses. Theorem A.4 (Convex Convergence of HRGrad). Suppose Assumption A.1 holds. Let each Li : RD → be convex and continuously differentiable, and let the PR m total gradient g = ∇θ L = i=1 ∇θ Li be L-Lipschitz continuous with constant L > 0. Then for any step size γ ≤ 2κ/L, the HRGrad update θ+ = θ − γgHRGrad satisfies (A.8)

Lγ ∥gHRGrad ∥2 ≤ L(θ). L(θ+ ) ≤ L(θ) − γκ 1 − 2κ

In particular, {L(θk )} is monotonically non-increasing, and the iterates converge to the global optimum of L or to a Pareto stationary point at which ∥gHRGrad ∥ = 0. Proof. Since g is L-Lipschitz, the descent lemma gives (A.9)

L(θ+ ) ≤ L(θ) + g ⊤ (θ+ − θ) +

L + ∥θ − θ∥2 . 2

Substituting θ+ − θ = −γgHRGrad : (A.10)

L(θ+ ) ≤ L(θ) − γg ⊤ gHRGrad +

Lγ 2 ∥gHRGrad ∥2 . 2

We establish a lower bound for g ⊤ gHRGrad . Write gHRGrad = ∥gHRGrad ∥gu and compute: g ⊤ gHRGrad = ∥gHRGrad ∥

m X

gi⊤ gu

i=1

(A.11)

≥ ∥gHRGrad ∥ ρmin

m X

∥gi ∥

(Assumption A.1)

i=1

= ∥gHRGrad ∥ ρmin · (A.12)

∥gHRGrad ∥ Sc

((P2))

= κ∥gHRGrad ∥2 .

Substituting (A.12) into (A.10): (A.13) L(θ+ ) ≤ L(θ) − γκ∥gHRGrad ∥2 +

Lγ Lγ 2 ∥gHRGrad ∥2 = L(θ) − γκ 1 − ∥gHRGrad ∥2 . 2 2κ

For γ ≤ 2κ/L the factor (1 − Lγ/(2κ)) ≥ 0, so L(θ+ ) ≤ L(θ) with equality iff ∥gHRGrad ∥ = 0. By (P2), ∥gHRGrad ∥ = 0 iff ∇θ Li (θ) = 0 for all i, which is the global optimum by convexity.

Z.Y. LIANG

Corollary A.5 (Conflict-free case). If all task gradients are non-conflicting (αi∗ = 0 for all i), then girot = gi , κ = 1, and the bound (A.8) tightens to L(θ+ ) ≤ L(θ) − γ(1 − Lγ/2)∥gHRGrad ∥2 with step size γ ≤ 2/L. Moreover, the Aggregate Product Identity (Lemma A.3(iv)) gives the exact equality g ⊤ gHRGrad = ∥gHRGrad ∥2 , so no alignment assumption is needed in this case. A.3. Convergence for Non-Convex Losses. Theorem A.6 (Non-Convex Convergence of HRGrad). Suppose Assumption A.1 holds. Let each Li be continuously differentiable (not necessarily convex), and let g = ∇θ L be L-Lipschitz continuous. Define L∗ = inf θ L(θ) > −∞ and α = min0≤k≤K−1 Sck > 0. Then for step size γ ≤ κ/L: (A.14)

min

0≤k≤K−1

∥∇θ L(θk )∥2 ≤

2[L(θ0 ) − L∗ ] K→∞ −−−−→ 0. γ κ α2 K

Consequently, every accumulation point of {θk } is a stationary point of L. Proof. Apply (A.10) at iteration k with the bound (A.12): (A.15)

Lγ k γκ k L(θk+1 ) ≤ L(θk ) − γ κ − ∥gHRGrad ∥2 ≤ L(θk ) − ∥g ∥2 , 2 2 HRGrad

where the P last step uses γ ≤ κ/L, giving κ − Lγ/2 ≥ κ/2. By (P2) and the triangle inequality i ∥gik ∥ ≥ ∥g k ∥: (A.16)

m X 2 k ∥gHRGrad ∥2 = Sck ∥gik ∥ ≥ (Sck )2 ∥g k ∥2 . i=1

Substituting (A.16) into (A.15) and telescoping from k = 0 to K − 1: K−1

(A.17)

γκ X k 2 k 2 (Sc ) ∥g ∥ ≤ L(θ0 ) − L(θK ) ≤ L(θ0 ) − L∗ . 2 k=0

Setting α = min0≤k≤K−1 Sck > 0 and dividing by γκα2 K/2: K−1

(A.18)

1 X k 2 2[L(θ0 ) − L∗ ] ∥g ∥ ≤ . K γκα2 K k=0

Taking the minimum on the left-hand side yields (A.14). As K → ∞ the bound vanishes, so lim inf k→∞ ∥g k ∥ = 0 and every accumulation point is a stationary point of L. Remark A.7 (Positivity of α). The rotation to H ensures Sck > 0 at each k prok,rot vided that the column rank condition on M = [U(g1k,rot ), . . . , U (gm )] in Lemma A.3 holds. Geometrically, this fails only if all rotated unit gradients become collinear, a codimension-(D − 1) event avoided generically during gradient flow. A.4. Per-Task Gradient Descent Compatibility. The above theorems esP tablish convergence on the total multi-task loss L = i Li via the manifold mapping coefficient κ = ρmin /Sc . A natural follow-up question is: does HRGrad’s isometric rotation preserve the optimization direction of each individual task loss? The next lemma answers this affirmatively.

HRGRAD

Lemma A.8 (Per-Task Rotation Compatibility). For each task i = 1, . . . , m, the isometric rotation Eq. (3.13) satisfies: gi⊤ girot = ∥gi ∥2 cos(αi∗ ) ≥ 0.

(A.19)

Proof. For a non-conflicting task: girot = gi , so gi⊤ girot = ∥gi ∥2 ≥ 0. For a conflicting task: gi⊤ girot = gi⊤ ∥gi ∥ ri (αi∗ ) = ∥gi ∥ gi⊤ cos(αi∗ )ḡi + sin(αi∗ )wi = ∥gi ∥ cos(αi∗ )gi⊤ ḡi + sin(αi∗ )gi⊤ wi . (A.20) Since gi = ∥gi ∥ḡi and wi ⊥ ḡi (Gram-Schmidt in Stage 2): gi⊤ ḡi = ∥gi ∥ and gi⊤ wi = ∥gi ∥ḡi⊤ wi = 0. Therefore gi⊤ girot = ∥gi ∥2 cos(αi∗ ). Because αi∗ ∈ [0, π/2], we have cos(αi∗ ) ≥ 0. Remark A.9. Lemma A.8 shows that the rotated gradient girot always lies in the same half-space as the original task gradient gi : gi⊤ girot ≥ 0. This guarantees that the HRGrad step −γgHRGrad does not reverse the individual task descent direction, even for the most severely conflicting tasks. In contrast, projection-based methods can produce P(gi ) with ∥P(gi )∥ → 0 when ϕ → π, erasing all task-specific information. Lemma A.10 (Proxy Loss Approximation Error). Let C ⊆ {1, . . . , m} denote ∗ = maxi∈C αi∗ . Then the index set of conflicting tasks and αmax P rot Pthe proxy gradient ĝ = i gi approximates the original aggregate gradient g = i gi with the error bound: ∗ X α ∥gi ∥. (A.21) ∥ĝ − g∥ ≤ 2 sin max 2 i∈C

Proof. Non-conflicting tasks contribute zero error since girot = gi . For a conflicting task i ∈ C: ∥girot − gi ∥2 = ∥∥gi ∥ ri (αi∗ ) − ∥gi ∥ḡi ∥2 = ∥gi ∥2 ∥ri (αi∗ ) − ḡi ∥2 ∗ α (A.22) = ∥gi ∥2 (cos αi∗ − 1)2 + sin2 αi∗ = 2∥gi ∥2 (1 − cos αi∗ ) = 4∥gi ∥2 sin2 2i . ∗ : By the triangle inequality and αi∗ ≤ αmax ∗ ∗ X X X α α (A.23) ∥ĝ − g∥ ≤ ∥girot − gi ∥ = 2 ∥gi ∥ sin 2i ≤ 2 sin max ∥gi ∥. 2 i∈C

i∈C

Remark A.11. When conflict is mild (αi∗ → 0 for all i), Eq. (A.21) gives ∥ĝ − g∥ → 0, multi-task loss Pso the proxy objective L̂ faithfully approximates the true 1 ∗ L= m L . The MER objective (Stage 3) explicitly minimizes α via the proximity i i i term ∥ri (αi∗ )−ḡi ∥2 = 2(1−cos αi∗ ), ensuring the smallest possible approximation error for a given conflict-resolution requirement. A.5. Angular Stability via the MER-MCAE Objective. We now analyze how the MER objective (Eq. (3.8)) determines the optimal rotation angles, and show that the resulting angular perturbation is locally stable. Lemma A.12 (MCAE Second-Order Angular Stability). The proximity term Pi (αi ) = ∥ri (αi ) − ḡi ∥2 = 2(1 − cos αi ) in the MER objective is strictly convex on [0, π/2]: (A.24)

∂ 2 Pi = 2 cos αi > 0, ∂αi2

αi ∈ [0, π2 ).

Z.Y. LIANG

Consequently, any deviation δαi from the optimal αi∗ incurs a quadratic MCAE penalty: (A.25)

Pi (αi∗ + δαi ) = Pi (αi∗ ) + Pi′ (αi∗ )δαi + cos(αi∗ ) (δαi )2 + O((δαi )3 ).

Proof. Direct computation: Pi (αi ) = 2(1 − cos αi ), Pi′ (αi ) = 2 sin αi , Pi′′ (αi ) = 2 cos αi . For αi ∈ [0, π/2), cos αi > 0, establishing strict convexity. The second-order Taylor expansion at αi∗ gives Eq. (A.25). Remark A.13. Lemma A.12 shows that: • Over-Correction (OCAE): setting αi > αi∗ incurs a penalty cos(αi∗ )(αi − αi∗ )2 > 0, quantifying the destruction of task-specific microscopic kinetic features. • Under-Correction (UCAE): setting αi < αi∗ similarly incurs a penalty, corresponding to insufficient rotation that leaves the gradient in conflict. (t) The adaptive step strategy αsteps (Stage 3) allocates more inner SGD iterations when the stiffness variance st is large, directly targeting a smaller δαi and hence a smaller quadratic MCAE penalty in stiff regimes. A.6. Convergence Rate and Task Scaling. We provide an explicit formula for the equal cosine similarity Sc and analyze how the convergence rate scales with the number of tasks m. Lemma A.14 (Explicit Formula for Sc ). Under the full column rank assumption on M , the equal cosine similarity is (A.26)

1 . Sc = p ⊤ ⊤ 1m (M M )−1 1m −1

. Let λ = (M ⊤ M )−1 1m ; then Proof. From Eq. (A.6), Sc = M (M ⊤ M )−1 1m ⊤ −1 ⊤ −1 ⊤ ⊤ −1 2 ⊤ ⊤ ⊤ 1m . Tak∥M λ∥ = λ M M λ = 1m (M M ) M M (M M ) 1m = 1⊤ m (M M ) ing the square root gives Eq. (A.26). Corollary A.15 (Convergence Rate Bound). If the m rotated unit gradients √ ⊤ {U(girot )}m i=1 are mutually orthonormal (i.e., M M = Im ), then Sc = 1/ m. In general, by the Cauchy-Schwarz inequality for positive definite matrices, Sc ≤ 1, with the non-convex convergence bound (Theorem A.6): (A.27) 2 m X √ 2m L̂(θ0 ) − L̂∗ k min gi (using α ≥ 1/ m in the orthonormal case). ≤ 1≤k≤K γK i=1 The rate thus scales as O(m/K) at most, degrading gracefully with the number of tasks. √ ⊤ ⊤ ⊤ −1 ⊤ Proof. √ For M M = Im : 1m (M M ) 1m = 1m 1m = m, so Sc = 1/ m. Setting α = 1/ m in the bound of Theorem A.6 gives Eq. (A.27). A.7. Comparison between HRGrad and baselines. A.7.1. HRGrad compared with PCGrad. PCGrad [49] resolves gradient conflicts by sequential orthogonal projection. Define the conflict-removal operator  ⊤ a − a b b, if a⊤ b < 0, (A.28) O(b, a) := ∥b∥2  a, otherwise,

HRGRAD

which removes the component of a conflicting with b. For each task i, P PCGrad m applies ĝi ← O(gj , ĝi ) over all j ̸= i in random order and sets gPCGrad = i=1 ĝi . The comparison with HRGrad turns on three structural deficiencies of PCGrad. First, although PCGrad is non-conflicting in the two-task case, this property does not extend ⊤ to general m. Indeed, for m = 2, a direct calculation gives gPCGrad g1 = ∥g1 ∥2 [1 − 2 Sc (g1 , g2 )] ≥ 0, confirming non-conflict. For m > 2, however, projecting ĝ1 onto g3⊥ after projecting it onto g2⊥ may reintroduce ĝ1⊤ g2 < 0; random task ordering may mitigate but cannot remove this failure mechanism. HRGrad avoids it by construction: every girot is deflected into H = K ∩ K∗ , so the dual-cone inclusion girot ∈ K∗ yields simultaneous pairwise non-conflict (girot )⊤ gjrot ≥ 0 for all i, j and the update-level bound (girot )⊤ gHRGrad ≥ 0 for all i (Lemma A.3(ii)). A second limitation is magnitude bias. Substituting (A.28) and using O(gi , gj ) ⊥ gi yields ⊤ (A.29) gPCGrad U(gi ) = ∥gi ∥ 1 − Sc2 (g1 , g2 ) , {i, j} = {1, 2}, ⊤ ⊤ so gPCGrad U(g1 ) / gPCGrad U(g2 ) = ∥g1 ∥/∥g2 ∥. Thus the update is biased toward the larger-norm task regardless of the conflict angle. By contrast, HRGrad satisfies the Equal Cosine Similarity property (Lemma A.3(iii)):

(A.30)

⊤ gHRGrad U(girot ) = 1, ⊤ gHRGrad U(gjrot )

∀ i ̸= j,

providing a norm-independent update direction critical in multiscale kinetic problems where ∥gmacro ∥ ≫ ∥gmicro ∥. Finally, PCGrad suffers from norm reduction under strong conflict. The projection satisfies q (A.31) ∥O(gj , gi )∥ = ∥gi ∥ 1 − Sc2 (gi , gj ) ≤ ∥gi ∥, with ∥O(gj , gi )∥ → 0 as Sc (gi , gj ) → −1, so ∥gPCGrad ∥ → 0 can arise far from any stationary point (spurious stagnation). In the multiscale setting where ∥gmicro ∥ ≪ ∥gmacro ∥, this norm collapse irreversibly suppresses high-frequency kinetic information. HRGrad eliminates this artifact through Isometric Fidelity (Lemma A.3(i)): (A.32)

∥girot ∥ = ∥gi ∥,

for all i, αi∗ ∈ [0, π/2],

and Lemma A.8 further guarantees gi⊤ girot = ∥gi ∥2 cos αi∗ ≥ 0, so the original task direction is never reversed. A.7.2. HRGrad compared with IMTL-G. IMTL-G [24] is designed to enforce equal projection lengths on the task-wise unit gradients, namely (A.33)

⊤ ⊤ gIMTL−G U(gi ) = gIMTL−G U(gj ),

∀ i, j = 1, . . . , m.

To achieve this, IMTL-G constructs a weighted combination of the original gradients,

(A.34)

gIMTL−G =

m X

αi gi ,

i=1

where the weights satisfy  ⊤ ⊤ [α2 , α3 , . . . , αm ] = g1⊤ UIMTL (DIMTL UIMTL )−1 ,    Pm  α1 = 1 − i=2 αi , (A.35) ⊤ ⊤ ⊤ ⊤  U  IMTL = [(U(g1 ) − U (g2 )) , (U(g1 ) − U (g3 )) , . . . , (U (g1 ) − U (gm )) ],   ⊤ DIMTL = [(g1 − g2 )⊤ , (g1 − g3 )⊤ , . . . , (g1 − gm )⊤ ].

Z.Y. LIANG

This objective is closer in spirit to HRGrad than PCGrad, since both methods seek an impartial update across tasks. The difference is geometric. IMTL-G enforces (A.33) by rescaling the original gradients themselves, whereas HRGrad first rotates each conflicting gradient into the harmonized cone and only then applies equal-cosine aggregation to the rotated set. Consequently, (A.33) controls only the final projection lengths; it does not alter the pairwise geometry of the aggregation inputs. HRGrad, by contrast, resolves pairwise conflicts before aggregation: girot ∈ H = K ∩ K∗ implies (girot )⊤ gjrot ≥ 0 for all i, j (Lemma A.3(ii)), while Lemma A.8 guarantees gi⊤ girot = ∥gi ∥2 cos αi∗ ≥ 0. For two tasks, the IMTL-G weights reduce to ∥g1 ∥ ∥g2 ∥ , , (A.36) [α1 , α2 ] = ∥g1 ∥ + ∥g2 ∥ ∥g1 ∥ + ∥g2 ∥ and therefore (A.37)

⊤ gIMTL−G U(g1 ) =

∥g1 ∥ ∥g2 ∥ 1 + Sc (g1 , g2 ) ≥ 0, ∥g1 ∥ + ∥g2 ∥

with the same identity holding for g2 . Thus IMTL-G is non-conflicting in the two-task setting. Its magnitude, however, is governed by the harmonic mean of ∥g1 ∥ and ∥g2 ∥: r (A.38)

∥gIMTL−G ∥ = 2

1 + Sc (g1 , g2 ) ∥g1 ∥ ∥g2 ∥ . 2 ∥g1 ∥ + ∥g2 ∥

Therefore, when ∥g1 ∥ ≫ ∥g2 ∥, the IMTL-G update magnitude tracks the smaller task gradient. In multiscale kinetic problems this may suppress the microscopic signal precisely when its norm is already much smaller than the macroscopic one. HRGrad avoids this collapse because it preserves every task norm through isometric rotation and then aggregatesPthe rotated gradients through the equal-cosine factor Sc , m yielding ∥gHRGrad ∥ = Sc i=1 ∥gi ∥ with Sc > 0 under the column-rank condition of Lemma A.3. The same two-task calculation also clarifies the connection with ConFIG: IMTL-G and ConFIG have the same update direction when m = 2, but their magnitudes differ by ∥gIMTL−G ∥ 2∥g1 ∥ ∥g2 ∥ = . ∥gConFIG ∥ (∥g1 ∥ + ∥g2 ∥)2

(A.39)

This ratio becomes small under strong norm imbalance, which explains why IMTLG, ConFIG, and HRGrad behave differently even when their two-task directions are closely related. A.7.3. HRGrad compared with ConFIG. ConFIG [26] constructs a conflictfree direction from the normalized gradient matrix M0 = [U (g1 ), . . . , U (gm )] through the pseudoinverse relation (A.40)

gConFIG =

m X

gi⊤ gu0 gu0 ,

gu0 = U M0 (M0⊤ M0 )−1 1m .

i=1

This is the same equal-cosine construction used in the original ConFIG formulation: the pseudoinverse determines a direction whose projections onto all normalized task gradients are equal, while the prefactor rescales the final update by the

HRGRAD

total projection length. By the Moore-Penrose identity, M0⊤ gu0 = Sc0 1m with Sc0 = ∥M0 (M0⊤ M0 )−1 1m ∥−1 > 0, so gi⊤ gConFIG = ∥gi ∥ ∥gConFIG ∥Sc0 > 0: ConFIG is nonconflicting with every gi and admits a convergence guarantee analogous to Theorems A.4–A.6 [26]. The structural relation with HRGrad is therefore explicit. ConFIG aggregates the original gradients through (A.41)

gConFIG =

m X

gi⊤ gu0 gu0 ,

gu0 = U M0 (M0⊤ M0 )−1 1m ,

i=1

whereas HRGrad uses the same pseudoinverse-based aggregation after rotating each conflicting gradient into the harmonized cone: (A.42)

gHRGrad =

m X

(girot )⊤ gu gu ,

† ⊤ gu = U (Mrot ) 1m ,

i=1

In the conflict-free case (αi∗ = 0, so girot = gi ), HRGrad reduces exactly to ConFIG (Mrot = M0 , Sc = Sc0 , gHRGrad = gConFIG ). HRGrad is therefore a conflict-aware extension of ConFIG, with SO(2) isometric rotation as the additional pre-processing stage. The key distinction between ConFIG and HRGrad lies in the treatment of gradient conflicts prior to aggregation. ConFIG applies the pseudoinverse aggregation directly to the original unit vectors U(gi ), which may be mutually conflicting. Although gConFIG is non-conflicting with each individual gi (gi⊤ gConFIG = ∥gi ∥ ∥gConFIG ∥ Sc0 > 0), no guarantee is placed on the pairwise relationships within the aggregation input itself: gu0 is computed from a potentially opposing configuration. HRGrad eliminates this limitation by resolving all pairwise conflicts before aggregation via Stages 1– 3. The dual-cone inclusion girot ∈ H ⊆ K∗ yields (girot )⊤ gjrot ≥ 0 for all i ̸= j (Lemma A.3(ii)), and Lemma A.8 guarantees gi⊤ girot = ∥gi ∥2 cos αi∗ ≥ 0, so the rotation never reverses any individual task gradient. This pre-resolution also determines Pm the numerical behavior under severe conflict. 0 ConFIG satisfies ∥g ∥ = S ConFIG c i=1 ∥gi ∥, whereas HRGrad satisfies ∥gHRGrad ∥ = Pm Sc i=1 ∥gi ∥; their stability diverges when gradients are nearly anti-parallel. For m = 2 with g1 ≈ −g2 , the ConFIG input matrix satisfies M0 ≈ [ḡ1 , −ḡ1 ] so that 1 −1 ⊤ (A.43) M 0 M0 ≈ , −1 1 rot with Mrot = [U (g1rot ), . . . , U (gm )].

which is rank-deficient: (M0⊤ M0 )−1 is ill-defined, Sc0 → 0, and ∥gConFIG ∥ → 0. Since HRGrad rotates both gradients into H before computing the pseudoinverse, the harmonized matrix Mrot = [U(g1rot ), U(g2rot )] has pairwise non-conflicting columns ((g1rot )⊤ g2rot ≥ 0), eliminating the anti-parallel degeneracy caused by the original gradients; under the column-rank condition in Lemma A.3, this yields Sc > 0. A.7.4. HRGrad compared with AlignGrad. AlignGrad [?] quantifies gradient alignment through the score m

(A.44)

A(v1 , . . . , vm ) = 2

2 1 X U(vi ) − 1 ∈ [−1, 1], m i=1

which equals the standard cosine similarity for m = 2 and equals 1 iff all unit gradients coincide. The direction maximizing the intra-step alignment score A(g1 , . . . , gm ) is

Z.Y. LIANG

the normalized sum of unit gradients, (A.45)

guA = U

m hX

i U(gi ) = U[M0 1m ],

i=1

and the corresponding update, scaled by the total projection length, is (A.46)

gAlignGrad :=

m X

gi⊤ guA guA .

i=1

Since guA is a unit vector, the Aggregate Product Identity gives ( i gi )⊤ gAlignGrad = ∥gAlignGrad ∥2 , so AlignGrad admits a convergence guarantee of the same structural P

form as Theorems A.4–A.6. The relation with ConFIG is immediate from comparing (A.40) with (A.46). Both methods aggregate the original gradients, but they use different directions: AlignGrad uses guA = U[M0 1m ], whereas ConFIG uses gu0 = U[M0 (M0⊤ M0 )−1 1m ]. When M0⊤ M0 = Im , these two choices coincide and guA = gu0 ; in general, ConFIG reweights correlated task gradients to enforce equal cosine similarity, whereas AlignGrad treats all tasks equally. This distinction is important because AlignGrad does not enforce uniform cosine similarity across tasks. By the Moore-Penrose identity, ConFIG satisfies U(gi )⊤ gu0 = Sc0 for all i, while HRGrad satisfies U(girot )⊤ gu = Sc for all i. For AlignGrad, however, the individual projections are Pm U(gi )⊤ j=1 U(gj ) ⊤ A Pm (A.47) U(gi ) gu = , j=1 U(gj ) which are not equal in general. A task whose unit gradient is more closely aligned with the mean direction receives a larger projection and therefore exerts disproportionate influence on the update. In multiscale kinetic problems, where macroscopic and microscopic gradients may have substantially different pairwise correlations, this mechanism systematically under-weights minority-direction tasks. The same averaging structure also limits the non-conflict guarantee beyond the two-task setting. For m = 2, one has gi⊤ guA ∝ 1 + cos θ12 > 0 whenever θ12 < π, so the update is non-conflicting. For m > 2, however, there is no guarantee that gi⊤ guA > 0 for every task: a minority gradient gk opposing the majority mean may satisfy gk⊤ guA < 0. HRGrad removes this difficulty by deflecting every girot into H before aggregation, which guarantees (girot )⊤ gjrot ≥ 0 for all i, j and all m (Lemma A.3(ii)). This P same loss of coherence appears in the severe-conflict regime. When g1 ≈ −g2 , A one has i U(gi ) ≈ 0, so gu becomes undefined and ∥gAlignGrad ∥ → 0. HRGrad avoids P this degeneracy by rotating all gradients into H before aggregation, ensuring ∥ i U(girot )∥ > 0 regardless of the original conflict angle. Finally, this degradation under conflict enters directly into the convergence-rate constant. The non-convex rate for AlignGrad has the same structural form as Theorem A.6, but with the HRGrad constant κ = ρmin /Sc replaced by κA = ρmin /αA , where Pm k i=1 U(gi ) (A.48) αA = min k m is the minimum per-step average alignment score. When conflicts are severe, αA → 0, the bound on mink ∥∇θ L(θk )∥2 deteriorates accordingly. By contrast, HRGrad

HRGRAD

guarantees Sck > 0 at every step (Remark A.7), so κ remains uniformly bounded and the convergence rate is preserved across all conflict regimes. A.7.5. Summary comparison. Table 1 summarizes the key theoretical properties of the four tabulated methods. Table 1 Comparison of gradient conflict resolution methods. Sc : equal cosine similarity constant; m: number of tasks; ✓: property is provably guaranteed; ×: not guaranteed.

Property Non-conflict update (2 tasks) Non-conflict update (m > 2 tasks) Pairwise non-conflict before aggregation Equal cosine with aggregation direction (uniform) Isometric Fidelity: ∥girot ∥ = ∥gi ∥ Per-task compatibility: gi⊤ girot ≥ 0 Aggregate Product Identity on aggregation inputs Convergence guarantee (convex) Convergence guarantee (non-convex) Uniform rate: α = mink Sck > 0 always Numerical stability under severe conflict Reduces to ConFIG when conflict-free †

PCGrad

AlignGrad

ConFIG

HRGrad

✓ × × ×

✓ ✓ × ✓

✓ ✓ ✓ ✓

× N/A ×

N/A N/A ✓

✓ ✓ ✓

× × × × N/A

✓ ✓ × × N/A†

✓ ✓ × × N/A

✓ ✓ ✓ ✓ ✓

AlignGrad = ConFIG when M0⊤ M0 = Im (orthonormal unit gradients).

Appendix B. Computation time. In this section, we present the computation times for solutions across various problems addressed in. Appendix C. Computational complexity and Knudsen numbers. In this section, slight deviations are observable at both t = 0 and t = 0.1 for different Knudsen numbers. These deviations can be attributed to the abrupt changes in the particle density function f across x = 0, where both the magnitude and dispersion of the particle density function undergo more pronounced shifts compared to those observed in Fig. Furthermore, for ε = 1.0, the deviations observed in the temperature may be linked to the distinct shape of the particle density function. Fig displays slices of the particle density function vx 7→ f at t = 0.1, x = 0.15, vy = vz = 0, for ε ∈ {1.0, 0.1, 0.01, 0.001}. For ε = 0.001, the density function exhibits a bell-shaped curve. In contrast, for ε ∈ {1.0, 0.01}, the function displays a bimodal distribution with two peaks, suggesting the need for a higher density of collocation points in the vx direction, thereby increasing the computational complexity. REFERENCES [1] R. Alonso, Boltzmann-type equations and their applications, Publicações Matemáticas do IMPA. [IMPA Mathematical Publications], Instituto Nacional de Matemática Pura e Aplicada (IMPA), Rio de Janeiro, 2015. 30o Colóquio Brasileiro de Matemática. [30th Brazilian Mathematics Colloquium]. [2] F. Bouchut, F. Golse, and M. Pulvirenti, Kinetic equations and asymptotic theory, Elsevier, 2000. [3] R. Caruana, Multitask learning, Machine learning, 28 (1997), pp. 41–75. [4] C. Cercignani and G. M. Kremer, The relativistic Boltzmann equation: theory and applications, vol. 22, Springer Science & Business Media, 2002.

Z.Y. LIANG

[5] J. Chen, Z. Ma, and K. Wu, A micro-macro decomposition-based asymptotic-preserving random feature method for multiscale radiative transfer equations, Journal of Computational Physics, (2025), p. 114103. [6] Z. Chen, J. Ngiam, Y. Huang, T. Luong, H. Kretzschmar, Y. Chai, and D. Anguelov, Just pick a sign: Optimizing deep multitask models with gradient sign dropout, Advances in Neural Information Processing Systems, 33 (2020), pp. 2039–2050. [7] F. de Avila Belbute-Peres, Y.-f. Chen, and F. Sha, Hyperpinn: Learning parameterized differential equations with physics-informed hypernetworks, The symbiosis of deep learning and differential equations, 690 (2021). [8] G. Dimarco and L. Pareschi, Numerical methods for kinetic equations, Acta Numerica, 23 (2014), pp. 369–520. [9] S. Duan, W. Wu, P. Hu, Z. Ren, D. Peng, and Y. Sun, Copinn: Cognitive physics-informed neural networks, in Forty-second International Conference on Machine Learning, 2025. [10] F. Filbet and S. Jin, A class of asymptotic-preserving schemes for kinetic equations and related problems with stiff sources, Journal of Computational Physics, 229 (2010), pp. 7625– 7648. [11] J. Han, J. Lu, and M. Zhou, Solving high-dimensional eigenvalue problems using deep neural networks: A diffusion monte carlo like approach, Journal of Computational Physics, 423 (2020), p. 109792. [12] X. Huang, Z. Ye, H. Liu, S. Ji, Z. Wang, K. Yang, Y. Li, M. Wang, H. Chu, F. Yu, et al., Meta-auto-decoder for solving parametric partial differential equations, Advances in Neural Information Processing Systems, 35 (2022), pp. 23426–23438. [13] L. Jiang, D. Meng, Q. Zhao, S. Shan, and A. Hauptmann, Self-paced curriculum learning, in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 29, 2015. [14] S. Jin, Asymptotic preserving (ap) schemes for multiscale kinetic and hyperbolic equations: a review, Lecture notes for summer school on methods and models of kinetic theory (M&MKT), Porto Ercole (Grosseto, Italy), (2010), pp. 177–216. [15] S. Jin, Z. Ma, and K. Wu, Asymptotic-preserving neural networks for multiscale timedependent linear transport equations, Journal of Scientific Computing, 94 (2023), p. 57. [16] S. Jin, K. Wu, et al., Asymptotic-preserving neural networks for multiscale kinetic equations, Communications in Computational Physics, 35 (2024), pp. 693–723. [17] A. Jungel, Transport Equations for Semiconductors, Springer-Verlag, Berlin Heidelberg, 2009. [18] A. Kendall, Y. Gal, and R. Cipolla, Multi-task learning using uncertainty to weigh losses for scene geometry and semantics, in Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 7482–7491. [19] D. P. Kingma and J. Ba, Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980, (2014). [20] Z. Li, Y. Wang, H. Liu, Z. Wang, and B. Dong, Solving the boltzmann equation with a neural sparse representation, SIAM Journal on Scientific Computing, 46 (2024), pp. C186–C215. [21] B. Lin, F. Ye, Y. Zhang, and I. W. Tsang, Reasonable effectiveness of random weighting: A litmus test for multi-task learning, arXiv preprint arXiv:2111.10603, (2021). [22] B. Liu, Y. Feng, P. Stone, and Q. Liu, Famo: Fast adaptive multitask optimization, Advances in Neural Information Processing Systems, 36 (2023), pp. 57226–57243. [23] B. Liu, X. Liu, X. Jin, P. Stone, and Q. Liu, Conflict-averse gradient descent for multi-task learning, Advances in Neural Information Processing Systems, 34 (2021), pp. 18878–18890. [24] L. Liu, Y. Li, Z. Kuang, J. Xue, Y. Chen, W. Yang, Q. Liao, and W. Zhang, Towards impartial multi-task learning, in International Conference on Learning Representations, iclr, 2021. [25] L. Liu, Y. Wang, X. Zhu, and Z. Zhu, Asymptotic-preserving neural networks for the semiconductor boltzmann equation and its application on inverse problems, Journal of computational physics, 523 (2025), p. 113669. [26] Q. Liu, M. Chu, and N. Thuerey, Config: Towards conflict-free training of physics informed neural networks, arXiv preprint arXiv:2408.11104, (2024). [27] S. Liu, E. Johns, and A. J. Davison, End-to-end multi-task learning with attention, in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2019, pp. 1871–1880. [28] L. Lu, X. Meng, Z. Mao, and G. E. Karniadakis, Deepxde: A deep learning library for solving differential equations, SIAM review, 63 (2021), pp. 208–228. [29] L. Lyu, K. Wu, R. Du, and J. Chen, Enforcing exact boundary and initial conditions in the deep mixed residual method, arXiv preprint arXiv:2008.01491, (2020). [30] A. Navon, A. Shamsian, I. Achituve, H. Maron, K. Kawaguchi, G. Chechik, and E. Fetaya, Multi-task learning as a bargaining game, arXiv preprint arXiv:2202.01017, (2022).

HRGRAD

[31] J. Oh, S. Y. Cho, S.-B. Yun, E. Park, and Y. Hong, Separable physics-informed neural networks for solving the bgk model of the boltzmann equation, SIAM Journal on Scientific Computing, 47 (2025), pp. C451–C474. [32] L. Pareschi and R. E. Caflisch, An implicit monte carlo method for rarefied gas dynamics: I. the space homogeneous case, Journal of Computational Physics, 154 (1999), pp. 90–116. [33] L. Pareschi and G. Russo, Time relaxed monte carlo methods for the boltzmann equation, SIAM Journal on Scientific Computing, 23 (2001), pp. 1253–1273. [34] N. Park, W. Cho, M. Jo, H. Lim, K. Lee, D. Lee, S. Hong, et al., Parameterized physicsinformed neural networks for parameterized pdes, in 41st International Conference on Machine Learning, ICML 2024, ICML, 2024. [35] M. Raissi, P. Perdikaris, and G. E. Karniadakis, Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations, Journal of Computational physics, 378 (2019), pp. 686–707. [36] O. Sener and V. Koltun, Multi-task learning as multi-objective optimization, Advances in neural information processing systems, 31 (2018). [37] Y. Song, H. Wang, H. Yang, M. L. Taccari, and X. Chen, Loss-attentional physics-informed neural networks, Journal of Computational Physics, 501 (2024), p. 112781. [38] T. Standley, A. Zamir, D. Chen, L. Guibas, J. Malik, and S. Savarese, Which tasks should be learned together in multi-task learning?, in International conference on machine learning, PMLR, 2020, pp. 9120–9132. [39] M. Toloubidokhti, Y. Ye, R. Missel, X. Jiang, N. Kumar, R. Shrestha, and L. Wang, Dats: Difficulty-aware task sampler for meta-learning physics-informed neural networks, in The Twelfth International Conference on Learning Representations, 2023. [40] R. Van Der Meer, C. W. Oosterlee, and A. Borovykh, Optimally weighted loss functions for solving pdes with neural networks, Journal of Computational and Applied Mathematics, 405 (2022), p. 113887. [41] S. Vandenhende, S. Georgoulis, W. Van Gansbeke, M. Proesmans, D. Dai, and L. Van Gool, Multi-task learning for dense prediction tasks: A survey, IEEE transactions on pattern analysis and machine intelligence, 44 (2021), pp. 3614–3633. [42] C. Villani, A review of mathematical topics in collisional kinetic theory, in Handbook of Mathematical Fluid Mechanics, S. Friedlander and D. Serre, eds., vol. I, North-Holland, 2002, pp. 71–305. [43] S. Wang, A. K. Bhartari, B. Li, and P. Perdikaris, Gradient alignment in physicsinformed neural networks: A second-order optimization perspective, arXiv preprint arXiv:2502.00604, (2025). [44] S. Wang, X. Yu, and P. Perdikaris, When and why pinns fail to train: A neural tangent kernel perspective, Journal of Computational Physics, 449 (2022), p. 110768. [45] E. Weinan, Principles of multiscale modeling, Cambridge University Press, 2011. [46] Z. Xiang, W. Peng, X. Liu, and W. Yao, Self-adaptive loss balanced physics-informed neural networks, Neurocomputing, 496 (2022), pp. 11–34. [47] Z.-Q. J. Xu, Y. Zhang, T. Luo, Y. Xiao, and Z. Ma, Frequency principle: Fourier analysis sheds light on deep neural networks, arXiv preprint arXiv:1901.06523, (2019). [48] J. Yao, C. Su, Z. Hao, S. Liu, H. Su, and J. Zhu, Multiadam: Parameter-wise scale-invariant optimizer for multiscale training of physics-informed neural networks, in International Conference on Machine Learning, PMLR, 2023, pp. 39702–39721. [49] T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn, Gradient surgery for multi-task learning, Advances in neural information processing systems, 33 (2020), pp. 5824–5836. [50] Z. Yuan, H. Liu, H. Zhou, D. Zhang, X. Zhang, H. Wang, and H. Xiong, Self-paced unified representation learning for hierarchical multi-label classification, in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 38, 2024, pp. 16623–16632. [51] Y. Zhang and Q. Yang, An overview of multi-task learning, National Science Review, 5 (2018), pp. 30–43. [52] Y. Zhang and Q. Yang, A survey on multi-task learning, IEEE transactions on knowledge and data engineering, 34 (2021), pp. 5586–5609. [53] D. Zhao, Y. S. Koh, G. Dobbie, H. Hu, and P. Fournier-Viger, Symmetric self-paced learning for domain generalization, in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 38, 2024, pp. 16961–16969.

Record · ID 138908 · SHA-256 159b77e86fc85377
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.