Proximal Gradient Method

The proximal gradient method's core application is solving composite optimization problems of the form: minimize f(x) + g(x), where f is differentiable (e.g., a data fidelity loss) and g is convex but potentially non-differentiable. Common non-differentiable regularizers g include:

• L1 Norm (Lasso): Promotes sparsity in model weights.
• Nuclear Norm: Encourages low-rank solutions, crucial for matrix completion and low-rank factorization.
• Group Lasso: Induces group-wise sparsity. The algorithm alternates between a gradient step on f and a proximal step (e.g., soft-thresholding for L1) on g, efficiently handling the non-smooth component.

This method is fundamental for recovering missing entries in a matrix under a low-rank assumption. The objective typically uses the Frobenius norm for the differentiable data-fitting term f and the nuclear norm (sum of singular values) as the non-differentiable regularizer g to promote low rank. The proximal operator for the nuclear norm is singular value thresholding, which shrinks singular values toward zero. This same framework powers Robust PCA, which separates a matrix into a low-rank component (modeled with a nuclear norm) and a sparse outlier component (modeled with an L1 norm).

Proximal gradient methods enable aggressive model compression by enforcing structured sparsity patterns during training, a technique known as learning-compression. Key applications include:

• Weight Pruning: Using L1 or group Lasso regularizers to drive unimportant weights to exactly zero, creating a sparse model.
• Low-Rank Factorization of Layers: Applying nuclear norm regularization on weight matrices to implicitly encourage low-rank factorizations without pre-defining the rank.
• Channel Pruning in CNNs: Using group Lasso on filters to remove entire channels from convolutional layers. The proximal step cleanly handles the non-differentiable sparsity-inducing penalty, resulting in models that are natively compressed and hardware-friendly.

FISTA is a momentum-accelerated variant of the proximal gradient method that dramatically improves convergence rates. It achieves an O(1/k²) convergence rate for the objective function value, compared to the standard method's O(1/k) rate. The algorithm introduces an auxiliary sequence of points and uses a specific momentum term, performing:

1. A gradient descent step on the smooth part f.
2. A proximal operator application to the non-smooth part g.
3. An update with Nesterov-type momentum. FISTA is the workhorse algorithm for large-scale L1-regularized problems (like LASSO) in high-dimensional statistics and compressed sensing.

The proximal gradient method is closely related to ADMM, another powerful optimization framework. Both can solve composite optimization problems. ADMM often decomposes the problem differently, using variable splitting and the method of multipliers. In many cases, a single iteration of ADMM involves solving subproblems that are themselves proximal operations. For problems with simple proximal operators, the proximal gradient method can be more straightforward and efficient. For problems with complex constraints or multiple non-smooth terms, ADMM's splitting strategy may be preferred. Understanding both provides a toolkit for designing custom training loops for compressed models.

While often implemented custom for research, core concepts are embedded in popular ML libraries:

• PyTorch: Custom optimizers can be built using torch.optim.Optimizer with a step() method that computes gradients for f and then applies a proximal operator (e.g., torch.nn.functional.soft_threshold for L1).
• TensorFlow: The tf.train.ProximalGradientDescentOptimizer is available for L1 and L2 regularization, implementing the proximal method for these specific cases.
• JAX: Its functional nature and jax.grad make it ideal for implementing custom proximal gradient loops with JIT compilation. The efficiency hinges on the proximal operator being computationally cheap, which is true for norms like L1, L2, and nuclear norm.

The proximal operator (prox-operator) is the core component of the proximal gradient method. For a function $g(x)$ and a scalar $\eta > 0$, it is defined as: $$\text{prox}_{\eta g}(v) = \arg\min_x \left( g(x) + \frac{1}{2\eta} |x - v|_2^2 \right)$$

• It acts as a generalized gradient step for potentially non-differentiable functions.
• For common regularizers like the L1 norm ($|x|_1$), the proximal operator is the soft-thresholding function, which sets small values to zero, inducing sparsity.
• For the nuclear norm (sum of singular values), the proximal operator applies soft-thresholding to the singular values of a matrix, promoting low-rank structure.

The nuclear norm (or trace norm) of a matrix $X$, denoted $|X|_*$, is the sum of its singular values. It is the convex envelope of the matrix rank function on the unit ball.

• Key Property: Minimizing the nuclear norm encourages matrices to be low-rank, analogous to how minimizing the L1 norm encourages vectors to be sparse.
• Use in Compression: In low-rank matrix completion and robust PCA, the nuclear norm is used as a convex regularizer. The proximal gradient method efficiently solves these problems because the proximal operator for the nuclear norm is computationally tractable (via SVD and soft-thresholding).

Iterative Hard Thresholding (IHT) is a related greedy algorithm for sparse approximation and low-rank recovery. It solves problems with explicit cardinality or rank constraints.

• Mechanism: Each iteration performs a gradient step on the smooth loss, followed by hard thresholding. For low-rank matrices, this means computing the SVD and keeping only the top-$k$ singular values/vectors, setting the rest to zero.
• Contrast with Proximal Gradient: IHT uses a hard, non-convex constraint, while proximal gradient (with nuclear norm) uses a soft, convex penalty. IHT can be faster but may converge to local minima, whereas proximal gradient offers convergence guarantees to a global minimum for convex problems.

The Alternating Direction Method of Multipliers (ADMM) is another powerful optimization algorithm for problems with separable objectives and constraints, commonly written as $\min f(x) + g(z)$ subject to $Ax + Bz = c$.

• Relation: Like proximal gradient, ADMM leverages proximal operators for the subproblems involving $f$ and $g$. It is particularly effective when the proximal operators for $f$ and $g$ are easy to compute separately, even if the combined problem is complex.
• Application: ADMM is widely used for distributed optimization and problems like Robust PCA, where the objective naturally splits into a low-rank term and a sparse term.

Majorization-Minimization (MM) is a general optimization framework that underlies many algorithms, including proximal gradient.

• Principle: At each iteration, it constructs a surrogate function (a majorizer) that is easier to minimize than the original objective and lies above it, touching at the current point.
• Connection: The proximal gradient method can be derived as an MM algorithm. The gradient step plus proximal step effectively minimizes a quadratic majorizer of the smooth function $f(x)$, augmented with the non-smooth $g(x)$. This perspective provides a clear path to proving its convergence guarantees.

Forward-Backward Splitting is the canonical name for the proximal gradient method, explicitly describing its two-step structure.

• Forward Step: A standard gradient descent step on the differentiable function $f(x)$: $x_k - \eta \nabla f(x_k)$.
• Backward Step: The application of the proximal operator for the non-differentiable function $g(x)$ to the result of the forward step.
• This splitting is optimal for problems where the gradient $\nabla f$ is Lipschitz continuous and the proximal operator of $g$ is inexpensive to compute, which is typical in compressed sensing and low-rank factorization tasks.