MODELING CONVERGENCE RATE ESTIMATES OF THE HEAVY BALL ALGORITHM VIA A RANDOM FOREST WITH SELF-ATTENTION

Abstract

This paper addresses the problem of estimating convergence rates of iterative optimization algorithms, which is essential for analyzing and accelerating the training of modern machine learning models. Classical approaches to convergence rate estimation for momentum-based methods typically rely on analytical bounds derived under restrictive assumptions or on extensive empirical experimentation, both of which are computationally expensive and poorly scalable. As an alternative, we propose a data-driven framework for modeling convergence rate estimates using supervised machine learning. The proposed approach is demonstrated for the Heavy Ball optimization method applied to quadratic objective functions with heterogeneous curvature. A synthetic dataset is constructed by simulating Heavy Ball trajectories under randomly sampled objective parameters, optimization hyperparameters, and normalized initial conditions. The input features capture key characteristics of the optimization process, while the target variable represents a logarithmic estimate of the objective function decay along the trajectory, serving as a proxy for the convergence rate. To model nonlinear dependencies between optimization parameters and convergence behavior, we employ a Random Forest model augmented with a self-attention mechanism. This architecture enables adaptive weighting of feature interactions and improves predictive accuracy without requiring explicit analytical modeling of the underlying dynamics. Numerical experiments show that the proposed model achieves a coefficient of determination exceeding 0.8 when predicting convergence-related quantities, indicating strong agreement with empirically observed behavior. The results suggest that machine learning-based surrogate models can substantially simplify convergence rate estimation and that the proposed methodology is applicable to other optimization algorithms and objective classes.