[RFC][WIP] Tensor Expression level automatic differentiation

[sgrechanik-h]

I'm working on automatic differentiation at the level of compute expressions, and I would like to share some progress and hear any comments. Currently the automatic differentiation works well enough for some operations, so that it is possible to train a simple model, here is a tutorial on how to do this. Yet, for many operations the performance is unacceptable, but I'm working on it.

My implementation mostly follows this paper. In this notebook I describe how my implementation works internally and give a list of operations which are known to work or not to work. Basically, the AD consists of two parts:

• The automatic differentiation itself which simply differentiates expressions according to the well-known rules and produces inefficient expressions. The code is here.
• A set of transformations to optimize the resulting inefficient expressions. The code is here.

All transformations work on the level of compute expressions (before scheduling). Their general goal is to eliminate summation over zeros by moving up conditional expressions of the form cond ? val : 0 and then using them to simplify iteration domains of reductions. Hopefully, these transformations may be useful for some other tasks besides AD when they are powerful enough. Currently the main problem is that they don't understand modular arithmetic (which is needed for differentiating dilated and strided convolutions and for the flattening operation).

The git branch
The squashed commit
The tutorial on training a simple model
The notebook describing some internals

[tqchen]

This looks great. I would also encourage us to think a bit more on how this can play together with the current system. Specifically, while being able to run AD in tensor expression level end to end is fun, we may not necessarily use this method, as AD on high-level IR still offers advantages of things like general graph optimizations and pick of algorithms(winograd vs direct), so technically we would still encourage the use of high-level IR AD.

On the other hand, this would become practical if we could have a proposal on how to write quick expression op, that can be integrated together with other high-level IR and its AD system, this would enable a powerful to write custom operators and automatically generate gradient for it

[jroesch]

@MarisaKirisame after PLDI could you provide more feedback.

[MarisaKirisame]

@jroesch yeah

[sgrechanik-h]

@tqchen A couple of months ago I tried to integrate this autodiff into NNVM, the attempt (now abandoned) is still preserved in some old commits:
https://github.com/sgrechanik-h/tvm/blob/109d2f785d3e73e56665fbc987f6bd0dc5823d60/nnvm/src/top/tensor/gradient.cc
The idea was to add a special operation called gradient which accepted as attributes the original operation name and the original operation's attributes. This operation's compute function consisted in automatically differentiating the original operation's compute. But I suppose there must be a better approach.

Being able to easily define new high-level ops just from tensor expressions sounds like a good idea. I definitely missed something like this in NNVM.

Also I think it would be nice to be able to override operations' attributes, like gradients, from python code. At least, this would be great for comparing automatic gradients with manually written gradients.

[tqchen]

FYI, we are moving toward NNVMv2(relay) and this is a good time to think about related questions as scaffolding are being done. Relay already support override op's attributes from python

[sergei-mironov]

On the other hand, this would become practical if we could have a proposal on how to write quick expression op, that can be integrated together with other high-level IR and its AD system, this would enable a powerful to write custom operators and automatically generate gradient for it

Indeed, I think that higher level IRs will always be able to decide which method to use. As far as Relay uses TVM under the hood, it should be in good position to make a choice: either substitute operations during building the backward path (e.g. Winograd -> matmul), switch to hand-written gradients or use a semantically-correct default implementation which is provided by the current approach.

[sergei-mironov]

Continuing discussion about training, started in #2122

@junrushao1994 @were @sgrechanik-h @szha I suggest we continue discussion related to training here.

So how could we do manual (or automatic) scheduling on a automatically generated backprop?

We plan to do the following regarding scheduling:

1. Provide some automatic scheduling out of the box
2. Provide API for manual scheduling of gradients. We could allow programmers to access gradient nodes by their source nodes.
3. Provide API to override node's automatic gradient with manually-defined gradient if it is required.

[were]

For 1., 2.
As far as I imagine, the API will look like:

diff = tvm.diff([op1, op2, op3])
sch = tvm.create_schedule(diff.op)

The differentiation is also some kind of op node so that we can unify the API to TVM general flow.
I am also working on some new op node scheduling thing. If this part requires some help, I am willing to contribute to this issue.

[sgrechanik-h]

I've updated our automatic differentiation branch. Now the result of differentiating flatten is acceptable, and operations like max pool work better as well. We have also improved the API, the simple use-case look pretty much the same up to function renaming:

[dw1, dw2, dw3] = tvm.differentiate(loss, [w1, w2, w3])

(The function differentiate is defined here, and here is a small tutorial).
However, now it is possible to get individual adjoints from the result:

x = tvm.placeholder((32, 3, 28, 28), name='x')
w1 = tvm.placeholder((10, 3, 3, 3), name='w1')
t1 = topi.nn.conv2d(x, w1, 1, 0, 1)
t2 = topi.nn.flatten(t1)
t3 = topi.sum(t2)

res = tvm.differentiate(t3)
res.adjoints[t1]

(It may be useful for manually scheduling intermediate tensors).
And it is also possible to override Jacobian computation for some tensors:

def mydiff(out, inp, head):
    return tvm.compute(inp.shape, lambda ax0, ax1, ax2, ax3: head[ax0, ax3 + ax2*26 + ax1*676])

res = tvm.differentiate(t3, [x, w1], manual={(t2, t1): mydiff})

(Which may be useful when autodiff does poor job).

[tqchen]

@sgrechanik-h Thanks for the hard work on proposing and bringing in the implementation, this seems to be an RFC that many folks are interested in.

I have no doubt on the possibility to make a correct and efficient tensor expression auto diff. The key issue we need to address first is how can we make a clear set of APIs and namespaces so that:

• We have a clear separation between relay's grad and the current tensor expression one
• Have a clear indication for frontend user to use relay's grad, while use tensor expression one for the operator impl
• Have a consistent API so a user does not get confused when using one or another.
• Because the primary support for end to end grad is in relay, we want to name the files clearly(e.g. tvm/autodiff-> tvm/expr_grad.h) to make that distinction, so it is more clear that relay is the entry point of tvm stack's differentiation.

These discussions and decision are important because they will affect how users will use the API, as per https://docs.tvm.ai/contribute/code_review.html#deliberate-on-api-and-data-structures

I suggest we discuss them, reach consensus before we move on

cc @merrymercy @jroesch @yzhliu @junrushao1994 @ZihengJiang @grwlf @dmlc/tvm-committer

[tqchen]

@sgrechanik-h would be great if you can edit the first post and list the choices or options clearly so that everyone can discuss in a coordinated way, see #2535