Hello,
class Subspace_model(nn.Module):
    """
    Wraps a model in order to train it in a subspace
    """
    def __init__(self, model, E):
        super(Subspace_model, self).__init__()
        self.model = model
        self.register_buffer("E", E)
        self.d_dim = E.shape[1]
        self.params_d = nn.Parameter(torch.zeros(self.d_dim))
        self.register_buffer("params_0", parameters_to_vector(model.parameters()))

    def forward(self, x):
        params_D = self.E @ self.params_d + self.params_0
        vector_to_parameters(params_D, self.model.parameters())
        return self.model.forward(x)

I want to train a model in an affine-linear subspace of the full-parameter space. For that, I created an embedding E that takes vectors from R^d to R^D. Here, d is the dimensionality of the affine-linear subspace, D is the dimensionality of the full parameter space.
params_0 are the initial D-dimensional parameters of the model.
Originally, my d-dimensional parameter-vector params_d is the zero vector. It is what I try to learn.
A forward pass works like this: First, I do E @ params_d + params_0 in order to get the new parameters of the model. Then I apply it to the input.
For the backward pass, I want to do the backward of the model in order to get the gradients with respect to the D-dimensional weight vector. Then I want to backpropagate this further through the matrix E in order to update the d-dimensional Vector params_d. params_d is the only vector I want to be updated in the optimizer step!
I think my code doesn’t work, and I’m not quite sure why. params_d doesn’t receive any gradient.
Potential reasons:

Maybe I need to register self.model somehow as a buffer as well, since I don’t want it to be updated. However, pytorch doesn’t allow that.
Maybe the gradient is not able to pass through the “vector_to_parameters” step. If so, how could this be solved?

Best regards,
LL