Robust in linear time

[pavelkomarov]

Section 4 of https://www.sciencedirect.com/science/article/pii/S0169743903001370 is about efficient algorithms. There are tricks to make the robust case run in linear complexity w.r.t. $N$ rather than cubic.

[pavelkomarov]

I have now implemented the robust solve in several ways. All return essentially the same answers.

First up is this sparse-matrices-based version:

from scipy import sparse, optimize

def sparse_smooth(y, A, Q, C, R, meas_huberM=1.345):
    """Solve the optimization problem for robust smoothing using sparse solvers to keep runtime down.
    """
    N = len(y)
    l = 1e-7 # regularization

    Q_sqrt_inv = sparse.csr_matrix(np.linalg.inv(sqrtm(Q)))
    R_sqrt_inv = sparse.csr_matrix(np.linalg.inv(sqrtm(R)))

    QA = Q_sqrt_inv @ sparse.csr_matrix(A)
    I1 = sparse.eye(N-1, N, k=0, format='csr') # 1s on main diagonal
    I2 = sparse.eye(N-1, N, k=1, format='csr') # 1s on shifted diagonal
    bigG = (sparse.kron(I2, Q_sqrt_inv) - sparse.kron(I1, QA)) # form sparse G

    I = sparse.eye(N, format='csr')
    bigR = sparse.kron(I, R_sqrt_inv)
    ytilde = bigR @ y

    bigC = sparse.kron(I, sparse.csr_matrix(C))
    Ctilde = bigR @ bigC

    GTG = (bigG.T @ bigG + l*sparse.eye(N*A.shape[0], format='csr')).tocsc()
    solver = sparse.linalg.factorized(GTG) # does LU factorization

    def obj_and_grad(v):
        u = solver(Ctilde.T @ v)
        return 0.5*(v @ Ctilde @ u + np.sum(v**2)) - ytilde @ v, v + Ctilde @ u - ytilde

    res = optimize.minimize(obj_and_grad, np.zeros(N), method='trust-constr', jac=True,
                        bounds=optimize.Bounds(-meas_huberM*np.ones(N), meas_huberM*np.ones(N)))
    vopt = res.x
    xopt = solver(Ctilde.T @ vopt) # vector of dimension N\nu
    return xopt.reshape(N, A.shape[0])

The idea here is to assume a quadratic penalty on the process term and Huber penalty on the measurements, expand to a form that covers all steps in a big sparse system (featuring matrix $G$), substitute a supremum expression over a new variable $v$ (from https://arxiv.org/pdf/1609.06369) in for the Huber loss, switch the order of the $\min$ and the $\sup$, solve for $x$ in terms of $v$, solve the supremum/maximization/-minimization for $v$, and then use that solution to get the optimal $x$. It's a bit of a mouthful when you say it that way. The actual optimization for $v$ can be accomplished with trust-constr, because it's limited to a box between $[-M, M]$, where $M$ is the Huber parameter. It takes an objective function that returns the value of the objective and the objective's gradient at a guess of $v$, where these are calculated using an intermediate sparse solve for $u$, which subsumes parts of the expressions involving inversion of $G^T G + \lambda \mathbb{I}$. This all turns out to be pretty expensive in practice, taking much longer than my CVXPY/CLARABEL implementation:

The horizontal axis is simulation length in "seconds" with $\Delta t = 0.01$, so $N = 100\times$ the tick marks, and the vertical axis is runtime in seconds on my machine.

This implementation is intended by construction to be theoretically linear in $N$, so it was disappointing to see this result, with even a bit of a polynomial bend in the runtime plot. But it made me curious about the CVXPY runtime, because at this level of zoom it also looks pretty linear, so perhaps CLARABEL was always doing the right thing.

Next, at the suggestion of @saravkin, I tried CP-V2, which works by iteratively stepping vectors $x$ and $\omega$ using proximity operators, which essentially smooth out/damp the piecewise loss. It requires derivation of those operators as well as a conjugate function, so it took some convex optimization knowledge to formulate. Here is a PDF with my math: CP-V2.pdf Beware there are errata: two stemming from https://arxiv.org/pdf/1609.06369, and one introduced by me. The two are corrected here:

And the last is a spurious factor of 2 in my derivation of the solution to $\text{prox}_{\tau g}$. These issues have been corrected in the code below. Tinkering with said code is what alerted me to their presence:

def chambolle_pock(y, A, Q, C, R, meas_huberM):
    N = len(y)
    
    Q_sqrt_inv = sparse.csr_matrix(np.linalg.inv(sqrtm(Q)))
    R_sqrt_inv = sparse.csr_matrix(np.linalg.inv(sqrtm(R)))

    QA = Q_sqrt_inv @ sparse.csr_matrix(A)
    I1 = sparse.eye(N-1, N, k=0, format='csr') # 1s on main diagonal
    I2 = sparse.eye(N-1, N, k=1, format='csr') # 1s on shifted diagonal
    bigG = (sparse.kron(I2, Q_sqrt_inv) - sparse.kron(I1, QA)) # form sparse G

    I = sparse.eye(N, format='csr')
    bigR = sparse.kron(I, R_sqrt_inv)
    r = bigR @ y

    bigC = sparse.kron(I, sparse.csr_matrix(C))
    K = bigR @ bigC

    eps = 1e-3
    U, S, VT = sparse.linalg.svds(K, k=1, which='LM') # 'LM' = largest magnitude
    L = S[0] # largest magnitude singular value, always very close to 1 for my systems
    tau = 1e3 - eps
    sig = 1/(L**2 * tau)
    print(L, tau, sig)

    GTG = (bigG.T @ bigG + sparse.eye(N*A.shape[0], format='csr')/tau).tocsc()
    solver = sparse.linalg.factorized(GTG) # LU factorization

    x = np.zeros(N * A.shape[0])
    w = np.zeros(N)
    x_prev = np.ones(x.shape)
    w_prev = np.ones(w.shape)

    #def huber(x, M):
    #    absx = np.abs(x)
    #    return np.where(absx <= M, 0.5*x**2, M*(absx - 0.5*M))

    #k = 0
    while np.linalg.norm(w - w_prev)/np.sqrt(N) + np.linalg.norm(x - x_prev)/np.sqrt(len(x)) > eps:
        #if k % 1000 == 0:
        #    obj = 0.5*np.linalg.norm(bigG @ x)**2 + np.sum(huber(K @ x - r, meas_huberM))
        #    print(obj)
        #k += 1

        v = w + sig*K @ (2*x - x_prev) - sig*r # this r needs the *sig
        w_prev = w
        w = np.clip(v/(1+sig), -meas_huberM, meas_huberM) # the +r here was spurious

        v = x - tau*K.T @ w
        x_prev = x
        x = solver(v/tau) # This is the big work

    return x.reshape(N, A.shape[0])

This is where the runtimes stood at this point:

The trouble with this method is that it requires picking $\epsilon$, $\tau$, and $\sigma$, and although I got it to run faster than my CVXPY implementation, and with equal reliability, it was only by playing with these parameters a while. Changing the simulation/data fed to the algorithm made no difference, but changing the order of the system did, meaning I'd have to sit there picking good parameters for the three different cases. Inconvenient, to say the least. Moreover, swapping out a loss function requires re-deriving different proximity operators, which is a bit fun to the mathematician in me, but is a somewhat unreasonable lift for standard practitioners. So users would be stuck with the cost formulation we give them. Indeed, we need to use a quadratic loss on the process term here to make the problem strongly convex, else Chambolle-Pock can not be applied successfully. This method is thus not extensible to situations with sparse state evolution, for instance.

I was prepared to stomach the above, but in the process of experimenting with longer data series I got a CVXPY warning saying my implementation wasn't vectorized. I spent some time trying to vectorize it only to end up with segfaults, but I finally found a matrix/vector shape disagreement that explains it. Reshaping, the code works in vectorized form, and it turns out to be wicked fast:

I tried using Huber loss on both process and measurement terms, and it's still blazingly fast, only a small factor slower than RTS. It turns out the professionals wrote a killer solver, and my main issue was problem compilation time, not solution time. This is the most easily extensible way to go, so the other two solvers won't make it in to the repo; I'll just leave them here for posterity and update the convex_smooth function to be properly vectorized and have the correct balance of loss terms per https://jmlr.org/papers/volume14/aravkin13a/aravkin13a.pdf.