Automatic Differentiation for indexed data

[hulihan-start]

Recently I built the forward pass for my custom function, but when I faced the backward pass, some problems happened.

For example, I have the embedding data $e$ and the indices $a$ and $b$, the function is e[a] * e[b] (not real, for example). The gradient I want to obtain is the $e[a].grad$ and $e[b].grad$, but the user guide didn't mention this situation. It only provides the backward creation for the given data (not indexing data).

Is there any suggestion or guide for my question? Thank you so much!

[hulihan-start]

I mean, should I write two or more functions for my case? One part is for indexing data, and the other part is for computation.

[roastduck]

You can only get gradients on e instead of e[a] or e[b]. This is not because the computation of gradients: FreeTensor can computes only gradients of e[a] and e[b]. This is because FreeTensor's functions only returns full tensors. If not, how do you expect FreeTensor to return gradients of e[a] and e[b] to you?

Yes, you can write two functions and separate indexing and computation.

Although feasible, writing two functions may have some negative performance impact: The indexed value must be saved back to slow memory (global memory for GPU) to be passed between functions. One possible way to work around this problem is to first derive the computing function, and inline it to the indexing function via ft.inlined_invoke. However, this is currently non-trivial and involves some low-level FreeTensor details. If you want to have a try, I can provide more information.

[hulihan-start]

Thanks a lot! For large-scale data during training, the better way is to return some of the gradients (I think, because we only need a small part of the total data and this could reduce memory traffic). For example, if I have more than 10GB of data on GPU, for each iteration I only need to update 4MB or less data, it's hard and unnecessary to return a whole gradient of 10GB of data and I only need to return 4MB of them and their indices. Would it be better for me to write the functions separately? BTW, I'm not sure if this is suitable for automatic differentiation.

It's a little difficult for me to implement based on some low-level FreeTensor functions but thanks for your enthusiasm and kind reply.

[roastduck]

Yes, I think you should write separated functions in your case. For performance impact, I mean FreeTensor cannot optimize across functions. The more small functions you write, the less FreeTensor can optimize. If each of your functions only contains few operations, you'd better turn to PyTorch.

[hulihan-start]

Thank you so much!