Nuclear norm

In this example, we use the PANOC solver to solve a problem with a nuclear norm as regularizer, in order to promote low-rank solutions.

\[\begin{aligned} & \minimize_X && \tfrac12 \normsq{AX - B}_F + \lambda \norm{X}_* \end{aligned}\]

Here, \(X \in \R^{m\times n}\), with a rank \(r\) that’s much smaller than its dimensions.

The JAX package is used to compute gradients (although we could easily do without, since the loss is quadratic).

# %% alpaqa nuclear norm example

from pprint import pprint

import jax
import jax.numpy as jnp
import jax.numpy.linalg as jla
from jax import grad, jit, random

import alpaqa as pa

jax.config.update("jax_enable_x64", True)

# %% Define the problem functions

# Quadratic loss plus nuclear norm regularization
#
# minimize  ½‖vec(AX - B)‖² + λ nucl(X)


# Returns the loss function with constant A and B
def loss(A, B):
    def _loss(X):
        err = (A @ X - B).ravel()
        return 0.5 * jnp.dot(err, err)

    return _loss


class MyProblem(pa.UnconstrProblem):
    def __init__(self, A, B, λ):
        self.rows, self.cols = A.shape[1], B.shape[1]
        super().__init__(self.rows * self.cols)
        f = loss(A, B)
        self.jit_loss = jit(f)
        self.jit_grad_loss = jit(grad(f))
        self.reg = pa.functions.NuclearNorm(λ, self.rows, self.cols)

    def eval_objective(self, x):  # Cost function
        # Important: use consistent order when reshaping or raveling!
        X = jnp.reshape(x, (self.rows, self.cols), order="F")
        return self.jit_loss(X)

    def eval_objective_gradient(self, x, grad_f):  # Gradient of the cost
        X = jnp.reshape(x, (self.rows, self.cols), order="F")
        grad_f[:] = self.jit_grad_loss(X).ravel(order="F")

    def eval_proximal_gradient_step(self, γ, x, grad, x_hat, p):
        # use the prox_step helper function to carry out a generalized
        # forward-backward step. This assumes Fortran order (column major),
        # so we have to use order="F" for all reshape/ravel calls
        return pa.prox_step(self.reg, x, grad, x_hat, p, γ, -γ)


# %% Generate some data

m, n = 30, 30
r = m // 5

key = random.PRNGKey(0)
key, *subkeys = random.split(key, 6)
# Random rank-r data matrix X = UV, then add some noise
U_true = random.uniform(subkeys[0], (m, r), minval=-1, maxval=1)
V_true = random.uniform(subkeys[1], (r, n), minval=-1, maxval=1)
B_rand = 0.01 * random.normal(subkeys[2], (m, n))
A = random.uniform(subkeys[3], (m, m), minval=-1, maxval=1)
X_true = U_true @ V_true
B = A @ X_true + B_rand  # Add noise

print("cond(A) =", jla.cond(A))
print("inf(B) =", jla.norm(B_rand.ravel(), jnp.inf))

prob = MyProblem(A, B, λ=1 * m)

# %% Solve the problem using alpaqa's PANOC solver

opts = {
    "max_iter": 2000,
    "stop_crit": pa.FPRNorm,
    "quadratic_upperbound_tolerance_factor": 1e-14,
}
direction = pa.LBFGSDirection({"memory": 100})
# direction = pa.AndersonDirection({"memory": 5})
solver = pa.PANOCSolver({"print_interval": 10} | opts, direction)

# Add callback to the solver
residuals = []


def callback(it: pa.PANOCProgressInfo):
    residuals.append(it.ε)


solver.set_progress_callback(callback)

# Add evaluation counters to the problem
cnt = pa.problem_with_counters(prob)
# Solve the problem
sol, stats = solver(cnt.problem, {"tolerance": 1e-10})
X_sol = jnp.reshape(sol, (m, n), order="F")

# %% Print the results

final_f = prob.eval_objective(sol)
print()
pprint(stats)
print()
print("Evaluations:")
print(cnt.evaluations)
print(f"Cost:          {final_f + stats['final_h']}")
print(f"Loss:          {final_f}")
print(f"Inf norm loss: {jla.norm((X_sol - X_true).ravel(), jnp.inf)}")
print(f"Regularizer:   {stats['final_h']}")
print(f"Rank:          {jla.matrix_rank(X_sol)}")
print(f"True rank:     {r}")
print(f"FP Residual:   {stats['ε']}")
print(f"Run time:      {stats['elapsed_time']}")
print(stats["status"])

# %% Plot the results

import matplotlib.pyplot as plt

plt.figure()
plt.semilogy(residuals, ".-")
plt.title("PANOC nuclear norm example: residuals")
plt.xlabel("Iteration")
plt.tight_layout()

if m * n <= 5000:
    plt.figure(figsize=(8, 5))
    plt.plot((A @ X_sol).ravel(), ".-", label="Estimated solution $ A\\tilde X $")
    plt.plot((A @ X_true).ravel(), "x:", label="True solution $ AX $")
    plt.plot(B.ravel(), "*-", label="Constant $ B $")
    plt.legend()
    plt.title("PANOC nuclear norm example: solution")
    plt.tight_layout()

plt.show()