Add orthogonal constraint optimizer

username: Akusa
cooked: <p>Say you want to optimize parameters of your model where your parameters are a tensor n x D.</p>
<pre><code class="lang-auto">var1 = torch.tensor(..., requires_grad = True) # Tensor of shape n x D
optimizer = torch.optim.Adam([var1], lr = 1e-5)
for ite in range(ite_max):
     optimizer.zero_grad()
     loss = loss_function(var1)
    optimizer.step()
</code></pre>
<p>I want to add an orthogonal constraint on my parameter var1, that is,</p>
<pre><code class="lang-auto">var1[:,i].T * var1[:,j] = kronecker_ij
</code></pre>
<p>where kronecker_ij = 1 if i = j, 0 otherwise.</p>
<p>How should I do it?</p>

username: KFrank
cooked: <p>Hi Akusa!</p>
<aside class="quote no-group quote-modified" data-username="Akusa" data-post="1" data-topic="121515" 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/akusa/48/41989_2.png" class="avatar"> Akusa:</div>
<blockquote>
<p>I want to add a orthogonal constraint on my parameter var1</p>
</blockquote>
</aside>
<p>You can add an orthogonal penalty to your loss function to push <code>var1</code><br>
towards being orthogonal.</p>
<p>Note, doing so will not force <code>var1</code> to be exactly orthogonal – if the rest<br>
of your loss function prefers a non-orthogonal <code>var1</code>, the optimization<br>
procedure will make a trade-off between getting the lowest value for your<br>
regular loss and having <code>var1</code> be orthogonal.  You can tune this trade-off<br>
by multiplying the penalty term with a weighting factor.  As you increase<br>
the weight of the penalty, the closer to orthogonal <code>var1</code> will become.</p>
<p>This script adds an orthogonality penalty to your code fragment:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.random.manual_seed (2021)

def ortho_penalty (t):
    return ((t @t.T - torch.eye (t.shape[0]))**2).sum()

n = 5
D = 7
lr = 0.1
ite_max = 200

var1 = torch.randn ((n, D), requires_grad = True) # Tensor of shape n x D
optimizer = torch.optim.Adam([var1], lr = lr)
for ite in range(ite_max):
    optimizer.zero_grad()
    loss = ortho_penalty (var1)
    if ite == ite_max - 1  or  ite % 20 == 0:
        print ('loss =', loss.item())
    loss.backward()
    optimizer.step()

print ('var1 = ...\n', var1)
print ('var1 @ var1.T = ...\n', var1 @ var1.T)
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">1.7.1
loss = 272.65423583984375
loss = 0.8204055428504944
loss = 0.2746223211288452
loss = 0.055152103304862976
loss = 0.0036679445765912533
loss = 0.0007598181255161762
loss = 7.743481546640396e-05
loss = 1.1900866411451716e-05
loss = 4.0019048697104154e-07
loss = 1.5665852970414562e-07
loss = 1.9611301027566697e-08
var1 = ...
 tensor([[-0.5975, -0.3254,  0.2185, -0.0053, -0.3276, -0.4471,  0.4268],
        [-0.1584, -0.7288, -0.2472, -0.2933, -0.1118,  0.2660, -0.4619],
        [-0.4144,  0.5079,  0.3184, -0.3310, -0.2745,  0.0373, -0.5317],
        [ 0.5968, -0.0300, -0.0399, -0.1651, -0.6691, -0.3961, -0.0974],
        [-0.0128,  0.0696,  0.0430,  0.0669, -0.5315,  0.7458,  0.3873]],
       requires_grad=True)
var1 @ var1.T = ...
 tensor([[ 1.0000e+00, -2.4647e-05, -1.3560e-05,  3.6400e-05,  3.8207e-05],
        [-2.4647e-05,  1.0000e+00, -7.5102e-06, -9.2573e-06, -1.5244e-05],
        [-1.3545e-05, -7.5102e-06,  1.0000e+00,  5.3346e-06, -4.1515e-05],
        [ 3.6404e-05, -9.2536e-06,  5.3346e-06,  1.0000e+00,  1.0524e-05],
        [ 3.8207e-05, -1.5244e-05, -4.1515e-05,  1.0554e-05,  9.9999e-01]],
       grad_fn=&lt;MmBackward&gt;)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Akusa
cooked: <p>Thank you! So, the easiest way is to add a penalty term on my loss function such as:<br>
<code>loss = loss_function(var1) + w*ortho_penalty(var1)</code><br>
where <code>w &gt; 0 </code> is the penalty weight.<br>
In this case, should not be “better” to use a Lagrangian multiplier to enforce the constraint?</p>