Coordinate Descent in Ordinary Least Squares not converging

Question

As a warm-up to writing my own elastic net solver, I'm trying to get a fast enough version of ordinary least squares implemented using coordinate descent.

So if $b = Ax$, I want to find $x$ that minimizes $$ \left\Vert b - Ax \right\Vert^2 $$ If $A_j$ is the $j$th column of $A$ and $A_{-j}$ is $A$ without column $j$, and the columns of $A$ are normalized so that $\left\Vert A_j \right\Vert^2 = 1$ for all $j$, the coordinate-wise update is then $$ x_j \gets A_j^T(b - A_{-j}x_{-j}) $$ I'm following along with these notes (page 9-10) but the derivation is simple calculus.

It's pointed out that instead of recomputing $A_j^T(b - A_{-j}x_{-j})$ all the time, a faster way to do it is with $$ x_j \gets A_j^Tr + x_j $$ where the total residual $r = b-Ax$ is computed once per convergence iteration. The equivalence of these update rules follows from noting that $Ax = A_jx_j + A_{-j}x_{-j}$ and rearranging terms.

My problem is that while the second method is indeed faster, it's wildly numerically unstable for me whenever the number of features isn't small compared to the number of samples. I was wondering if anyone might have some insight as to why that's the case. I should note that the first method, which is more stable, still starts disagreeing with more standard methods as the number of features approaches the number of samples.

Below is some Julia code for the two update rules:

function OLS_builtin(A,b)
    x = A\b
    return(x)
end

function OLS_coord_descent(A,b)    
    N,P = size(A)
    x = zeros(P)
    for iter in 1:1000
        for j = 1:P 
            x[j] = dot(A[:,j], b - A[:,1:P .!= j]*x[1:P .!= j])
        end    
    end
    return(x)
end

function OLS_coord_descent_fast(A,b) 
    N,P = size(A)
    x = zeros(P)
    for iter in 1:1000
        r = b - A*x
        for j = 1:P
            x[j] += dot(A[:,j],r)
        end    
    end
    return(x)
end

I generate data with the following:

n = 100
p = 50
σ = 0.1
ϐ = float([i*(-1)^i for i in 1:10])

β = append!(ϐ,zeros(Float64,p-length(ϐ)))
X = randn(n,p); X .-= mean(X,1); X ./= sqrt(sum(abs2(X),1))
y = X*β + σ*randn(n); y .-= mean(y);

I get good agreement between OLS_coord_descent(X,y) and OLS_builtin(X,y), whereas OLS_coord_descent_fast(X,y)returns exponentially large values for the regression coefficients. When p is less than about 20, OLS_coord_descent_fast(X,y) agrees with the other two.

Answer 1

The quick answer: You forgot to update r after each x[j] update. Following is the fixed function which behaves like OLS_coord_descent:

function OLS_coord_descent_fast(A,b) 
    N,P = size(A)
    x = zeros(P)
    for cycle in 1:1000
        r = b - A*x
        for j = 1:P
            x[j] += dot(A[:,j],r)
            r -= A[:,j]*dot(A[:,j],r)   # Add this line
        end    
    end
    return(x)
end

Answer source.