The automatic differentiation framework stores the actual random numbers generated during the forward pass and uses them directly during the backward pass. So you shouldn't have to do anything special and you can make use of automatic differentiation in your custom layer (by not defining your own backward function).
You can also read a bit more about automatic differentiation in MATLAB here: https://www.mathworks.com/help/deeplearning/ug/deep-learning-with-automatic-differentiation-in-matlab.html
There it says: "In other words, automatic differentiation evaluates derivatives at particular numeric values; it does not construct symbolic expressions for derivatives." Maybe this piece of information helps with understanding the behaviour when using random numbers.
I also put together a small example to illustrate what I mean. It is quite simplified from your example but hopefully you can use it to better understand or play around with autodiff framework and can transfer the idea to your implementation:
% Construct a simple network with some learnable layers and a custom layer
% in the middle which sets some channels of the data to zero.
layers = [ featureInputLayer(10)
fullyConnectedLayer(5,Name="fc1")
randomChannelDropLayer(5,"channelDrop")
fullyConnectedLayer(1,Name="fc2")];
net = dlnetwork(layers);
in = dlarray(rand(10,3),'CB');
% Now let's compute gradients. Note that the custom layer does not specify
% a backward function and hence automatic differentiation is used.
for i=1:5
disp("Execution #"+i)
% Every time we do a forward pass a different channel is dropped. The
% gradient of the custom layer's output with respect to its input
% contains zeros in the same channel that was randomly dropped during
% forward.
[layerOutput,layerGrad] = dlfeval(@customLayerGradients,net,in)
end
function [layerOutput,grad] = customLayerGradients(net,in)
% Compute gradients of the custom layer output with respect to its input.
% This gradient is used for backpropagation through the whole network.
[layerInput,layerOutput] = net.forward(in,Outputs=["fc1","channelDrop"]);
combinedOutput = sum(layerOutput,'all');
grad = dlgradient(combinedOutput,layerInput);
end
And here is the definition of the custom layer:
classdef randomChannelDropLayer < nnet.layer.Layer
% randomChannelDropLayer sets one randomly selected input channel to
% all zeros during training. The data passes through the layer
% unchanged during prediction.
properties
NumChannels
end
methods
function layer = randomChannelDropLayer(numChannels,name)
layer.NumChannels = numChannels;
layer.Name = name;
end
function Y = forward(layer,X)
channelToDrop = randi(layer.NumChannels,1);
Y = X;
Y(channelToDrop,:,:) = 0;
end
function Y = predict(~,X)
Y = X;
end
end
end
