Simultaneous Inverse Matrix Computations in TensorFlow for Orthogonal Matching Pursuit

Question

I am currently working on creating a version of Orthogonal Matching Pursuit (OMP) that can be performed simultaneously on different patches, utilizing the power of TensorFlow for optimization through static graph compilation.

I have provided the pseudocode for a single step of the main loop of the algorithm below:

support = tf.TensorArray(dtype=tf.int64, size=1, dynamic_size=True)

# For each element in patches, compute the projection on the dictionary using only TensorFlow API
dot_products = tf.abs(tf.matmul(A, patches, transpose_a=True))
max_index = tf.argmax(dot_products, axis=1)
support = support.write(i, max_index + 1)
support = support.write(i + 1, max_index + 1)
idx = support.stack()
non_zero_rows = tf.reduce_all(tf.not_equal(idx, 0), axis=1)

idx = tf.boolean_mask(idx, non_zero_rows) - 1
A_tilde = tf.gather(A, idx, axis=1)
m, n = A.shape
selected_atoms = tf.matmul(A_tilde, A_tilde, transpose_a=True)

After obtaining selected_atoms, which is a 3D tensor consisting of l matrices of size nxm, I need to solve the least square problem |patches - selected_atoms * x_coeff| ** 2 To do this, I need to compute the inverse of selected_atoms.T @ selected_atoms. Is there a method using TensorFlow's API to simultaneously compute the l inverse matrices and store them in a 3D tensor of shape lxnxm? I would greatly appreciate any guidance.

Thank you.

Answer 1

To get the pseudoinverse, can use the pinv function, it does what you want (for least squares, pseudo-inverse is well-defined whereas inverse might not be).

import tensorflow as tf
x_batch = tf.constant([[[1.,  0.4,  0.5],
                 [0.4, 0.2,  0.25],
                 [0.5, 0.25, 0.35]],
                [[0.5,  0,  0],
                 [0, 0.5,  0],
                 [0, 0, 0.5]]])
tf.linalg.pinv(x_batch)

the function call returns:

tf.Tensor: shape=(2, 3, 3), dtype=float32, numpy=
array([[[ 4.9999995e+00, -9.9999933e+00, -4.2915344e-06],
        [-9.9999905e+00,  6.6666672e+01, -3.3333347e+01],
        [-6.1988831e-06, -3.3333344e+01,  2.6666681e+01]],

       [[ 2.0000000e+00,  0.0000000e+00,  0.0000000e+00],
        [ 0.0000000e+00,  2.0000000e+00,  0.0000000e+00],
        [ 0.0000000e+00,  0.0000000e+00,  2.0000000e+00]]], dtype=float32)

If you only want the coefficients of the regression, you don't need to store pseudo_inverse, but the pinv*patches (which might save a lot of memory) see: http://www.seas.ucla.edu/~vandenbe/133A/lectures/ls.pdf