On $O(n)$ Algorithms for Projection onto the Top-$k$-sum Sublevel Set

1. Problem Definition

The paper addresses the Euclidean projection onto the top-$k$-sum sublevel set. Given a scalar budget $r$, an index $k \in \{1, \dots, n\}$, and an input vector $x_0 \in \mathbb{R}^n$, the goal is to solve the following strongly convex quadratic program:

$$\begin{aligned} \min_{x \in \mathbb{R}^n} \quad & \frac{1}{2} \|x - x_0\|_2^2 \\ \text{subject to} \quad & T^{(k)}(x) := \sum_{i=1}^k \vec{x}_i \leq r \end{aligned}$$

where $\vec{x}$ denotes the vector $x$ sorted in non-increasing order ($\vec{x}_1 \geq \vec{x}_2 \geq \dots \geq \vec{x}_n$), and $T^{(k)}(x)$ is the sum of the $k$ largest components.

Context & Motivation:

• Superquantile Optimization: The top-$k$-sum operator is directly related to the superquantile (also known as Conditional Value-at-Risk or CVaR). This projection is a critical subroutine in solving composite optimization problems involving risk-averse constraints (e.g., robust design, fair machine learning, and safety-critical systems).
• Existing Limitations:
  • General Solvers (e.g., Gurobi): While applicable, they are too slow for large-scale problems ($n \sim 10^7$) and often yield inexact solutions, complicating the computation of generalized Jacobians required for second-order methods.
  • Grid-Search Methods: Existing specialized methods (e.g., from [38]) have a complexity of $O(k(n-k))$. When $k$ is a linear fraction of $n$ (common in practice), this degrades to $O(n^2)$, making it prohibitive for large $n$.
  • Semismooth Newton Methods: While finite-termination is guaranteed, the floating-point complexity is unspecified, and they often require full sorting of the input vector.

2. Methodology

The authors propose two finite-termination algorithms that achieve $O(n)$ complexity for sorted input vectors, with constants independent of $k$. Both algorithms exploit the structure of the Karush-Kuhn-Tucker (KKT) conditions when the input is sorted.

A. Parametric Linear Complementarity Problem (PLCP) Approach

• Concept: This method treats the projection problem as a parametric LCP involving a Z-matrix (a matrix with non-positive off-diagonal elements).
• Mechanism:
  1. The summation constraint is penalized in the objective function with a parameter $\lambda \geq 0$.
  2. The problem is reformulated as finding a $\lambda$ such that the budget constraint is satisfied while maintaining the isotonic (ordering) constraints.
  3. The algorithm uses a pivoting method (specialized for Z-matrices) to trace the solution path as $\lambda$ increases from 0 to $+\infty$.
  4. Due to the tridiagonal structure of the associated matrix, each pivot step takes $O(1)$ time.
  5. The algorithm terminates once the budget constraint is met.
• Complexity: $O(n)$ for sorted inputs.

B. Early-Stopping Grid-Search (ESGS) Approach

• Concept: This is an optimized variant of the grid-search method introduced in [38]. Instead of exhaustively searching all pairs of indices $(k_0, k_1)$ that define the "breakpoints" of the solution, it exploits monotonicity properties of the KKT residuals.
• Mechanism:
  1. The algorithm starts at a specific index pair $(k-1, k)$.
  2. It traverses the search space by either decrementing $k_0$ or incrementing $k_1$ based on which KKT conditions are violated.
  3. Key Insight: The KKT conditions exhibit global monotonicity. If a condition fails at a certain index, the algorithm can prove that all indices "behind" it in the search direction are also invalid. This allows the algorithm to skip large portions of the grid search.
  4. The path moves strictly "up" or "right" in the $(k_0, k_1)$ space, guaranteeing termination in at most $n$ steps.
• Complexity: $O(n)$ for sorted inputs.

C. Handling Unsorted Inputs (Partial Sorting)

• For unsorted inputs, a full sort ($O(n \log n)$) is traditionally required. However, the authors observe that the solution only perturbs the top $k$ elements (specifically up to index $\bar{k}_1$).
• They propose an online sorting strategy using a max-heap. Elements are extracted only as needed during the algorithm's execution.
• Complexity: $O(\bar{k}_1 \log n)$, where $\bar{k}_1 \leq n$ is the index of the last modified element. This is significantly faster than full sorting when $\bar{k}_1 \ll n$.
• They also provide a rigorous check to verify if a candidate permutation (e.g., from a previous iteration in a sequence of problems) is sufficient, avoiding re-sorting entirely in many iterative scenarios.

3. Key Contributions

1. Two New $O(n)$ Algorithms: Introduction of PLCP and ESGS, which are the first methods to guarantee $O(n)$ complexity for the sorted version of this problem with constants independent of $k$.
2. Theoretical Rigor: Proofs of finite termination and complexity bounds, leveraging Z-matrix properties and KKT monotonicity.
3. Partial Sorting Optimization: A novel procedure to handle unsorted inputs with cost proportional to the number of affected elements rather than the full dimension, crucial for iterative applications.
4. Extension to Vector-$k$-Norm: The methods are adaptable to the projection onto the Ky-Fan norm ball (vector-$k$-norm), a related but slightly more complex problem.

4. Numerical Results

The authors implemented the algorithms in Julia and compared them against:

• GRID: The standard grid-search method ($O(k(n-k))$).
• SSN: Semismooth Newton method.
• Gurobi (GRBS/GRBU): Commercial QP solver (sorted and unsorted formulations).

Key Findings:

• Speed: The proposed methods (ESGS and PLCP) are orders of magnitude faster than existing methods.
  • For $n=10^7$ and $k=10^4$, ESGS/PLCP solve the problem in ~0.05 seconds.
  • The SSN method takes ~1 second.
  • Gurobi and the Grid-search method take minutes to hours (or time out).
• Scalability: The proposed methods exhibit linear scaling ($O(n)$), whereas the grid-search method exhibits quadratic scaling ($O(n^2)$).
• Robustness: ESGS consistently outperforms PLCP by a factor of $\approx 2$ in difficult instances (requiring many pivots) and remains competitive in easier instances.
• Sorting Overhead: For large $n$, the cost of full sorting exceeds the time taken by the proposed algorithms to solve the projection on a sorted vector, highlighting the value of the partial sorting optimization.

5. Significance

This work provides a practical, high-performance oracle for superquantile-constrained optimization.

• Enabling Large-Scale Applications: By reducing the projection cost from minutes/hours to milliseconds, it makes iterative methods (like Augmented Lagrangian or Projected Gradient Descent) feasible for problems with millions of variables ($n \sim 10^7$).
• Impact on Machine Learning & Risk Management: It facilitates the training of robust, fair, and risk-averse models where CVaR constraints are essential but previously computationally prohibitive at scale.
• Algorithmic Efficiency: It demonstrates that specialized combinatorial structures (Z-matrices and KKT monotonicity) can be exploited to bypass generic convex optimization bottlenecks, offering a blueprint for solving other structured projection problems.

The code is publicly available at https://github.com/jacob-roth/top-k-sum.