Hi Victor,
I’m trying to implement a regulariser described in this paper: [2105.05612] Evading the Simplicity Bias: Training a Diverse Set of Models Discovers Solutions with Superior OOD Generalization.
The shared features are a function of the input and model parameters, so the gradient of head output w.r.t. shared features is a function of the model parameters as well.
Anyways, I was able to fix the issue. I was using ReLU activation in the MLP heads, and the second derivative of the loss w.r.t. the shared feature representation involves taking the double derivative of the ReLU function, which is always zero. Changing the activations to Softplus (which has non-zero second derivative) resolved this issue—the gradients now back-propagate to the base model parameters as expected.