How to train model in subspace of parameter space

username: Leon_Lang
cooked: <p>Hello,</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>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.<br>
params_0 are the initial D-dimensional parameters of the model.<br>
Originally, my d-dimensional parameter-vector params_d is the zero vector. It is what I try to learn.</p>
<p>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.<br>
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. <strong>params_d is the only vector I want to be updated in the optimizer step!</strong></p>
<p>I think my code doesn’t work, and I’m not quite sure why. params_d doesn’t receive any gradient.</p>
<p>Potential reasons:</p>
<ul>
<li>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.</li>
<li>Maybe the gradient is not able to pass through the “vector_to_parameters” step. If so, how could this be solved?</li>
</ul>
<p>Best regards,<br>
LL</p>

username: albanD
cooked: <p>Hi,</p>
<p>Could you show us the <code>vector_to_parameters</code> function please?</p>

username: Leon_Lang
cooked: <p>Hi albanD,</p>
<p>it is this one:</p>
<p><a href="https://pytorch.org/docs/stable/_modules/torch/nn/utils/convert_parameters.html" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/docs/stable/_modules/torch/nn/utils/convert_parameters.html</a></p>
<p>Here is the source code:</p>
<pre><code class="lang-auto">def vector_to_parameters(vec, parameters):
    r"""Convert one vector to the parameters

    Arguments:
        vec (Tensor): a single vector represents the parameters of a model.
        parameters (Iterable[Tensor]): an iterator of Tensors that are the
            parameters of a model.
    """
    # Ensure vec of type Tensor
    if not isinstance(vec, torch.Tensor):
        raise TypeError('expected torch.Tensor, but got: {}'
                        .format(torch.typename(vec)))
    # Flag for the device where the parameter is located
    param_device = None

    # Pointer for slicing the vector for each parameter
    pointer = 0
    for param in parameters:
        # Ensure the parameters are located in the same device
        param_device = _check_param_device(param, param_device)

        # The length of the parameter
        num_param = param.numel()
        # Slice the vector, reshape it, and replace the old data of the parameter
        param.data = vec[pointer:pointer + num_param].view_as(param).data

        # Increment the pointer
        pointer += num_param
</code></pre>

username: albanD
cooked: <p>As you can see, this function is using <code>.data</code> which means that it break the computational graph. This is not expected to work with differentiable operations.</p>
<p><code>nn.Parameters()</code> are built only to be trainable parameters, not intermediary Tensors that require gradients.<br>
One way to go around this is to manually delete the old parameters and just put your new weights there:</p>
<pre><code class="lang-auto">del mod.weight # Remove it from the Parameter list
mod.weight = your_new weights
# mod.weight is not just a Tensor, not a Parameter anymore.</code></pre>

username: KFrank
cooked: <p>Hi Leon!</p>
<aside class="quote no-group" data-username="Leon_Lang" data-post="1" data-topic="64132" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/leon_lang/48/18798_2.png" class="avatar"> Leon_Lang:</div>
<blockquote>
<p>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.</p>
</blockquote>
</aside>
<p>I don’t really follow your embedding scheme, and I can’t comment<br>
on your code.</p>
<p>But if I understand goal correctly, perhaps you can project your<br>
gradients onto your desired subspace after an unmodified<br>
<code>backward()</code> step.</p>
<p>The common optimizers use some variant of gradient descent<br>
where <code>params -= learning_rate * gradient</code> (where <code>gradient</code><br>
might have some momentum history in it, but this doesn’t change<br>
the idea).</p>
<p>So just do the same update step, but replace <code>gradient</code> with a<br>
<code>projected_gradient</code> that lies in your subspace.  (You would<br>
have tweak or write your own optimizer to do this.)</p>
<p>Alternatively, you could perform the standard update (and not have<br>
to modify the optimizer), and then project your new parameters to<br>
lie in your subspace.  (These two approaches are equivalent*, but I<br>
think in terms of projecting the gradient rather than the parameters<br>
for some reason.)</p>
<p>*)  “Project gradient” and “project (updated) parameters” are equivalent<br>
if your subspace is understood to contain the zero vector.  From your<br>
stated initial condition, I understand this to be the case.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Leon_Lang
cooked: <p>Dear albanD and KFrank,</p>
<p>thanks for your answers!</p>
<p><a class="mention" href="/u/kfrank">@KFrank</a>: The idea to change the optimizer (in my case SGD) worked in my case.</p>
<p>In order to do so, I split up the matrix E into several components, one for each array in the parameters of the model. Then, I could manually backpropagate the gradients of all the arrays through the embedding. I then re-embedded this into the full parameter space.<br>
Mathematically, that corresponds to multiplying the gradient with E*E.T, which is probably what you meant by “projection”.</p>
<p>In case you’re interested: I’m reproducing the results from this paper on the intrinsic dimension of objective landscapes: <a href="https://arxiv.org/abs/1804.08838" rel="nofollow noopener">https://arxiv.org/abs/1804.08838</a></p>
<p>Best regards,<br>
Leon</p>