Enable MXNet to compute second (or higher) order gradients for a list of operators using autograd package.
Table of Contents |
---|
Link to dev List discussion
https://lists.apache.org/thread.html/464712f0136fb51916ca9f1b702b99847e108dbdbd0b6a2b73fc91f1@%3Cdevfc9088b6c6345abe3338e6f411f94a274a9e6f01ecf3351d328be34b@%3Cdev.mxnet.apache.org%3E
Developer
Feature Shepherd
Problem
Currently only a very limited number of operators (such as exp) that support second or higher order gradient calculation. For other operators, if users try to get the second order gradient of an operator, MXNet would issue an error message such as "mxnet.base.MXNetError: [23:15:34] src/pass/gradient.cc:192: Operator _backward_XXX is non-differentiable because it didn't register FGradient attribute." This is because MXNet backend does not implement the FGradient function for the backward node of the operator and therefore cannot support second (and higher) order of gradient. See the example of sin operator below. The gradient function is registered through a hidden operator _backward_sin. However, the _backward_sin operator does not register FGradient function.
Higher order gradient calculation is required for many applications such as adaptive learning rate optimization [1], WWGAN-GAN GP network [2], network neural architecture search [3], and etc. Implementing higher order gradient can help unlock these applications and improve the usability and popularity of Apache MXNet framework.
...
This project will implement a set of operators to support adaptive learning rate optimization proposed by [1] . The following operators will be supported in this project. We will also keep an up-to-date document of the supported operators in MXNet repo.
Operator | Second order support? | ETA |
---|---|---|
exp | Y | |
elemwise_mul | N | 04/2019 |
sin | N | 04/2019 |
cos | N | 04/2019 |
relu | N | 04/2019 |
dot | N | 04/2019 |
negative | N | 04/2019 |
softmax | N | 05/2019 |
sigmoid | N | 05/2019 |
...
Case I: Operators with simple FGradient function: many simple tensor arithmetic operators are implemented using simple backward functions. E.g. the sin(x) operator shown below
To support higher order gradient, we will replace the “_backward_sin” operator with an explicit cos(x) operator. Because cos(x) operator will have its own FGradient function, by calling the FGradient functino function recursively we can achieve many order of derivative of the sin(x) operator. A preliminary implementation as proof-of-concept is here.
Case II: Operators with complex FGradient function: in some other operators the derivative is not so straightforward or can be easily represented analytically. E.g. the softmax operator. The FGradient function is registered using a class:
For this type of operators, we have two options.
- Replace the hidden “_backward_xxx” with a simple function with other forward operators just like we did with sin(x) operator. This needs efficient representation of the analytical derivative function for each operator and requires many a lot of expertise on a per operator base.
- Register a FGradient function inside the “_backward_xxx” operator. This can be simpler in terms of implementation and less destructive to current code design pattern. However, it limits the gradient to second order and cannot automatically extend to higher order of gradients.
Case III: Operators implemented using third party library such as cuDNN and MKLDNN: many neural network specific operators are using implementations directly from cuDNN and/or MKLDNN libraries if available, such as batchnorm, convolution, pooling etc. An example of convolution operator using MKLDNN library is shown below:
In such case cases, we will have to depend on the third party library to provide higher order gradient support for these operators if users install MXNet library with other third party packages. For the first phase of this project, since we are dealing with multilayer perceptron network, we do not have to implement higher order derivative for these operators. We will survey the related third party packages and incorporate their higher order gradient implementation in MXNet.
...
There should be no performance impact. We will track the performance through nightly Benchmark AI performance test.
Test Plan
Add unit tests for each operator that have the second order gradient support
...
1) Use finite-difference approximation to calculate second order gradients. The pro is its generic for almost most all operators and developers do not need to implement second order gradient calculation one by one. The con is picking the finite-difference value becomes another hyperparameter to tune and requires expertise from users.
2) Implement a code generation mechanism which will make it easier to create the differentiable backward operators. Currently Pytorch is using a code generation mechanism in which backward pass is implemented with Variables and defined in a yaml file, thus fully differentiable. We need to think big and provide a scalable implementation that allows for easy extensibility and external contributions. Moreover, also if implementation of a differentiable backward pass is easy to code and review, it will multiply developer productivity and maintainability.
...
[1] Forward and Reverse Gradient-Based Hyperparameter Optimization
[2] Improved Training of Wasserstein GANGANs
[3] Neural Architecture Search with Reinforcement Learning
...