How to set the diagonal entries of Linear layer weight matrix to be always negative?

bala_nbk: <p>I want to create a neural network with a single linear layer whose weight matrix has diagonal entries that will always be negative (even during training). I tried various means, but nothing seems to work. The closest I have got to is the below code:</p>
<pre><code class="lang-auto">class IdentityMask(nn.Module):
    def __init__(self, n):
        super(IdentityMask, self).__init__()
        self.n = n
        self.weight = nn.Parameter(torch.ones(n, n), requires_grad=False)
        torch.diagonal(self.weight).fill_(-1.0)

class LinearLayer(nn.Module):
    def __init__(self, n):
        super(LinearLayer, self).__init__()
        self.identity_mask = IdentityMask(n)
        self.linear = nn.Linear(n, n, bias=False)
        nn.init.normal_(self.linear.weight, mean=0.0, std=0.01)

    def forward(self, x):
        self.linear.weight.data *= self.identity_mask.weight
        x = self.linear(x)
        return x
</code></pre>
<p>What should I do to get the diagonal entries of <code>self.linear.weight</code> to be negative?</p>

bala_nbk: <p>OK. I figured out what is the problem. The line <code>self.linear.weight.data *= self.identity_mask.weight</code> multiplies the diagonal entries by -1 for every forward pass. But what we want is diagonal entry multiplied by -1 only when the entry turns positive. So I have modified the code as follows:</p>
<pre><code class="lang-auto">class LinearLayer(nn.Module):
    def __init__(self, n):
        super(LinearLayer, self).__init__()
        self.linear = nn.Linear(n, n, bias=False)
        nn.init.normal_(self.linear.weight, mean=0.0, std=0.01)

    def forward(self, x):
        diagonal = torch.diag(self.linear.weight)
        for i in range(diagonal.size(0)):
            if diagonal[i] &gt; 0:
                diagonal[i] *= -1.0
                self.linear.weight.data[i][i] = diagonal[i]
        print("Diagonal entries:", diagonal)
        x = self.linear(x)
        return x
</code></pre>
<p>No need for a mask layer. This code works. But if there are any other better ways to achieve my objective, please let me know.</p>

albanD: <p>note that a more efficient implementation of this can be done:</p>
<pre><code class="lang-python">import torch

a = torch.randn(10, 10)

print(a.diagonal())
a.diagonal().mul_(a.diagonal().sign())
print(a.diagonal())
</code></pre>
<p>Also you can use <code>with torch.no_grad()</code> instead of <code>.data</code>.</p>
<p>Finally this looks like a good candidate for a Parametrization (<a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrize.register_parametrization.html?highlight=parametrization#torch.nn.utils.parametrize.register_parametrization" class="inline-onebox">torch.nn.utils.parametrize.register_parametrization — PyTorch 2.0 documentation</a>) if you want to be able to do this without having to rewrite a Module by hand.</p>

KFrank: <p>Hi Bala (and Alban)!</p>
<aside class="quote no-group" data-username="bala_nbk" data-post="1" data-topic="177339" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/71c47a/48.png" class="avatar"> bala_nbk:</div>
<blockquote>
<p>What should I do to get the diagonal entries of <code>self.linear.weight</code> to be negative?</p>
</blockquote>
</aside>
<p>My intuition is that it is better to smoothly map an unconstrained trainable<br>
parameter (that runs over <code>(-inf, inf)</code>) to a new tensor whose diagonal<br>
is negative (and runs over <code>(-inf, 0.0)</code>, rather than brute-force flip the<br>
sign of the diagonal.  It is straightforward and conceptually satisfying to<br>
train the unconstrained parameter and understand the negative-diagonal<br>
as an intermediate result.</p>
<p>Suppose during training your optimizer moves a slightly negative diagonal<br>
entry to a slightly positive value.  You then flip it back to negative,  But on<br>
the next iteration, the optimizer moves it back to a positive value.  While<br>
conceptually acceptable, it just seems to me that this is likely to throw a<br>
little bit of sand in the optimization process (and possibly confuse fancier<br>
optimizers such as <code>Adam</code>),</p>
<p><code>-exp()</code> is a well-behaved function maps to strictly negative values.</p>
<p>Here is an illustration:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; preWeight = torch.randn (5, 5, requires_grad = True)   # unconstrained trainable parameter
&gt;&gt;&gt; preWeight                                              # unconstrained diagonal -- can be positive
tensor([[ 0.4305, -0.3499,  0.4749,  0.9041, -0.7021],
        [ 1.5963,  0.4228, -0.6940,  0.9672, -0.5319],
        [ 0.8088, -0.1603,  0.8184, -0.6093,  0.8177],
        [ 0.1459, -0.9558, -1.3761,  1.3246, -0.0744],
        [ 0.5472,  1.6779,  0.8275, -1.0542, -0.7374]], requires_grad=True)
&gt;&gt;&gt; weight = preWeight.clone()
&gt;&gt;&gt; weight.diagonal().copy_ (-preWeight.diagonal().exp())
tensor([-1.5380, -1.5262, -2.2668, -3.7607, -0.4784],
       grad_fn=&lt;AsStridedBackward0&gt;)
&gt;&gt;&gt; weight                                                 # derived weight tensor with negative diagonal
tensor([[-1.5380, -0.3499,  0.4749,  0.9041, -0.7021],
        [ 1.5963, -1.5262, -0.6940,  0.9672, -0.5319],
        [ 0.8088, -0.1603, -2.2668, -0.6093,  0.8177],
        [ 0.1459, -0.9558, -1.3761, -3.7607, -0.0744],
        [ 0.5472,  1.6779,  0.8275, -1.0542, -0.4784]], grad_fn=&lt;CopySlices&gt;)
&gt;&gt;&gt; x = torch.randn (5, 5)
&gt;&gt;&gt; (weight @ x).sum().backward()
&gt;&gt;&gt; preWeight.grad                                         # gradients flow back to trainable parameter
tensor([[-1.2142,  0.9742,  3.3650, -1.4189,  2.5436],
        [ 0.7895, -1.4869,  3.3650, -1.4189,  2.5436],
        [ 0.7895,  0.9742, -7.6279, -1.4189,  2.5436],
        [ 0.7895,  0.9742,  3.3650,  5.3361,  2.5436],
        [ 0.7895,  0.9742,  3.3650, -1.4189, -1.2168]])
</code></pre>
<p>(Also, as Alban notes, you could register such a mapping as a<br>
parameterization.)</p>
<p>Best.</p>
<p>K. Frank</p>

bala_nbk: <p>Hi,</p>
<p>I am revisiting this problem after a long time. I am essentially trying to follow the idea of using a smooth map ( -exp() ) applied to the diagonal elements, suggested by KFrank. But now I am getting a different error that seems to do with autograd and back propagation. The error that I get is: <strong>Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</strong></p>
<p>Here is the new code:</p>
<pre><code>class NegDiagLinear(nn.Module):
    def __init__(self, in_features, out_features):
       super(NegDiagLinear, self).__init__()
       self.in_features = in_features
       self.out_features = out_features
       self.pre_weight = nn.Parameter(torch.Tensor(out_features, in_features))
       nn.init.normal_(self.pre_weight, mean=0, std=0.01)
       self.weight = self.pre_weight.clone()
    
    def forward(self, input):
       self.weight.diagonal().copy_ (-self.pre_weight.diagonal().exp())
       return input @ self.weight

class simple_model(nn.Module):
    def __init__(self):
       super().__init__()
       self.linear = NegDiagLinear(2,2)

    def forward(self, x):
       return self.linear(x)

model = simple_model()

define_criterion = torch.nn.MSELoss()

SGD_optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

for epoch in range(18):
   SGD_optimizer.zero_grad()
   predict_y = model(x) 
   loss = define_criterion(predict_y, y)  
   loss.backward() 
   SGD_optimizer.step()
   print('epoch {}, loss function {}'.format(epoch, loss.item()))
</code></pre>
<p>My initial guess is that this is due to clone() applied to self.pre_weight(). So I dabbled around this trying to see if I could get around this problem. But I can’t seem to get rid. Please help me to solve this problem. If you could also provide an explanation of the cause of this error, it will be nice.</p>

KFrank: <p>Hi Bala!</p>
<aside class="quote no-group quote-modified" data-username="bala_nbk" data-post="5" data-topic="177339" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/71c47a/48.png" class="avatar"> bala_nbk:</div>
<blockquote>
<p><strong>Trying to backward through the graph a second time</strong></p>
<p>Here is the new code:</p>
<pre><code>class NegDiagLinear(nn.Module):
    def __init__(self, in_features, out_features):
       super(NegDiagLinear, self).__init__()
       self.in_features = in_features
       self.out_features = out_features
       self.pre_weight = nn.Parameter(torch.Tensor(out_features, in_features))
       nn.init.normal_(self.pre_weight, mean=0, std=0.01)
       self.weight = self.pre_weight.clone()
    
    def forward(self, input):
       self.weight.diagonal().copy_ (-self.pre_weight.diagonal().exp())
       return input @ self.weight
</code></pre>
</blockquote>
</aside>
<p>Your problem is that:</p>
<pre data-code-wrap="python"><code class="lang-python">      self.weight = self.pre_weight.clone()
</code></pre>
<p>creates the parts of the computation graph that connects the off-diagonal<br>
elements of <code>weight</code> to <code>pre_weight</code> only once (in <code>__init__()</code>).</p>
<p>Then:</p>
<pre data-code-wrap="python"><code class="lang-python">       self.weight.diagonal().copy_ (-self.pre_weight.diagonal().exp())
</code></pre>
<p>creates the computation graph for the diagonal elements every time you call<br>
forward.</p>
<p>But your first call to <code>.backward()</code> frees the whole computation graph, including<br>
that for the off-diagonal elements, and the off-diagonal piece is never rebuilt.  So<br>
your second call to <code>.backward()</code> raises the “backward through the graph a second<br>
time” error.</p>
<p>Try:</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self, input):
       weight = self.pre_weight.clone()
       weight.diagonal().copy_ (-self.pre_weight.diagonal().exp())
       return input @ weight
</code></pre>
<p><code>forward()</code> now rebuilds the whole computation graph – both the off-diagonal and<br>
diagonal parts – every time it is called.</p>
<p>Note that <code>weight</code> is no longer a property of <code>NegDiagLinear</code> – it’s now just a<br>
local variable in <code>forward()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>