scipy.optimize.minimize fails to converge but result is OK

Question

I am trying to optimize a non-linear least squares problem with scipy.optimize.minimize. I have simplified my actual problem down to the case where I am just computing the top 'principal components' like in a PCA analysis and the method reports failure to converge (as it does for my full model). However, the result matches reasonably closely the expected result that you get through SVD decomposition as usually done in PCA. Note that my full problem has no closed-form solution, unlike this simplified one.

My questions are: can I fix the convergence issue? If not, can I ignore the non-convergence status (in the more complicated models)?

In short, I want to find the loadings of the top r PCA components of a matrix, i.e. a k x r orthogonal matrix L such that XL maximizes its variance (or X L L^T is as close to X as possible). I parametrize the orthogonal matrices with a skew-symmetric matrix C and use the matrix exponential map to convert it to L.

I have done the following 'obvious' checks:

1. NaNs/Infs - I don't see how any could occur and I don't observe any in the outputs.
2. Needing to rescale the loss function to be not huge/not small. Loss is around 10-20 in the range I am dealing with. The Jacobian is pretty small but non-zero at the final iteration (values around 1e-3 to 1e-5).
3. Non-differentiable or poorly behaved loss function. It's squared sum residuals and it looks OK to me. The matrix exponential should be well-behaved (finite) in the region where we are in, with the matrix norms all about 1 or lower.
4. Alternative optimization algorithms. None that I've tried have worked, including BFGS, CG, Newton-CG, and Nalder-Mead (which I think would be too slow for my real data anyway).

Below is a reproducible example on some small simulated dataset. I am using JAX to get the Jacobians (similarly, using jax.scipy.optimize.minimize gives convergence problems too, returning with status=3 indicating that the zoom step of BFGS failed). Disabling JAX jit doesn't change the results.

import jax
import jax.scipy.optimize
from jax import numpy as jnp
import scipy.optimize
import numpy

# Simulate data with one large PC in the first two components
N = 30
scores = numpy.linspace(-1,1, N)
simd = numpy.concatenate([
    [10*scores],
    [10*scores],
    numpy.random.normal(size=(1,N))*5,
    numpy.random.normal(size=(5,N))*0.01
], axis=0).T

# Optimal solution, via SVD
def low_rank_loadings(X, r):
    ''' give best rank r loadings to approximation X '''
    u,d,vt = numpy.linalg.svd(X, full_matrices=False)
    #return u[:,:r] @ numpy.diag(d[:r]) @ vt[:r,:]
    return vt[:r,:].T

# Loss function
@jax.jit
def eval(C, L0, X):
    def func(i, ssr):
        x = X[[i],:]
        L = jax.scipy.linalg.expm(C) @ L0
        return ssr + jnp.linalg.norm(x - x @ L @ L.T)**2
    ssr = jax.lax.fori_loop(
        0,
        len(X),
        func,
        init_val = jnp.asarray([0.]),
    )
    return ssr / len(X)
def extract(vars, L0):
    # Pull out our C matrix from a flattened vector
    k, r = L0.shape
    N = k*r - (r*(r+1)//2) # Num free vars per matrix
    C_ = vars[0:N]
    # Convert to rectangular lower triangular matrices
    idxs = jnp.tril_indices(n=k, k=-1, m=r)
    Ctri = jnp.zeros((k,r)).at[idxs].set(C_)
    # Convert to the right coordinates to match L0
    C = Ctri @ L0.T - L0 @ Ctri.T
    return C

def jax_pca(X, r):
    ''' X = data matrix, r = rank PCA to compute '''
    n,k = X.shape
    L0 = jnp.eye(k, r) # Start with projection to first r components; a bad PCA
    X = jnp.asarray(X)
    def f(vars):
        C = extract(vars, L0)
        return eval(C, L0, X)[0]
    N = k*r - (r*(r+1)//2) # Num free vars per matrix
    res = scipy.optimize.minimize(
        f,
        x0 = numpy.zeros(N),
        jac = jax.jacrev(f),
        method = "BFGS",
    )
    return (L0, extract(res.x, L0), res)
L0, C, res = jax_pca(simd, 1)
print(f"Success? {res.success} -  {res.message}")
loadings = jax.scipy.linalg.expm(C) @ L0
optimal_loadings = low_rank_loadings(simd, 1)
print("Computed:")
print(loadings)
print("Optimal:")
print(optimal_loadings)

Which gives output:

Success? False -  Desired error not necessarily achieved due to precision loss.
Computed:
[[ 6.9389027e-01]
 [ 6.9399828e-01]
 [-1.9204833e-01]
 [-3.3489120e-04]
 [ 3.8082333e-05]
 [ 1.3443755e-04]
 [ 1.8192278e-04]
 [ 2.3052096e-04]]
Optimal:
[[ 6.93945195e-01]
 [ 6.93945195e-01]
 [-1.92041436e-01]
 [-3.36661681e-04]
 [ 3.91581503e-05]
 [ 1.34102010e-04]
 [ 1.81875549e-04]
 [ 2.31636442e-04]]