Add eager mode gradients for ops.

[nsthorat]

From @nsthorat on January 18, 2018 15:56

The infrastructure for eager mode is now ready for gradient methods to be filled in!

Eager mode provides a new set of methods on NDArrayMath which allows the user to eagerly compute gradients. Most users will use an optimizer like this:

const weights = dl.variable(Array2D.randNormal([784, 10]));
const cost = optimizer.minimize(() => {
  const batch = data.nextTrainBatch(BATCH_SIZE);
  const y = math.matMul(batch.xs, weights);
  const loss = math.mean(math.softmaxCrossEntropyWithLogits(labels, ys));
  return loss;
});

You'll notice that there is no use of the Graph, we simply use ops on NDArrayMath directly inside of an optimizer.minimize() method.

You can find a full example of training MNIST in eager mode here: https://github.com/PAIR-code/deeplearnjs/blob/master/demos/mnist_eager/model.ts

As part of NDArrayMath we expose several new methods. The important ones are these:

• math.gradients(f: () => cost, xs) which executes f() (which produces a scalar value) and returns the gradient of the output of f with respect to xs (which can be an NDArray or a string => NDArray map).
• math.valueAndGradients(f: () => cost, xs) which is the same as math.gradients() but also returns the output of f().
• math.vjp(f: () => y, x, dy) which computes a vector-jacobian product - it is similar to gradients, but allows f() to produce a non-scalar value and lets the user provide a dy. This is useful to compute a subset of backpropagation, or to tests gradients of a single op with a provided dy (this is how we unit test).
• math.customGradient(f: () => {value, gradients}, xs) which allows the user to provide a custom gradient of an arbitrary function closure instead of using the default gradients of the ops in the function. We use this for numerical stability for ops like softmaxCrossEntropy, and for mean / sum so we can compute a faster gradient (instead of the combination of gradients of the kernels they use). Most of the time, you shouldn't need to use this.

Now that these methods exist and are relatively stable, we can flush out gradients for kernels and ops!

To add gradients for kernels, we simply need to add a derivative function to the executeKernel calls inside of NDArrayMath. An example:

const der = (dy: Array2D<'float32'>, y: Array2D) => {
  return {
    a: () => this.matMul(dy, b, MatrixOrientation.REGULAR, MatrixOrientation.TRANSPOSED),
    b: () => this.matMul(a, dy, MatrixOrientation.TRANSPOSED, MatrixOrientation.REGULAR)
  };
};
return this.backendEngine.executeKernel(
  'MatMul', {inputs: {a, b}, args: {aOrientation, bOrientation}}, der);

The derivative is an function that takes dy, and y, and returns an object whose keys are the inputs (as defined by the inputs argument to executeKernel) and returns a function that returns the derivative with respect to that input. These derivatives should not call executeKernel, rather call math ops directly (this is so we can compute second order gradients).

Two example PRs adding gradients:
tensorflow/tfjs-core#521
tensorflow/tfjs-core#544

Note that we have lots of gradients in the Graph layer already, we just need to move them over to the gradients defined in eager mode.

Here is the list of ops and whether the gradient has been implemented:

• abs
• acos
• add
• argmax
• argmin (not important)
• asin
• atan
• avgPool
• batchNormalization
• cast
• ceil
• clip
• clone
• ceil
• concat
• conv1D
• conv2D
• conv2DDerBias (would be a second order der)
• conv2DDerFilter (would be a second order der)
• conv2DDerInput (would be a second order der)
• conv2DTranspose
• cos
• cosh
• depthwiseConv2D
• divide
• elu
• eluDer
• equal
• exp
• floor
• greater
• greaterEqual
• leakyRelu
• less
• lessEqual
• localResponseNormalization
• log
• logicalOr
• logicalAnd
• matMul (needs derivatives when using transposed bit)
• max
• maximum
• maxPool
• maxPoolBackprop (would be second order der)
• min
• minimum
• minPool
• multinomial
• multiply
• neg
• notEqual
• oneHot
• pad
• pow (half implemented, needs broadcast + derB)
• prelu
• preluDer (would be second order der)
• relu
• reshape
• resizeBilinear3D
• reverse
• selu
• sigmoid
• sin
• sinh
• slice
• softmax
• softmaxCrossEntropyWithLogits
• sqrt
• square
• step
• sub
• sum
• tan
• where
• tanh
• tile
• topK
• transpose

Copied from original issue: tensorflow/tfjs-core#561

[nsthorat]

From @dsmilkov on January 18, 2018 16:15

For those that are interested:

• leave a comment here to claim one or several ops
• make a PR that adds gradients for those ops
• After we merge your PR, we'll check the checkbox

I'd suggest to start with the unary ops.

[nsthorat]

From @israelvicars on January 19, 2018 17:20

I'll start abs. Thanks for tweeting the invitation to contribute.

[nsthorat]

From @manrajgrover on January 28, 2018 10:27

@nsthorat @dsmilkov A little guidance required here. What should the gradient of comparison and logical operations look like? Any good sources for the same?

[nsthorat]

From @gena on January 28, 2018 19:36

Added implementation for sigmoid, PR #603.

[nsthorat]

From @manrajgrover on February 9, 2018 13:57

@nsthorat We can tick off ceil, clip, cosh, floor, maximum, minimum, selu, sigmoid, sinh, softmax and tanh.

prelu, elu are partially implemented. Only gradients w.r.t. alpha need to be added.

leakyRelu and step gradients are in progress.

Since logical and comparison operations don't have an actual gradient, should we return NaN, zeros or not pass gradient function toexecuteKernel?

Also, looks like unary and binary operations are now almost covered. Would be needing pointers on next set to focus on.

[nsthorat]

From @dsmilkov on February 9, 2018 22:51

@manrajgrover Great q. For now, let's leave out logical and comparison ops and not pass a gradient function to executeKernel. Next would be reverse, slice, pad, concat (in that order). Thanks!

[nsthorat]

From @dsmilkov on February 19, 2018 3:30

I'll be taking reverse, slice, pad and concat (already started work). Thanks!

[nsthorat]

From @easadler on April 5, 2018 1:14

Seems like this may be a little out of date, but I'd like to help if you still would like help. I think oneHot is still available.

[nsthorat]

A high priority one would be batchNorm (which you'd return gradients for all of the parameters - this should be pretty straightforward) or resizeBilinear, if you're willing to take it on! resizeBilinear has a big filed here: #38

[nsthorat]

From @easadler on April 5, 2018 14:53

I will give batchNorm a shot!

[tafsiri]

I will work on resizeBilinear

[jgartman]

I can finish up pow.

[jgartman]

I'll finish matMul as well.

[jgartman]

I'll try localResponseNormalization

[nsthorat]

Closing this out in favor of individual issues.

[generic-github-user]

Is there an updated list of which functions have gradients implemented and which do not?