Just wondering if anyone in the pytorch community had come across this physics paper from 2020? They discovered a method of speeding up Principle Component Analysis due to its equivalence to the Schrodinger equation. Its only fully accurate up to second order, but higher order terms are generally small enough for it to matter less. In many use cases a faster PCA is well worth a minor accuracy penalty.
To try it out I went to Versatile Diffusion and substituted the elements of the torch.pca_lowrank() function for the schrodinger equation.
def solve_schrodinger_eq(u, s, v, potential_func, num_eigenstates):
# Convert the variance to potential
potential = potential_func(s)
# Compute the inverse mass from the correlation kernel
inverse_mass = torch.diag(s ** -1)
# Compute the eigenvalues and eigenvectors of the Hamiltonian matrix
hamiltonian = torch.mm(torch.mm(u, torch.diag(potential)), v.t()) + inverse_mass
eigenvalues, eigenvectors = torch.symeig(hamiltonian, eigenvectors=True)
# Sort the eigenvalues and eigenvectors in ascending order
eigenvalues, indices = eigenvalues.sort()
eigenvectors = eigenvectors[:, indices]
# Select the num_eigenstates lowest-energy eigenvalues and eigenvectors
eigenvalues = eigenvalues[:num_eigenstates]
eigenvectors = eigenvectors[:, :num_eigenstates]
# Normalize the eigenvectors
norms = torch.sqrt(torch.sum(eigenvectors ** 2, dim=0))
eigenvectors = eigenvectors / norms
# Return the eigenvalues and eigenvectors (renamed as eigen-energies and eigenstates, respectively)
return eigenvalues, eigenvectors
# Example usage
x_input = torch.rand(10, 10)
q = 0.5
niter = 10
num_eigenstates = 5
# Compute the PCA
u, s, v = torch.pca_lowrank(x_input, q=q, center=False, niter=niter)
# Define the potential function
potential_func = lambda x: -torch.log(x)
# Solve the Schrodinger equation
eigenvalues, eigenvectors = solve_schrodinger_eq(u, s, v, potential_func, num_eigenstates)
# Rename the eigenvectors as eigenstates
eigenstates = eigenvectors
# Print the results
The other beauty of this approach is it should work extremely quickly in qbits.
The authors also did a neat little video where they explain what they found.
Thanks for sharing this approach. I get a few errors trying to execute it:
R = torch.randn(n, q, dtype=dtype, device=A.device)
TypeError: randn(): argument 'size' must be tuple of ints, but found element of type float at pos 2
which points to the
q=0.5 argument. Also:
RuntimeError: This function was deprecated since version 1.9 and is now removed. Please use the `torch.linalg.eigh` function instead.
is raised. Which PyTorch version did you use to execute the code?
apologies. ptrblack. Here is a direct comparison in 1.12.0
# Create some test data.
variances = torch.tensor([1.0, 2.0, 3.0])
correlation_kernel = torch.tensor([0.5, 0.3, 0.2])
def schrodinger_pca(variances, correlation_kernel):
Runs the Schrödinger equation for a quantum steady state problem.
variances: A tensor of variances.
correlation_kernel: A tensor of correlation kernels.
A tuple of eigen-energies and eigenstates.
# Convert the variances and correlation kernels to Schrödinger equation parameters.
potential = torch.diag(variances)
inverse_mass = torch.diag(1.0 / correlation_kernel)
# Solve the Schrödinger equation.
eigen_energies, eigenstates = torch.linalg.eigh(potential @ inverse_mass)
return eigen_energies, eigenstates
# Run the Schrödinger PCA function.
start = time.time()
eigen_energies_schrodinger, eigenstates_schrodinger = schrodinger_pca(variances, correlation_kernel)
end = time.time()
time_schrodinger = end - start
# Run the original PCA function.
start = time.time()
A = torch.diag(variances) @ torch.diag(1.0 / correlation_kernel)
pca = torch.pca_lowrank(A)
end = time.time()
time_pca = end - start
# Print the results.
print("Schrödinger Egeinenergies", eigen_energies_schrodinger)
print("Schrödinger Egein states", eigenstates_schrodinger)
# Print the eigenvalues.
print("PCA Eigenvalues:", pca)
# Print the eigenvectors.
print("PCA Eigenvectors:", pca)
# Print the results.
print("Schrödinger PCA time:", time_schrodinger)
print("Lowrank PCA time", time_pca)
# Compare the results.
print("Are the results the same?:", torch.allclose(eigen_energies_schrodinger, pca))
The key thing I’m conveying here is how the same inputs can be used for both calculations. There is a bunch of existing literature on quantum PCA, but seeing this paper made me realise it was applicable to pytorch tensors using our existing code bases.
When you run this code you’ll see the results aren’t an exact match. The schrodinger version is giving back exact principal components not estimated weightings. My understanding is that the larger the sample data the closer the lowrank function’s estimate will get to the schrodinger answer.
In addition to exact answer, the key benefit of the Schrodinger version is that it is very amenable to calculation on Qbits, since it is literally a quantum waveform calculation. This is such an advantage that the speed advantages of lowrank are eclipsed and you’re left with a solution which is both more accurate and faster. This is because wave form collapse of a Qbit array doesn’t scale as slowly with input vectors, they are computed simultaneously, so very large input datasets can be processed near instantly.
In the case of a workload like AI training this may drastically reduce the training and processing times. Theoretically by orders of magnitude in some circumstances. There are similar applications for climate modelling, financial modelling or other high dimensionality problems.
For all these use cases it seems a fairly simply pytorch interface layer could shift a lot of the workload from classical to quantum computing using an existing codebase.