When doing SGD, we can split up the update rule to each parameter tensor in a for loop. Hence, for every layer in a feedforward neural network, we would update weights and biases.
However, I want to implement a different version of SGD where only the k parameters corresponding to the largest gradient entries are actually updated by SGD. The crucial thing here: These should be the k maximal gradients over the entire model parameters. If k=1 then this would mean we only want to update the parameter in the network corresponding to the largest gradient entry, and not to the largest entry in each layer.
How would I find the maximal elements of the entire gradient vector efficiently?
Thanks a lot!