Initialize weights using the matrix multiplication result from two nn.Parameter

username: YJHuang
cooked: <p>I have two tensor matrix, A $\in R^{nxm})$, and B  $\in R^{mx1}$</p>
<pre><code class="lang-auto">a = nn.Parameter(A, requires_grad=True)
b = nn.Parameter(B, requires_grad=True)
</code></pre>
<p>Is it possible to use the matrix multiplication result from A*B = C $\in R^{nx1}$ as the initialize weights for the nn.Linear layer?</p>
<pre><code class="lang-auto">linear_weights = nn.Parameter(torch.matmul(a,b), requires_grad=True)
linear_layer.weight = linear_weights
</code></pre>
<p>Another two questions to the gradient updates for parameters $a$ and $b$.</p>
<ol>
<li>Will the gradient updates the parameters $a$ and $b$ using above code?</li>
<li>If I change the required_grad=False to Parameter $a$ at the beginning, after that the $linear_weights$ has set to requires_grad=True. Will the Parameter $a$ still be frozen?</li>
</ol>

username: pascal_notsawo
cooked: <p>For your questions, just do some test to get the answer. Here is the code that supports the answers below.</p>
<pre><code class="lang-auto">import torch
class Model(torch.nn.Module) :
    def __init__(self) :
        super().__init__()
        torch.manual_seed(0)
        n, m = 1, 1
        A = torch.rand((n, m))
        B = torch.rand((m, 1))
        # AB ~ n x 1
        linear_layer = torch.nn.Linear(n, 1)

        flag = False
        if flag :
            # will raise this exception
            # TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)
            linear_layer.weight = torch.matmul(A,B)
        else :
            self.a = torch.nn.Parameter(A, requires_grad=True) # or A if you don't want to track the evolution of its value
            self.b = torch.nn.Parameter(B, requires_grad=True) # or B if you don't want to track the evolution of its value
            self.linear_weights = torch.nn.Parameter(torch.matmul(self.a,self.b), requires_grad=True)
            linear_layer.weight = self.linear_weights

        self.linear_layer = linear_layer
        
    def forward(self, x) :
        return self.linear_layer(x)  

def mae(y, y_pred) :
    """mean absolute error"""
    return (y - y_pred).abs().mean()


# data 
x = torch.tensor([1., 3, 3]).unsqueeze(dim=1) # unsqueeze to set batch_size to 3 ~~ [[1.], [3.], [3.]]
y = 2 * x
# model
model = Model()
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

model.state_dict()
"""
OrderedDict([('a', tensor([[0.4963]])),
             ('b', tensor([[0.7682]])),
             ('linear_weights', tensor([[0.3812]])),
             ('linear_layer.weight', tensor([[0.3812]])),
             ('linear_layer.bias', tensor([-0.7359]))])
"""

optimizer.zero_grad()
y_pred = model(x) # tensor([[-0.3547], [ 0.4078], [ 0.4078]], grad_fn=&lt;AddmmBackward&gt;)
loss = mae(y_pred, y) # tensor(4.5131, grad_fn=&lt;MeanBackward0&gt;)

loss.backward() 

optimizer.step()

model.state_dict()
"""
OrderedDict([('a', tensor([[0.4963]])), # no change
             ('b', tensor([[0.7682]])), # no change
             ('linear_weights', tensor([[0.3912]])), # change, because is the same variable as linear_layer.weight 
             ('linear_layer.weight', tensor([[0.3912]])), # change
             ('linear_layer.bias', tensor([-0.7259]))]) # change
"""
</code></pre>
<ul>
<li>
<p>Is it possible to use the matrix multiplication result from A*B as the initialize weights for the nn.Linear layer?<br>
No, <code>AxB</code> is still a tensor; so if you do it you will get an error like <code>TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)</code></p>
</li>
<li>
<p>Will the gradient updates the parameters <code>a</code> and <code>b</code> using above code?<br>
No, they are not involved in the computation graph.<br>
<code>linear_weights</code>, the result of their multiplication, is another full variable, which intervenes naturally in the graph through <code>linear_layer.weight</code>, with which it shares the same memory address (I prefer to say that it is the same variable as <code>linear_layer.weight</code>)<br>
If you want it to change, do this in <code>forward</code>:  <code>return torch.matmul(self.a,self.b)*x</code> (they intervene directly in the computation graph, unlike above)<br>
Then, after the forward and backward pass, you will have :</p>
</li>
</ul>
<pre><code class="lang-auto">model.state_dict()
"""
OrderedDict([('a', tensor([[0.5063]])), # change
             ('b', tensor([[0.7782]])), # change
             ('linear_weights', tensor([[0.3812]])), # no change
             ('linear_layer.weight', tensor([[0.3812]])), # no change
             ('linear_layer.bias', tensor([-0.7359]))]) # no change
"""
</code></pre>
<ul>
<li>If I change the <code>required_grad=False</code> to Parameter <code>a</code> at the beginning, after that the <code>linear_weights</code> has set to <code>requires_grad=True</code>. Will the Parameter <code>a</code> still be frozen?<br>
Yes (see the code below)</li>
</ul>
<pre><code class="lang-auto">n, m = 1, 1
A = torch.rand((n, m))
B = torch.rand((m, 1))
a = torch.nn.Parameter(A, requires_grad=False) # or A
b = torch.nn.Parameter(B, requires_grad=True) # or B
linear_weights = torch.nn.Parameter(torch.matmul(a,b), requires_grad=True)

a
"""
Parameter containing:
tensor([[0.3074]])
"""

linear_weights
"""
Parameter containing:
tensor([[0.1949]], requires_grad=True)
"""
</code></pre>

username: YJHuang
cooked: <p>What if I want to assign to Conv2d layer’ weights instead  nn.Linear layer? Is it still possible to update the a and b?</p>
<p>As you said, in linear layer, we can apply torch.matmul in forward(),<br>
How about using the conv2d operator with initialize weight by torch.matmul(a,b), can i also do something similar?</p>

username: pascal_notsawo
cooked: <p>If you want a parameter to be updated, make sure that it is of type nn.Parameter, that its attribute <code>requires_grad</code> is equal to True and that it is part of the computation graph.</p>

username: YJHuang
cooked: <p>Thank you <a class="mention" href="/u/pascal_notsawo">@pascal_notsawo</a>,</p>
<p>Is it possible for me to initialize the weight of a 1x1 pointwise nn.Conv2d kernel by using the torch.matmul(a,b) result? I want the nn.Parameter(a), nn.Parameter(b) be update too.</p>
<p>So after the loss.backward() and optimizer.step(), I expect the torch.matmul(a’, b’), where a’ and b’ are the updated nn.Parameters, will have the same value as the updated 1x1 pointwise nn.Conv2d.weight.</p>

username: pascal_notsawo
cooked: <p>You are welcome <a class="mention" href="/u/yjhuang">@YJHuang</a></p>
<p>It gets a little complicated.<br>
If you look at the <code>forward</code> method of <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/conv.html#Conv2d" rel="noopener nofollow ugc"><code>nn.Conv2d</code></a>, you will notice this:</p>
<pre><code class="lang-auto">return self._conv_forward(input, self.weight, self.bias)
</code></pre>
<p>So try to inherit the <code>nn.Conv2d</code> class and modify the <code>forward</code> method by replacing <code>self.weight</code> with <code>torch.matmul(self.a,self.b)</code> and observe the parameter values as I did above.</p>
<p>I dare to believe that it will be ok, since going further down than that we will run into C++: the basic convolution operations of pytorch are implemented in C++ (<a href="https://github.com/pytorch/pytorch/blob/1465970a343e61f2f2b104859ca7f5d7e03f5d02/aten/src/ATen/native/Convolution.cpp#L539" class="inline-onebox" rel="noopener nofollow ugc">pytorch/Convolution.cpp at 1465970a343e61f2f2b104859ca7f5d7e03f5d02 · pytorch/pytorch · GitHub</a>). So one of the only alternatives if it doesn’t work is to do everything from scratch in python (forward method), by googling you can find stable implementations.</p>

username: YJHuang
cooked: <p>The example you provided is very easy to follow and to observe the parameter values. I will also check these links and try to inherit the <code>nn.Conv2d</code> class.<br>
Thank you!!</p>

username: YJHuang
cooked: <p>It seems that the <code>torch.matmul()</code> will return a tensor value. So after I perform <code>torch.matmul(nn.Parameter(a), nn.Parameter(b))</code> and assign its value to <code>nn.Conv2d.weight</code>, it makes a copy. Is there any method to perform the operator like <code>torch.matmul</code> but return a <code>nn.Parameter</code> as well?</p>

username: pascal_notsawo
cooked: <p>In my case it works.<br>
Look at the following code, mainly the values of <code>a</code> and <code>b</code> before and after the forward-backward pass + optimizer gradient step (I have directly extracted some code snippets from one of my text classification problems)</p>
<pre><code class="lang-auto">import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn.common_types import _size_2_t

class Conv2d(torch.nn.Conv2d) :
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: _size_2_t,
        stride: _size_2_t = 1,
        padding: _size_2_t = 0,
        dilation: _size_2_t = 1,
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = 'zeros' 
    ) :
        super(Conv2d, self).__init__(
                    in_channels = in_channels,
                    out_channels = out_channels,
                    kernel_size = kernel_size,
                    stride = stride,
                    padding = padding,
                    dilation = dilation,
                    groups = groups,
                    bias = bias,
                    padding_mode = padding_mode
        )
        # self.weight.shape ~ n_filters x in_channels x filter_size x emb_dim
        filter_size, emb_dim = kernel_size
        assert self.weight.shape == torch.Size([n_filters, in_channels, filter_size, emb_dim]) # n_filters x in_channels x filter_size x emb_dim
        
        # Here I just want to make sure that AxB has the same dimensions as self.weight (It's up to you to make sure that it is the same on your side)
        torch.manual_seed(0) 
        intermediate_dim = 7
        A = torch.rand((n_filters, in_channels, filter_size, intermediate_dim))
        B = torch.rand((intermediate_dim, emb_dim))
        # AxB ~ n_filters x in_channels x filter_size x emb_dim)
        self.a = torch.nn.Parameter(A, requires_grad=True) 
        self.b = torch.nn.Parameter(B, requires_grad=True)
        # self.weight = torch.matmul(self.a, self.b)
    def forward(self, input: Tensor) -&gt; Tensor:
        #return self._conv_forward(input, self.weight, self.bias)
        return self._conv_forward(input, torch.matmul(self.a, self.b), self.bias) 


## Models
in_channels = 1
n_filters = 2
emb_dim = 7
filter_size = 3
model = Conv2d(in_channels = in_channels, out_channels = n_filters, kernel_size = (filter_size, emb_dim)) 

n_labels = 3
pred_layer = torch.nn.Linear(n_filters, n_labels)

## Data
torch.manual_seed(0) 
bs, slen = 5, 6
x = torch.rand((bs, slen, emb_dim))
y = torch.empty(bs, dtype=torch.long).random_(n_labels)

# optimizer
#optimizer = torch.optim.Adam(list(model.parameters()) + list(pred_layer.parameters()), lr=1e-2)
# The update of the parameters of the classification layer is not of interest to us here
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) 

model.state_dict()
"""
OrderedDict([('weight',
              tensor([[[[-0.0464,  0.0466, -0.1422, -0.0112,  0.1562, -0.0224,  0.0061],
                        [-0.0188,  0.0442,  0.1388,  0.2067,  0.1386,  0.2072, -0.0158],
                        [-0.1960, -0.1035,  0.1486, -0.0014, -0.1085, -0.1672, -0.2042]]],
              
              
                      [[[-0.1842, -0.0443,  0.1197,  0.1180, -0.2105,  0.1361, -0.1708],
                        [-0.0461, -0.0885, -0.0420, -0.0428, -0.1958, -0.1884, -0.0341],
                        [ 0.0028, -0.0991,  0.0822, -0.1964, -0.0147,  0.1919, -0.0890]]]])),
             ('bias', tensor([0.1971, 0.0790])),
             ('a',
              tensor([[[[0.4963, 0.7682, 0.0885, 0.1320, 0.3074, 0.6341, 0.4901],
                        [0.8964, 0.4556, 0.6323, 0.3489, 0.4017, 0.0223, 0.1689],
                        [0.2939, 0.5185, 0.6977, 0.8000, 0.1610, 0.2823, 0.6816]]],
              
              
                      [[[0.9152, 0.3971, 0.8742, 0.4194, 0.5529, 0.9527, 0.0362],
                        [0.1852, 0.3734, 0.3051, 0.9320, 0.1759, 0.2698, 0.1507],
                        [0.0317, 0.2081, 0.9298, 0.7231, 0.7423, 0.5263, 0.2437]]]])),
             ('b',
              tensor([[0.5846, 0.0332, 0.1387, 0.2422, 0.8155, 0.7932, 0.2783],
                      [0.4820, 0.8198, 0.9971, 0.6984, 0.5675, 0.8352, 0.2056],
                      [0.5932, 0.1123, 0.1535, 0.2417, 0.7262, 0.7011, 0.2038],
                      [0.6511, 0.7745, 0.4369, 0.5191, 0.6159, 0.8102, 0.9801],
                      [0.1147, 0.3168, 0.6965, 0.9143, 0.9351, 0.9412, 0.5995],
                      [0.0652, 0.5460, 0.1872, 0.0340, 0.9442, 0.8802, 0.0012],
                      [0.5936, 0.4158, 0.4177, 0.2711, 0.6923, 0.2038, 0.6833]]))])
"""

## Zero all the gradients
optimizer.zero_grad()


## Forward pass
x = x.unsqueeze(dim=1) # bs x 1 x slen x emb_dim
conved = model(x) # bs x n_filters x https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html ?
conved = F.relu(conved).squeeze(3) # bs x n_filters x (slen - emb_dim - filter_size + 1) ? https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
pooled = F.max_pool1d(conved, conved.shape[2]).squeeze(2) # bs x n_filters
pooled = F.dropout(pooled, p = 0.1) # bs x n_filters
y_pred = pred_layer(pooled)

## Loss and backward pass
loss = F.cross_entropy(input=y_pred, target=y) # tensor(13.4064, grad_fn=&lt;NllLossBackward&gt;)
loss.backward() 


## optimizer gradient step
optimizer.step()

model.state_dict()

"""
OrderedDict([('weight',
              tensor([[[[-0.0464,  0.0466, -0.1422, -0.0112,  0.1562, -0.0224,  0.0061],
                        [-0.0188,  0.0442,  0.1388,  0.2067,  0.1386,  0.2072, -0.0158],
                        [-0.1960, -0.1035,  0.1486, -0.0014, -0.1085, -0.1672, -0.2042]]],
              
              
                      [[[-0.1842, -0.0443,  0.1197,  0.1180, -0.2105,  0.1361, -0.1708],
                        [-0.0461, -0.0885, -0.0420, -0.0428, -0.1958, -0.1884, -0.0341],
                        [ 0.0028, -0.0991,  0.0822, -0.1964, -0.0147,  0.1919, -0.0890]]]])),
             ('bias', tensor([0.1871, 0.0690])),
             ('a',
              tensor([[[[0.4863, 0.7582, 0.0785, 0.1220, 0.2974, 0.6241, 0.4801],
                        [0.8864, 0.4456, 0.6223, 0.3389, 0.3917, 0.0123, 0.1589],
                        [0.2839, 0.5085, 0.6877, 0.7900, 0.1510, 0.2723, 0.6716]]],
              
              
                      [[[0.9052, 0.3871, 0.8642, 0.4094, 0.5429, 0.9427, 0.0262],
                        [0.1752, 0.3634, 0.2951, 0.9220, 0.1659, 0.2598, 0.1407],
                        [0.0217, 0.1981, 0.9198, 0.7131, 0.7323, 0.5163, 0.2337]]]])),
             ('b',
              tensor([[ 0.5746,  0.0232,  0.1287,  0.2322,  0.8055,  0.7832,  0.2683],
                      [ 0.4720,  0.8098,  0.9871,  0.6884,  0.5575,  0.8252,  0.1956],
                      [ 0.5832,  0.1023,  0.1435,  0.2317,  0.7162,  0.6911,  0.1938],
                      [ 0.6411,  0.7645,  0.4269,  0.5091,  0.6059,  0.8002,  0.9701],
                      [ 0.1047,  0.3068,  0.6865,  0.9043,  0.9251,  0.9312,  0.5895],
                      [ 0.0552,  0.5360,  0.1772,  0.0240,  0.9342,  0.8702, -0.0088],
                      [ 0.5836,  0.4058,  0.4077,  0.2611,  0.6823,  0.1938,  0.6733]]))])
"""
</code></pre>
<p>a and b have indeed changed.</p>

username: YJHuang
cooked: <aside class="quote no-group" data-username="pascal_notsawo" data-post="9" data-topic="120557">
<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/pascal_notsawo/48/22829_2.png" class="avatar"> pascal_notsawo:</div>
<blockquote>
<pre><code class="lang-auto">def forward(self, input: Tensor) -&gt; Tensor:
        #return self._conv_forward(input, self.weight, self.bias)
        return self._conv_forward(input, torch.matmul(self.a, self.b), self.bias) 
</code></pre>
</blockquote>
</aside>
<p>Thanks, it works.<br>
I realize that my My forward() function was defined as:</p>
<pre><code class="lang-auto">def forward(self, input: Tensor) -&gt; Tensor:
        self.weight = torch.matmul(self.a, self.b)
        return self._conv_forward(input, self.weight, self.bias) 
        #return self._conv_forward(input, torch.matmul(self.a, self.b), self.bias) 
</code></pre>
<p>and it will return an error message:</p>
<pre><code class="lang-auto">TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)
</code></pre>
<p>I don’t know what is the difference between the two forward functions. Only yours work! I really appreciate your help!!</p>

username: pascal_notsawo
cooked: <p>I explained it above when answering your second question in my first post.<br>
<code>self.weight = torch.matmul(self.a, self.b)</code> creates a whole other variable, and it is this variable that is involved in the computation graph, not <code>a</code> and <code>b</code>. But then if you want to know how <code>a</code> and <code>b</code> would be updated, you can calculate it yourself after the backward :</p>
<pre><code class="lang-auto">a.grad = dloss/da = dloss/dweigths * dweigths/da = weigths.grad * f(b), with f(b) = b or b^T if I am not mistaken
a = optimizer(a, learning_rate, a.grad), for example a = a - learning_rate*a.grad 
</code></pre>
<p>With <code>self.weight = torch.matmul(self.a, self.b)</code>, it’s a bit like saying to the computation graph: hello dear graph, don’t worry about <code>a</code> and <code>b</code> when retropagating the gradient, they are now represented by <code>self.weight</code> (self). This fails because at this moment <code>torch.matmul(self.a, self.b)</code> is only a normal tensor, hence :</p>
<pre><code class="lang-auto">TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)
</code></pre>
<p>But you can do as below in the <code>forward</code> function so that the error disappears, but <code>a</code> and <code>b</code> won’t change anymore, but rather <code>self.weight</code> (I remind again and again that it is because it intervenes in the computation graph, and that <code>a</code> and <code>b</code> do not intervene anymore)</p>
<pre><code class="lang-auto">self.weight = torch.nn.Parameter(torch.matmul(self.a, self.b), requires_grad=True)
return self._conv_forward(input, self.weight, self.bias)
</code></pre>
<p>Do it as below (without <code>self</code>) and you will see that <code>a</code> and <code>b</code> will change :</p>
<pre><code class="lang-auto">weight = torch.matmul(self.a, self.b)
return self._conv_forward(input, weight, self.bias) # equivalent to return self._conv_forward(input, torch.matmul(self.a, self.b), self.bias)
</code></pre>
<p>Here it works because <code>weight</code> (without <code>self</code>) is just a bait, it’s just a temporary variable nothing more, the graph doesn’t even know it (if I’m not mistaken), but rather <code>a</code> and <code>b</code>. It doesn’t prevent <code>a</code> and <code>b</code> from intervening in the graph.</p>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/innocent.png?v=9" title=":innocent:" class="emoji" alt=":innocent:"> <img src="https://discuss.pytorch.org/images/emoji/apple/innocent.png?v=9" title=":innocent:" class="emoji" alt=":innocent:">This famous computation graph …</p>