---
title: Generalized Primal Averaging (GPA)
url: https://www.emergentmind.com/topics/generalized-primal-averaging-gpa
type: topic
---

# Generalized Primal Averaging (GPA)

Generalized Primal Averaging (GPA) is a family of optimization techniques that implement weighted or exponential averaging of primal iterates within first-order or primal-dual algorithms. GPA arises in both distributed and non-distributed settings, seeking to combine strong theoretical guarantees with simplicity, acceleration, and communication-efficiency. GPA interfaces tightly with dual averaging, Nesterov-type momentum, and modern stochastic optimization, and underpins the convergence theory of several widely-used optimizers in machine learning.

## 1. Mathematical Formulation and Core Principles

GPA algorithms maintain one or more sequences of primal variables, updating them via recursive weighted averaging of model iterates or search directions. A basic GPA update for iterates $x^{(t)}$ is
\[
x^{(t+1)} = \mu_x \, x^{(t)} + (1-\mu_x) \, z^{(t+1)}
\]
where $\mu_x \in [0,1)$ denotes the exponential moving average (EMA) coefficient and $z^{(t+1)}$ is the update from the base optimizer. Typical implementations also feature a "gradient computation" point,
\[
y^{(t)} = \mu_y \, x^{(t)} + (1-\mu_y) \, z^{(t)},
\]
decoupling the score at which gradients are computed ($y^{(t)}$) from where the model evaluation occurs. This decoupling generalizes classical averaging schemes such as Polyak–Ruppert, Nesterov acceleration, and momentum [2512.17131][2010.00406].

In distributed convex optimization, GPA manifests in frameworks such as CoCoA⁺ [1502.03508], where local subproblems yield primal and dual updates that are aggregated globally with tunable weight. In constrained convex programs, primal recovery is effected by weighted Fenchel-type averaging over primal approximations [1511.03305].

## 2. Distributed Primal-Dual Algorithms and Communication Efficiency

In the distributed setting, GPA is exemplified by the CoCoA⁺ framework for empirical risk minimization [1502.03508]. The primal problem is
\[
\min_{w\in\mathbb R^d} \quad \frac{1}{n} \sum_{i=1}^n \ell_i(w^T x_i) + \frac{\lambda}{2}\|w\|^2,
\]
with dual
\[
\max_{\alpha\in\mathbb R^n} \quad -\frac{1}{n}\sum_{i=1}^n \ell_i^*(-\alpha_i) - \frac{\lambda}{2} \left\| \frac{1}{\lambda n} A \alpha \right\|^2.
\]
GPA-type global update rules interpolate between conservative averaging,
\[
\alpha^{(t+1)} = \alpha^{(t)} + \frac{1}{K} \sum_{k=1}^K \Delta\alpha_{[k]}^{(t)}
\]
and aggressive (additive) aggregation,
\[
\alpha^{(t+1)} = \alpha^{(t)} + \sum_{k=1}^K \Delta\alpha_{[k]}^{(t)},
\]
where $K$ is the number of workers. Setting the aggregation parameter $\nu=1$ yields the GPA/CoCoA⁺ variant, which provably achieves iteration complexity independent of $K$ under suitable quadratic-separability conditions.

Empirically, CoCoA⁺ converges in roughly half the communication rounds of the original CoCoA. Strong scaling is observed in large clusters; increasing $K$ can even decrease time-to-accuracy, whereas naive averaging degrades linearly with $K$. Safe settings ($\nu=1$, $\sigma'=K$) require no tuning and are robust across data regimes [1502.03508].

## 3. Weighted Primal Averaging and Accelerated Primal-Dual Algorithms

GPA is foundational in alternating minimization algorithms (AMA) for linearly constrained convex programs [1511.03305]. Consider
\[
\min_{u\in U, v\in V} \quad g(u) + h(v) \quad \text{s.t.} \quad Au + Bv = c,
\]
with no strong convexity required. The primal recovery step computes weighted averages over primal candidates,
\[
\bar{u}^k = \frac{1}{S_k}\sum_{i=0}^k w_i \tilde{u}^i, \quad \bar{v}^k = \frac{1}{S_k}\sum_{i=0}^k w_i \tilde{v}^i,
\]
where $w_i$ are positive weights (such as dual step-sizes or accelerated weights).

Non-accelerated primal-dual AMA with GPA yields an $\varepsilon$-solution in $O(L D_U / \varepsilon^2)$ iterations, while accelerated variants achieve $O(\sqrt{L D_U} / \varepsilon)$, both without requiring strong convexity. These rates are optimal for black-box first-order models (Nemirovskii–Yudin lower bounds) [1511.03305].

## 4. Momentum, Non-Convex Optimization, and Lyapunov Analysis

GPA directly specializes to stochastic primal averaging (SPA) forms, which are equivalent to SGD with momentum ("heavy-ball") under a specific reparameterization [2010.00406]. The core update involves two sequences,
\[
\begin{aligned}
z_{k+1} &= z_k - \eta_k \nabla f(x_k, \xi_k), \\
x_{k+1} &= (1 - c_{k+1}) x_k + c_{k+1} z_{k+1},
\end{aligned}
\]
with the equivalence given by $\alpha_k = \eta_k c_{k+1}$, $\beta_k = \frac{\eta_{k-1}(1-c_k)}{\eta_k}$.

Lyapunov-based analysis yields tight $O(1/\sqrt{T})$ bounds in the non-convex stochastic setting, matching that of SGD but with reduced variance and improved constants. The negative term proportional to $\|x_k-x_{k-1}\|^2$ cancels a substantial portion of the noise, explaining the empirical advantage of momentum especially in early epochs.

Practical guidance includes simultaneous scheduling of learning rates and averaging weights: for example, upon decaying $\eta$ by a factor $\phi$, increase $c$ by $\phi$ (i.e., decrease momentum coefficient). The marginal benefit of momentum vanishes after a few training stages.

## 5. GPA in Large-Scale Machine Learning and LLM Training

Recent work extends GPA to wrap modern optimizers (AdamW, etc.) via EMA, delivering strong empirical performance in non-distributed single-worker settings [2512.17131]. A canonical algorithm maintains three points: $y^{(t)}$ (gradient computation), $z^{(t)}$ (update), and $x^{(t)}$ (model average),
```python
y(t) = μ_y * x(t) + (1–μ_y) * z(t)
g(t) = ∂f(y(t); ξ(t))
d(t) = BaseOpt(g(t))
z(t+1) = (1–γ(t)λ) * z(t) + γ(t) * d(t)
x(t+1) = μ_x * x(t) + (1–μ_x) * z(t+1)
```
where $\mu_x$ and $\mu_y$ are independent EMA coefficients. By unifying the averaging schemes of Schedule-Free (SF) and single-worker DiLoCo, GPA removes the need for two-loop structures and additional hyperparameters.

On Llama-160M and Llama-1B models, GPA-AdamW yields up to 24% steps-to-baseline-loss speedup compared to AdamW, and outperforms single-worker DiLoCo at optimal inner-steps. On ImageNet ViT workloads, similar speedups and top-1 accuracy gains are reported across batch sizes.

Theoretical guarantees show that GPA inherits the online regret bound $O(\sqrt{T})$ from its base optimizer and can further accelerate convergence depending on the values of $\mu_x$ and $\mu_y$ due to non-negative Bregman-divergence terms.

## 6. Hyperparameter Selection, Implementation, and Practical Guidelines

Implementing GPA requires storing at most one additional buffer beyond the base optimizer state, which can often be reduced further by implicit reconstruction ($x = y/\mu_y + (1-1/\mu_y)z$). Standard hyperparameter schedules (AdamW betas, learning rates) can be retained. Tuning $\mu_x$ and $\mu_y$ translates directly from DiLoCo—given DiLoCo parameters $(H, \mu)$, set $\mu_x \approx \mu^{1/H}$, $\mu_y \approx \mu$.

In distributed settings, choose aggregation $\nu=1$, separability parameter $\sigma'=K$, and local-solver steps $H$ to balance compute and communication. For unconstrained convex problems, weighted averaging with step-weight $w_k$ yields optimal convergence rates when choosing $\gamma = \varepsilon/D_U$.

The online duality gap is recommended as a universal stopping criterion. No strong convexity is required for rate guarantees in constrained programs or distributed dual methods.

## 7. Limitations, Extensions, and Open Problems

GPA’s convergence guarantees in the non-distributed setting are presently limited to the averaged-iterate regime under convexity assumptions; extension to nonconvex last-iterate guarantees remains open [2512.17131]. Full compatibility with advanced base optimizers and cross-region distributed regimes demands further theoretical analysis. While GPA provides optimal rates and flexibility, the introduction of learning rate schedules may re-introduce complexity absent in pure uniform-averaging schemes. 

In summary, GPA unifies and extends weighted averaging, momentum, and primal-dual update schemes, underpinning scalable, communication-efficient, and empirically superior optimization algorithms in modern machine learning, convex and nonconvex alike.

Source: https://www.emergentmind.com/topics/generalized-primal-averaging-gpa