GradPower

Implements GradPower, a sign-power gradient transformation applied before a base optimizer's update.

GradPower is a lightweight, plug-in modification that reshapes the raw gradient through an elementwise signed power before it enters the usual momentum and adaptive-scaling machinery. With exponent $p>1$ it sharpens large coordinates and suppresses small ones, while $p<1$ does the reverse; the sign is preserved so the descent direction per coordinate is unchanged. The example below instantiates it inside AdamW (AdamPower), where the transformation costs a single extra line over vanilla Adam.

$$\begin{aligned} g_t &\leftarrow |g_t|^{p}\,\mathrm{sign}(g_t) \\ m_t &= \beta_1 m_{t-1} + (1-\beta_1)\,g_t, \qquad \hat{m}_t = \frac{m_t}{1-\beta_1^{t}} \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)\,g_t^{2}, \qquad \hat{v}_t = \frac{v_t}{1-\beta_2^{t}} \\ \theta_t &= \theta_{t-1} - \eta_t\!\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} + \lambda\,\theta_{t-1}\right) \end{aligned}$$

where $\theta$ are the parameters, $\eta_t$ the learning rate, $g_t$ the gradient, $p>0$ the power exponent (e.g. $p=1.2$), $m_t,v_t$ the first and second moments with decay rates $\beta_1,\beta_2$, $\lambda$ the decoupled weight decay, and $\epsilon$ a stability constant; all powers and the sign act elementwise.

Reference: Mingze Wang, Jinbo Wang, Jiaqi Zhang, Wei Wang, Peng Pei, Xunliang Cai, Weinan E, Lei Wu, "GradPower: Powering Gradients for Faster Language Model Pre-Training", arXiv 2025. https://arxiv.org/abs/2505.24275

---

Back to the Canon