title: Fast Matrix Multiplication and Neural Networks
body: <p>I've used neural networks a little, but not much.  So as an attempt to increase my level of comfort, I decided to use one to play around with one of my favorite math problems: fast matrix multiplication.
  The standard algorithm takes O(n^3) to multiply two nxn matrices.  The Strassen algorithm does it in O(n^2.8).  The algorithms based off of work by Coppersmith and Winograd get down to O(n^2.373) but are impractical due to the large constant factor.
  There's a lot of wiggle room in between the latter two.  In particular, if you can multiply two 4x4 matrices using 48 multiply operations or less, you've done better than Strassen.</p>
  
  <p>So here's my setup: I have two (pseudo-randomly generated) nxn matrices, A and B.  One neural network takes NMULT linear combinations of elements of A and NMULT linear combinations of B, multiplies them together pointwise and then takes n^2 linear combinations of the output, trying to reconstruct the product AB.  The loss is the sum-of-squares error over the entries.
  The adversarial network takes two random matrices A' and B', and outputs softsign(A' + A_offset) and softsign(B' + B_offset), with loss function = -1 * sum-of-squares error of the other network.</p>
  
  <p>I alternate between 3 steps of training: training the fast-matrix-multiply network on random input matrices A and B, training the adversarial network on random input matrices A' and B', and training the fmm network on the output of the adversarial network.</p>
  
  <p>It doesn't work.  Not only can I not do better than Strassen, I can't even reproduce basic matrix multiplication!  That is, if I take n = 2 and NMULT = 8, I don't get down to 0 error.</p>
  
  <p>I know there are other (potentially better) ways of solving this problem than using neural networks -- I'm only doing this as a learning method.  Can anyone give me suggestions as to how to fix this?</p>
  
  <p>See code below:</p>
  
  <pre><code>import numpy as np
  import tensorflow as tf
  
  epochs=1000
  tot_batch = 1000
  learning_rate = 0.01
  
  MATRIX_SIZE = 2
  NMULTS = 8
  
  nvals = MATRIX_SIZE * MATRIX_SIZE
  
  # These are the inputs to the adversarial NN generating our input matrices A&amp;B.
  a_inputs = tf.placeholder(tf.float32, [None, nvals])
  b_inputs = tf.placeholder(tf.float32, [None, nvals])
  
  adv_a_icpt = tf.Variable(tf.random_normal([nvals]))
  adv_b_icpt = tf.Variable(tf.random_normal([nvals]))
  
  a_vector = tf.nn.softsign(a_inputs + adv_a_icpt)
  b_vector = tf.nn.softsign(b_inputs + adv_b_icpt)
  
  # These are the two adversarial matrices we are multiplying; all entries
  # are in [-1, 1]. This makes normalizing the error easier.
  a_matrix = tf.reshape(a_vector, [-1, MATRIX_SIZE, MATRIX_SIZE])
  b_matrix = tf.reshape(b_vector, [-1, MATRIX_SIZE, MATRIX_SIZE])
  
  # This is the product A * B.
  m_matrix = tf.matmul(a_matrix, b_matrix)
  
  # This is what the fast-matrix-multiply NN will be predicting.
  m_vector = tf.reshape(m_matrix, [-1, nvals])
  
  fmm_a_wts = tf.Variable(tf.random_normal([nvals, NMULTS]))
  fmm_b_wts = tf.Variable(tf.random_normal([nvals, NMULTS]))
  fmm_output_wts = tf.Variable(tf.random_normal([NMULTS, nvals]))
  
  # This is the output of the fast-matrix-multiply NN.
  def fmm_output(input_a_vec, input_b_vec):
    hidden_a_inputs = tf.matmul(input_a_vec, fmm_a_wts)
    hidden_b_inputs = tf.matmul(input_b_vec, fmm_b_wts)
    hidden_output = tf.multiply(hidden_a_inputs, hidden_b_inputs) 
    return tf.matmul(hidden_output, fmm_output_wts)
  
  # Treat each element of the input arrays as having a variance of O(1). Then
  # the output array elements have a variance of O(MATRIX_SIZE).
  loss_adv = tf.divide(
      tf.losses.mean_squared_error(m_vector, fmm_output(a_vector, b_vector)),
      MATRIX_SIZE)
  abs_err_vec_adv = tf.abs(tf.subtract(m_vector, fmm_output(a_vector, b_vector)))
  mean_abs_err_adv = tf.reduce_mean(abs_err_vec_adv, reduction_indices=[1])
  
  m_rand = tf.matmul(tf.reshape(a_inputs, [-1, MATRIX_SIZE, MATRIX_SIZE]),
                     tf.reshape(b_inputs, [-1, MATRIX_SIZE, MATRIX_SIZE]))
  loss_rand = tf.divide(
      tf.losses.mean_squared_error(tf.reshape(m_rand, [-1, nvals]),
                                   fmm_output(a_inputs, b_inputs)),
      MATRIX_SIZE) 
  
  
  optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
  
  train_ADV = optimizer.minimize(-loss_adv,
                                 var_list=[adv_a_wts, adv_b_wts,
                                           adv_a_icpt, adv_b_icpt])
  
  train_FMMA = optimizer.minimize(loss_adv,
                                  var_list=[fmm_a_wts, fmm_b_wts,
                                            fmm_output_wts])
  
  train_FMMR = optimizer.minimize(loss_rand,
                                  var_list=[fmm_a_wts, fmm_b_wts,
                                            fmm_output_wts])
  
  
  init = tf.global_variables_initializer()
  
  with tf.Session() as sess:
    sess.run(init)
  
    adv_batch_size = 100
    fmm_batch_size = 100
  
    for epoch in range(epochs):
      adv_loss = 0.0
      rand_loss = 0.0
      for i in range(tot_batch):
        # Run the fast-matrix-multiply NN training against random inputs.
        batch_a_inputs = np.random.uniform(low=-1., size=[fmm_batch_size, nvals])
        batch_b_inputs = np.random.uniform(low=-1., size=[fmm_batch_size, nvals])
        _, rerr = sess.run([train_FMMR, loss_rand],
                           feed_dict={ a_inputs : batch_a_inputs,
                                       b_inputs : batch_b_inputs })
  
        # Run the adversarial NN training.
        batch_a_inputs = np.random.normal(size=[adv_batch_size, nvals])
        batch_b_inputs = np.random.normal(size=[adv_batch_size, nvals])
        sess.run(train_ADV, feed_dict={ a_inputs : batch_a_inputs,
                                        b_inputs : batch_b_inputs })
  
        # Run the fast-matrix-multiply NN training against adversarial inputs.
        batch_a_inputs = np.random.normal(size=[fmm_batch_size, nvals])
        batch_b_inputs = np.random.normal(size=[fmm_batch_size, nvals])
        _, aerr, mae = sess.run([train_FMMA, loss_adv, mean_abs_err_adv],
                           feed_dict={ a_inputs : batch_a_inputs,
                                       b_inputs : batch_b_inputs })
  
        adv_loss += aerr / tot_batch
        rand_loss += 3.0 * rerr / tot_batch
        if i % 200 == 0:
          print("Batch " + str(i) + ", mean abs error: " + str(mae[0:4]))
      print("Epoch: " + str(epoch) + ", rand loss = " + str(rand_loss) +
            ", adv loss = " + str(adv_loss))
  </code></pre>