Maximizing the entropy would be forcing the random binary variable towards 0.5, i.e centering the ReLU input distributions median around 0.
Some very random ideas:

You essentially want to penalize whenever the Median of x deviates from 0. Maybe an easier ways to do this that propagates gradients of the correct direction directly towards whatever goes into the ReLU. A Sigmoid on the inputs with very high temperature H i.e if z = ReLU(x), p = Sigmoid(H*x) would tend steeply towards 0 or 1. Adding binary entropy on p to loss could be done manually and I imagine H controls the steepness of penalty of deviations from 0.5.

If you want to use Distributions to encapsulate this initialize a dist=Bernoulli(logits=H*x) in each forwardprop and add dist.entropy() to the loss function.


A second idea is to use the Straight-through Gumbel trick (similar to what I did with the ECDF but added stochastic noise). Someone else should answer exactly how this is done properly. I imagine expected value of solutions are closely related to solution 1)


A third idea could be to add Gumbel- or Normal noise to avoid saturated ReLUs I’m sure there’s lots of research on this.