Backprop & Mean Pooling in Rev-MViT

Hi! Thanks for open-sourcing the reversible ViT architecture. I am finding it very useful for contrastive learning--e.g SimCLR--where large batch sizes are needed.

I had 2 questions that I hoped you could help me with:
1. I'm a bit confused about the back-propagation mechanism.

My assumption is that the `dY_1` and `dY_2` that are returned from backward pass are the derivatives of the inputs of the reversible transformation with respect to the final loss. I.e,

$$
I =
\begin{bmatrix}
x_{1} \\
x_{2} \\
\end{bmatrix} \\ \xrightarrow[\text{T}]{}
\begin{bmatrix}
x_{1} + F(x_2) \\
x_{2} + \\
\end{bmatrix} \\ \xrightarrow[T_2]{}
\begin{bmatrix}
x_{1} + F(x_2) \\
x_{2} + G(x_1 + F(x_2)) \\
\end{bmatrix} \\ =
O
$$

the `backward_pass` [function](https://github.com/facebookresearch/SlowFast/blob/64f8d7e1f888c789f6e6ef881fc520cc88589cd0/slowfast/models/reversible_mvit.py#L547) gets the derivatives of vector $O$ with respect to the final loss (and vector $O$ itself) & returns the derivatives of vector $I$ with respect to the final loss (and vector $I$ itself).

I understand this line:
https://github.com/facebookresearch/SlowFast/blob/ad73dcfa43b56dfeec96c551128ec615a1080997/slowfast/models/reversible_mvit.py#L583

But, I do not understand this line:
https://github.com/facebookresearch/SlowFast/blob/ad73dcfa43b56dfeec96c551128ec615a1080997/slowfast/models/reversible_mvit.py#L607

Apologies for the imprecise terminology here, but it seems like $x_2$ has a non-linear relationship with $Y_2$ (i.e, incrementing $x_2$ by 1 does not increment $Y_2$ by 1) due to the $G(x_1 + F(x_2))$ term. Thus, I don't understand how we can just use the `dY_2` term. I (think) I understand the `X_2.grad` term, since the gradient must flow backwards through `Y_1` and then through $F(x_2)$, but the first term doesn't make sense to me.

Let me know if this question makes sense.

2. In your implementation, you have the option to [choose whether to use mean pooling](https://github.com/facebookresearch/SlowFast/blob/64f8d7e1f888c789f6e6ef881fc520cc88589cd0/slowfast/models/video_model_builder.py#L1176). However, it seems like the difference between the 2 choices is whether normalization happens before/after the fusion + mean pooling process. Given that mean pooling happens regardless of which option is chosen, I am wondering why the option is called "use_mean_pooling".

Thanks again for taking the time to answer these questions!