[Python] Convergence of Coordinate Descent

I am trying to implement the algorithms described in the paper Learning Fast Approximations of Sparse Coding and completed the ISTA and FISTA algorithms and the dictionary update rule using projected block coordinate descent.

I ran into a problem while implementing Coordinate Descent (Algorithm 2 in the paper) for approximating the optimal sparse codes though, since the code does not seem to converge. A similar problem occurred with ISTA too, but I quickly figured out that the learning rate was smaller than the largest eigenvalue of DTD, which violates one of the convergence criteria. However, since the Coordinate Descent algorithm does not have such hyperparameters except the regularization, I am confused.

I got some suggestions to try tweaking the regularization parameter, but it does not seem to help a bit.

Minimal reproducible example


import torch
import itertools

def shrink(a, b):
return torch.sign(a) * torch.maximum(torch.abs(a) - b, torch.tensor(0.0))

def change(a, b):
return torch.sum(torch.abs(a - b))

def CoD(x : torch.Tensor,
h_dim : int,
D : torch.Tensor,
regularization=0.5,
frequency=None) -> torch.Tensor:
"""
Implements the Coordinate Descent algorithm to find the optimal
sparse codes for the given input vector x and a given
dictionary matrix D. This function follows notation corresponding to
the paper "Learning Fast Approximations of Sparse Coding"

aram x: The input vector
aram h_dim: Dimension of sparse code required. In
the overcomplete case, should be greater than or equal to
dim(x).
aram D: the dictionary matrix
aram regularization: parameter which controls the relative
weight given to sparsity vs reconstruction loss
aram lr: the learning rate
aram frequency: the number of iterations after which an update is printed
:return: An approximation to the optimal sparse code of the
given input h_opt.
"""
n = x.shape[0]
m = h_dim
assert D.shape == (n, m), (f"D should be matrix of shape (x_dim, h_dim), where "
f"x_dim={n} and h_dim={m}, but dimensions of D are {D.shape}")
z = torch.zeros(h_dim)
z_old = z.clone()

S = torch.eye(h_dim) - torch.matmul(torch.transpose(D, 0, 1), D)
B = torch.matmul(D.T, x)

counter = itertools.count(start=0, step=1)
while True:
z = shrink(B, regularization)
k = torch.argmax(torch.abs(z - z_old))
for j in range(m):
B[j] = B[j] + S[j][k] * (z[k] - z_old[k])
if change(z, z_old) < 0.01:
break
z_old = z.clone()
iter_num = next(counter)
if frequency is not None and iter_num % frequency == 0:
print(f"Coordinate Descent: Iteration {iter_num}")
return shrink(B, regularization)

print(CoD(torch.randn(10), 10, torch.randn(10, 10), frequency=100))


I tried to implement the algorithm almost exactly as described in the paper. Can someone help me with the debugging process?

Continue reading...