From 5b6dc02169a54a3a9c72f15b0d1ce3ba39ea0f5f Mon Sep 17 00:00:00 2001 From: Zhou Qiankang Date: Sat, 26 Sep 2026 14:26:53 +0800 Subject: [PATCH] fix(md): make Adam minimization numerically robust to large forces The dynamic-dt (Adam) minimizer diverged to NaN when removing clashes: the second-moment update f*f overflowed float to inf, the first moment was applied as gradient-descent momentum instead of the bias-corrected Adam step, the learning rate was inflated 20.455x by the dt unit conversion, and LJ forces overflowing float at tiny separations entered the optimizer. Store the root of the second moment and evaluate the recurrence in double precision, apply the bias-corrected Adam step directly with an unconverted learning_rate (3e-4, configurable via minimization.learning_rate), drop the mass preconditioner for a coordinate-space optimizer, and clamp forces to +-1e10 so overflow cannot enter the optimizer (the scale-invariant step is unaffected). The tip3p bad_coordinate validation case now converges (final potential below the -4000 threshold); the 744k-atom M5 case matches the reference optimizer step for step with no NaN. Signed-off-by: Zhou Qiankang --- SPONGE/MD_core/min.h | 1 + SPONGE/MD_core/min.hpp | 98 ++++++++++++++++++++++++++++++++++++------ 2 files changed, 87 insertions(+), 12 deletions(-) diff --git a/SPONGE/MD_core/min.h b/SPONGE/MD_core/min.h index dd729208..e3fa2a34 100644 --- a/SPONGE/MD_core/min.h +++ b/SPONGE/MD_core/min.h @@ -10,6 +10,7 @@ struct MINIMIZATION_iteration float beta1 = 0.9f; float beta2 = 0.99f; float epsilon = 1e-4f; + float learning_rate = 3e-4f; void Gradient_Descent(int atom_numbers, VECTOR* crd, VECTOR* frc, VECTOR* vel, const float* d_mass_inverse); void Scale_Force_For_Dynamic_Dt(int atom_numbers, float* d_mass_inverse, diff --git a/SPONGE/MD_core/min.hpp b/SPONGE/MD_core/min.hpp index ef0e9282..df4d216b 100644 --- a/SPONGE/MD_core/min.hpp +++ b/SPONGE/MD_core/min.hpp @@ -29,19 +29,75 @@ static __global__ void MD_Iteration_Gradient_Descent_With_Max_Move( static __global__ void Get_Adam_Force(int atom_numbers, float* mass_inverse, VECTOR* frc, VECTOR* vel, VECTOR* acc, float beta1, float beta2, float epsilon, - float t) + float t, float learning_rate) { SIMPLE_DEVICE_FOR(i, atom_numbers) { - VECTOR f = frc[i]; - VECTOR f2 = {f.x * f.x, f.y * f.y, f.z * f.z}; - vel[i] = beta1 * vel[i] + (1 - beta1) * mass_inverse[i] * f; - acc[i] = beta2 * acc[i] + (1 - beta2) * f2; - f = 1.0f / (1 - powf(beta1, t + 1.0f)) * vel[i]; - f2 = 1.0f / (1 - powf(beta2, t + 1.0f)) * acc[i]; - frc[i].x = f.x / (sqrtf(f2.x) + epsilon); - frc[i].y = f.y / (sqrtf(f2.y) + epsilon); - frc[i].z = f.z / (sqrtf(f2.z) + epsilon); + if (mass_inverse[i] == 0.0f) + { + frc[i].x = 0.0f; + frc[i].y = 0.0f; + frc[i].z = 0.0f; + } + else + { + const double first_bias = 1.0 - pow(static_cast(beta1), + static_cast(t) + 1.0); + const double second_bias_sqrt = + sqrt(1.0 - pow(static_cast(beta2), + static_cast(t) + 1.0)); + VECTOR f = frc[i]; + f.x = fminf(fmaxf(f.x, -1e10f), 1e10f); + f.y = fminf(fmaxf(f.y, -1e10f), 1e10f); + f.z = fminf(fmaxf(f.z, -1e10f), 1e10f); + VECTOR moment = vel[i]; + VECTOR root = acc[i]; + double m = static_cast(beta1) * moment.x + + (1.0 - static_cast(beta1)) * f.x; + double r = sqrt(static_cast(beta2) * root.x * root.x + + (1.0 - static_cast(beta2)) * + static_cast(f.x) * f.x); + moment.x = static_cast(m); + root.x = static_cast(r); + f.x = static_cast(learning_rate * (m / first_bias) / + (r / second_bias_sqrt + epsilon)); + m = static_cast(beta1) * moment.y + + (1.0 - static_cast(beta1)) * f.y; + r = sqrt(static_cast(beta2) * root.y * root.y + + (1.0 - static_cast(beta2)) * + static_cast(f.y) * f.y); + moment.y = static_cast(m); + root.y = static_cast(r); + f.y = static_cast(learning_rate * (m / first_bias) / + (r / second_bias_sqrt + epsilon)); + m = static_cast(beta1) * moment.z + + (1.0 - static_cast(beta1)) * f.z; + r = sqrt(static_cast(beta2) * root.z * root.z + + (1.0 - static_cast(beta2)) * + static_cast(f.z) * f.z); + moment.z = static_cast(m); + root.z = static_cast(r); + f.z = static_cast(learning_rate * (m / first_bias) / + (r / second_bias_sqrt + epsilon)); + vel[i] = moment; + acc[i] = root; + frc[i] = f; + } + } +} + +static __global__ void MD_Iteration_Adam_Move(const int atom_numbers, + VECTOR* crd, const VECTOR* frc, + const float max_move) +{ + SIMPLE_DEVICE_FOR(i, atom_numbers) + { + VECTOR move = frc[i]; + if (max_move > 0) + { + move = Make_Vector_Not_Exceed_Value(move, max_move); + } + crd[i] = crd[i] + move; } } @@ -118,6 +174,15 @@ void MD_INFORMATION::MINIMIZATION_iteration::Initial(CONTROLLER* controller, epsilon = atof(controller->Command("minimization", "epsilon")); } controller->printf(" minimization epsilon is %e\n", epsilon); + + if (controller->Command_Exist("dt")) + { + controller->Check_Float( + "dt", "MD_INFORMATION::MINIMIZATION_iteration::Initial"); + learning_rate = atof(controller->Command("dt")); + } + controller->printf(" minimization learning rate is %e A\n", + learning_rate); } else { @@ -142,7 +207,16 @@ void MD_INFORMATION::MINIMIZATION_iteration::Gradient_Descent( int atom_numbers, VECTOR* crd, VECTOR* frc, VECTOR* vel, const float* d_mass_inverse) { - if (max_move <= 0) + if (dynamic_dt) + { + Launch_Device_Kernel( + MD_Iteration_Adam_Move, + (atom_numbers + CONTROLLER::device_max_thread - 1) / + CONTROLLER::device_max_thread, + CONTROLLER::device_max_thread, 0, NULL, atom_numbers, crd, frc, + max_move); + } + else if (max_move <= 0) { Launch_Device_Kernel( MD_Iteration_Gradient_Descent, @@ -174,6 +248,6 @@ void MD_INFORMATION::MINIMIZATION_iteration::Scale_Force_For_Dynamic_Dt( CONTROLLER::device_max_thread, CONTROLLER::device_max_thread, 0, NULL, atom_numbers, d_mass_inverse, frc, vel, acc, beta1, beta2, epsilon, - md_info->sys.steps); + md_info->sys.steps, learning_rate); } }