Implement default preconditioner for non AddedDiagLazyTensors?

[jacobrgardner]

Right now, we precondition linear solves and log determinants for AddedDiagLazyTensors. This instance of the preconditioner uses the Woodbury formula to solve with L_{k}L_{k}^{\top} + s*I, and the matrix determinant lemma to get the log determinant of that matrix (see the NIPS paper).

Absent an added diagonal, I think L_{k} could still be useful for preconditioning at least linear systems by solving a least squares problem:

# New proposed default for LazyTensor
def _preconditioner(self):
    piv_chol_self = pivoted_cholesky(...)
    def precond_closure(v):
        # Solve least squares problem with L_{k}L_{k}^{\top}
    return precond_closure, None

Here's the catch: I think this default could only be used for inv_matmul, not inv_quad_log_det. At least, I'm not sure how to get a good log determinant correction out of this like we do for the added diagonal case. However, this could make our inv_matmul calls for non AddedDiagLazyTensors all magically much better (of course we should test this).

@gpleiss

[jacobrgardner]

If I've done math right, one possible inner definition of precond_closure(v) is:

def precond_closure(v):
    q, r = torch.qr(piv_chol_self) # Decompose L_{k}, use to solve with L_{k}L_{k}^{\top}
    new_rhs = q.t().matmul(v)
    solve = q.matmul(torch.trtrs(torch.trtrs(new_rhs, r), r, transpose=True))
    return solve

[jacobrgardner]

@gpleiss Initial results suggest this is in fact a terrible idea :-)

I'll keep pushing on it and try to understand exactly why this is failing.