Wrong MaxPool gradient

[yiyuezhuo]

I tried to reused PyTorch weights in Flux and made the output fully consistent. But unfortunately, gradient results seem wrong. Finally I figured out that it's due to "Strided" MaxPool's wrong gradient. Following gists can be used to show it:

https://gist.github.com/yiyuezhuo/94a08637a6a80ac4a25eb89e5e605e16

compare_master.ipynb gist show present implementation comparison results.

By just replace it = with +=. The most inconsistency are removed.

But some inconsistency still hold. It's due to PyTorch's row major and Flux/Julia's column major difference. They follow their natural order to find the first element to apply gradients. So replace the iteration order can remove almost all inconsistency:

But this modification may damage performance, and just get vain gain.

Present unit test failed to detect this error, since reshape(1:n, ...) style example is unable to check the case in which an entry is maximum on many windows. The = override other gradients passed to it though.

I have prepared two branches to Pull Request:

fix_maxpool_gradient
fix_maxpool_gradient_further (change iteration order)

Which one should I PR? (Well, since this issue is ignored entirely, maybe no PR will be accepted. But the gradient is still wrong, a solution will be fine imo.)

[CarloLucibello]

Hi,
I have a couple of questions to understand the issue:

• Is the maxpool mathematical operation correctly implemented in NNlib for any combination of its input arguments?
• Then the gradient implementation should give the corresponding mathematical gradient. I'm quite surprised we don't test the gradient implementations in NNlib against numerical gradient.
• So you are saying that the gradient implementation is wrong whenever there is receptive field overlap, right? Then, you should submit the PR implementing the correct gradient, we want to be correct, performance considerations come after. Possibly test the implementation against numerical gradient.

[yiyuezhuo]

Set beta=1 instead of default 0 may get the correct result. Since

beta = 0  # default value
dx[x_idxs...] += dy_idx*alpha + beta*dx[x_idxs...]

and

beta = 1
dx[x_idxs...] = dy_idx*alpha + beta*dx[x_idxs...]  # old implementation

are equivalent.

According to exported interface, beta can be defined by kwargs. But Flux's call doesn't pass any kwargs. If set beta=1 is an expected behavior, then the error should be attributed to Flux instead of NNlib.

But according to treatment on the meanpool , it seems like that the beta should be modified in this function as well rather than requiring the caller to do something. So "completed" code like:

        if $(name == :max)
            alpha = T(1)
        end

should fix it as well.

I decided to replace = with += instead of adding a block like above or fixing the Flux end. Because I notice that testing for "strided"(overlapped) doesn't pass beta and the comment saying:

                    # Maxpooling does not support `beta` right now.  :(
                    #else
                    #    dx[x_idxs...] = T(0) + beta*dx[x_idxs...]

So I leave beta as default value 0. But I still wonder if maybe fixing beta from somewhere is a better solution.

Yes, present unit tests use just inputs like [1,2,3,4,5] and predefined answers. Interestingly, I just made the exact same misleading mistake before checking unit test code. It's such a tempting but misleading input.

Yes. Present unit test try to test overlap case but [1,2,3,4,5] style input cover up the problem. I will add testing for numerical gradients and mentioned cases.

[yiyuezhuo]

Fixed a order mistake in gist. Now further version is not required:
https://gist.github.com/yiyuezhuo/ade7f497e983a1aba56533c1b1056214

master (present) scatter:

master (present) difference heatmap:

present scatter:

proposed difference heatmap:

[CarloLucibello]

This is an example reproducing the issue:

julia> using Zygote, NNlib, FiniteDiff

julia> x = rand(10,10,1,1);

julia> g1 = gradient(x -> sum(maxpool(x,(2,2), pad=0, stride=2)), x)[1];

julia> g2 = FiniteDiff.finite_difference_gradient(x -> sum(maxpool(x,(2,2), pad=0, stride=2)), x);

julia> g1 ≈ g2
true

julia> g1 = gradient(x -> sum(maxpool(x,(2,2), pad=0, stride=1)), x)[1];

julia> g2 = FiniteDiff.finite_difference_gradient(x -> sum(maxpool(x,(2,2), pad=0, stride=1)), x);

julia> g1 ≈ g2
False

julia> g1
[:, :, 1, 1] =
 0.0  0.0  0.0  0.0  1.0  0.0  0.0  1.0  0.0  0.0
 0.0  1.0  0.0  1.0  0.0  1.0  0.0  1.0  0.0  1.0
 0.0  0.0  1.0  0.0  1.0  0.0  0.0  0.0  1.0  0.0
 0.0  1.0  0.0  0.0  0.0  1.0  0.0  1.0  0.0  1.0
 1.0  0.0  0.0  1.0  0.0  1.0  0.0  0.0  0.0  1.0
 0.0  0.0  1.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0
 0.0  1.0  0.0  0.0  1.0  1.0  0.0  1.0  0.0  1.0
 1.0  0.0  1.0  1.0  1.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  1.0  1.0  0.0  1.0
 0.0  1.0  1.0  1.0  0.0  1.0  1.0  0.0  1.0  0.0

julia> g2
10×10×1×1 Array{Float64,4}:
[:, :, 1, 1] =
 0.0  0.0  0.0  0.0  2.0  0.0  0.0  2.0  0.0  0.0
 0.0  3.0  0.0  1.0  0.0  2.0  0.0  1.0  0.0  2.0
 0.0  0.0  3.0  0.0  4.0  0.0  0.0  0.0  2.0  0.0
 0.0  4.0  0.0  0.0  0.0  3.0  0.0  4.0  0.0  1.0
 1.0  0.0  0.0  3.0  0.0  2.0  0.0  0.0  0.0  1.0
 0.0  0.0  4.0  0.0  0.0  0.0  0.0  4.0  0.0  0.0
 0.0  1.0  0.0  0.0  1.0  4.0  0.0  2.0  0.0  2.0
 2.0  0.0  2.0  4.0  1.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  1.0  2.0  0.0  2.0
 0.0  2.0  1.0  1.0  0.0  2.0  1.0  0.0  1.0  0.0

@yiyuezhuo how is your fix doing on the above example? could you submit a PR? it's easier to discuss a PR than a gist

[yiyuezhuo]

@CarloLucibello The fix can pass this test. A PR has been submitted.