Broken Adjoint Calculation in StackMap

[clutzweiler]

Hello GBI Team,

thanks for creating and maintaining this neat Matlab library!

I was facing some issues regarding the adjoint calculation in applyJacobianT_ of StackMap:

function x = applyJacobianT_(this, y, v)
    % Reimplemented from :class:`Map`   
    x = cell(this.numMaps,1);
    for n = 1:this.numMaps
        x{n} = this.alpha(n) .* this.mapsCell{n}.applyJacobianT(y,v);
    end
    x = cat(length(this.sizeout),x{:});
end

• Size related run-time failure.
• Functional issue: the forward pass at its core is broadcast -> per-channel-op-application -> concat; so the backward pass should be slice -> per-channel-op-transpose-application -> sum-reduction.

function x = applyJacobianT_(this, y, v)
    % Reimplemented from :class:`Map`
    
    % The adjoint op of a StackMap (:= apply & concat) is slice &
    % apply & summation.
    % Therefore we need to slice the upstream gradient along its
    % last dimension. To be working with arbitrary dimension, we
    % could e.g. use approaches like:
    %     https://stackoverflow.com/questions/19955653/matlab-last-dimension-access-on-ndimensions-matrix
    % For now, we take the simple path just flatten it to 2-D,
    % followed by another reshape op back to basic.

    % TODO(cl): quick'n'dirty, to be tested for robustness!
    
    full_dim = size(y);
    plain_dim = full_dim(1:end - 1);
    cat_dim = full_dim(end);
    
    y_flattened = reshape(y, [], cat_dim);
    
    x = cell(this.numMaps,1);

    for n = 1:this.numMaps
        y_n = reshape(y_flattened(:, n), plain_dim);
        x{n} = this.alpha(n) .* this.mapsCell{n}.applyJacobianT(y_n,v);
    end

    x = cat(length(this.sizeout),x{:});
    x = sum(x, length(this.sizeout));
end

What's your thoughts on a preferred solution? Let me know! Of course, this needs testing and cleaning-up...

Somewhat related question. Calculating the adjoint of StackMap * UpstreamOp, the adjoint of UpsteamOp gets executed multiple times, once for each channel of StackMap. Any better solution than wrapping StackMap functionality to first execute and sum all channel contributions before passing the gradient to the upstream node?

Best,
CL

[esoubies]

Hi CL,

Thanks a lot for having pointed out this bug. Clearly this StackMap has not been tested so far. Your proposition is perfect and you can do a pull-request for it. Actually, this is related to the issue #9 that is in the pipeline for a while but is still open...
Concerning you related question, I do not see why the adjoint of UpstreamOp gets executed multiple times as it applies to the output of the applyJacobian_T of StackMap which (with your code, but the previous version was wrong) perform the sum at the end. Did I miss something ?

Thx again,
Best
Emmanuel