Hi Zeyuyun!



 zeyuyun1:

we can add |X^T X - I| in our loss.
…
orthogonal_loss = reg * sym.abs().sum()



While this “orthogonality penalty” is zero if, and only, if X is orthogonal,
and is positive otherwise, it doesn’t work well for orthogonalizing X
with gradient descent, because its structure doesn’t fit* well with the
natural geometry of the set of orthogonal matrices.
Replacing the sum of absolute values with the sum of squares gives
a much more natural penalty that works well with gradient descent.
Please take a look at the modified version of your script, below, that
successfully applies the sum-of-squares orthogonality penalty to the
optimization of model.weight.


  param = model.weight.data



Also, as written, param is no longer part of the computation graph (It has
been set to .data.), so gradients of orthogonal_loss won’t backpropagate
and won’t affect the training of model.weight.
Finally, I substantially increase reg, the weight of the orthogonality penalty
in the total loss.


#make a random vector
X = torch.rand(30,500).to(device)
#make a random orthogonal matrix
rho = torch.nn.init.orthogonal_(torch.empty(500, 500)).to(device)

The optimal weight for the model is certainly rho, which will gives 0 loss. However, it doesn’t seem to converge to it.


The problem here is that with only 30 vectors to train on, you don’t have
enough information to reconstruct the 500x500 matrix rho.
There are a couple more issues:


#X_rotated will be the target vector
X_target = X@rho



Here you right-multiply X by rho.  Not a big deal, but it means that you
will be training model.weight to become equal to rho.T (rather than
rho).
Lastly, orthogonal matrices come in two disconnected sets:  those with
determinant +1 and those with determinant -1.  Although the optimizer
takes finite steps, these steps are “small,” so the optimization procedure
evolves model.weight continuously from its initial to its final value.
If you happen to choose rho with determinant +1, but also happen to
initialize model.weight with determinant -1, the optimizer will not** be
able to train model.weight to become rho.
*)  Crudely, you can think of the geometry of the set of orthogonal matrices
as being like a circle, with a rotational symmetry.   The sum of absolute
values singles out the x and y directions, not respecting that symmetry.
The sum of squares respects the rotational symmetry.  This is part of the
reason that the sum of squares works better as a penalty.
**)   Hypothetically, the optimizer could take a large, fortuitous step that
flips the sign of the determinant of model.weight, but this would be
highly unlikely.
Here is the modified script, with some informative comments:
import torch
torch.__version__

torch.random.manual_seed (2021)

nOrt = 500
nBatch = 30   # not enough rotated vectors to recover 500x500 rho

nEpoch = 1000
nFreq = 100

# make a batch of random vectors
X = torch.rand (nBatch, nOrt)

print ('X.shape =', X.shape)

# make a random orthogonal matrix with determinant +1
rho = torch.nn.init.orthogonal_ (torch.empty (nOrt, nOrt))
if  torch.slogdet (rho)[0] < 0.0:
    rho *= -1   # flip determinant -- assumes that rho is odd-dimensional

# the target vector will be the rotated X
# note that X is right-multiplied by rho
# this will cause model.weight to train to match rho.T, rather than rho itself
X_target = X @ rho

# the model is simply a single Linear layer (with bias)
model = torch.nn.Linear (nOrt, nOrt)
# initialize model.weight to another orthogonal matrix with determinant +1 to match det (rho)
wInit = torch.nn.init.orthogonal_ (torch.empty (nOrt, nOrt))
if  torch.slogdet (wInit)[0] < 0.0:
    wInit *= -1   # flip determinant -- assumes that wInit is odd-dimensional

with torch.no_grad():
    _ = model.weight.copy_ (wInit)

optimizer = torch.optim.Adam (model.parameters(), lr=0.001)
loss_fn = torch.nn.MSELoss (reduction='sum')

# reg = 0.001   # weight orthogonality-penalty much more heavily
reg = 1000.0
for t in range (nEpoch):
    # param = model.weight.data   # breaks computation graph --  model.weight not optimized
    param = model.weight
    sym = torch.mm (param, torch.t (param))
    sym -= torch.eye (param.shape[0])
    # ls_ort = sym.abs().sum()   # poor match to geometry of orthogonal matrices
    ls_ort = sym.pow (2.0).sum()
    ls_fit = loss_fn (model(X), X_target)
    ls_tot = ls_fit + reg * ls_ort
    with torch.no_grad():
        ldet = torch.slogdet (param)[1]
    optimizer.zero_grad()
    ls_tot.backward()
    optimizer.step()
    if t <= 10  or  t % nFreq == 0  or  t >= nEpoch - 1:
        print ('epoch:{:4d}'.format (t), ' ls_tot:{:9.3f}'.format (float (ls_tot)), ' ls_fit:{:9.3f}'.format (float (ls_fit)), ' ls_ort: {:.2e}'.format (float (ls_ort)), ' ldet: {: .1e}'.format (float (ldet)))

# check rho and model.weight determinants
print ('slogdet (rho)    =  {: 1.0f}  {: .3e}'.format (torch.slogdet (rho)[0], torch.slogdet(rho)[1]))
print ('slogdet (param)  =  {: 1.0f}  {: .3e}'.format (torch.slogdet (param)[0], torch.slogdet(param)[1]))

# check fit of model.weight to rho
with torch.no_grad():
    fit = (param - rho).abs().sum()
    fitT = (param - rho.T).abs().sum()

print ('fit = {:13.5f}, fitT = {:13.5f}'.format (fit, fitT))

And here is its output:
X.shape = torch.Size([30, 500])
epoch:   0  ls_tot: 9968.365  ls_fit: 9968.365  ls_ort: 2.27e-10  ldet:  6.9e-05
epoch:   1  ls_tot: 6984.131  ls_fit: 6536.737  ls_ort: 4.47e-01  ldet: -4.1e-01
epoch:   2  ls_tot: 5456.405  ls_fit: 4554.480  ls_ort: 9.02e-01  ldet: -9.2e-01
epoch:   3  ls_tot: 4708.714  ls_fit: 3511.573  ls_ort: 1.20e+00  ldet: -1.4e+00
epoch:   4  ls_tot: 4356.467  ls_fit: 3003.020  ls_ort: 1.35e+00  ldet: -1.8e+00
epoch:   5  ls_tot: 4166.570  ls_fit: 2764.880  ls_ort: 1.40e+00  ldet: -2.1e+00
epoch:   6  ls_tot: 4033.207  ls_fit: 2654.856  ls_ort: 1.38e+00  ldet: -2.1e+00
epoch:   7  ls_tot: 3915.477  ls_fit: 2601.749  ls_ort: 1.31e+00  ldet: -1.9e+00
epoch:   8  ls_tot: 3786.879  ls_fit: 2565.707  ls_ort: 1.22e+00  ldet: -1.7e+00
epoch:   9  ls_tot: 3633.864  ls_fit: 2528.301  ls_ort: 1.11e+00  ldet: -1.4e+00
epoch:  10  ls_tot: 3459.893  ls_fit: 2488.102  ls_ort: 9.72e-01  ldet: -1.2e+00
epoch: 100  ls_tot:  633.859  ls_fit:  630.421  ls_ort: 3.44e-03  ldet: -6.8e-02
epoch: 200  ls_tot:   47.161  ls_fit:   46.997  ls_ort: 1.64e-04  ldet: -6.4e-03
epoch: 300  ls_tot:    3.470  ls_fit:    3.461  ls_ort: 9.53e-06  ldet: -8.9e-04
epoch: 400  ls_tot:    0.318  ls_fit:    0.317  ls_ort: 7.55e-07  ldet: -2.2e-04
epoch: 500  ls_tot:    0.037  ls_fit:    0.036  ls_ort: 7.71e-07  ldet:  2.4e-05
epoch: 600  ls_tot:    0.014  ls_fit:    0.011  ls_ort: 3.07e-06  ldet: -2.9e-05
epoch: 700  ls_tot:    0.005  ls_fit:    0.004  ls_ort: 7.76e-07  ldet:  5.3e-05
epoch: 800  ls_tot:    0.007  ls_fit:    0.005  ls_ort: 1.90e-06  ldet: -6.9e-05
epoch: 900  ls_tot:    0.028  ls_fit:    0.016  ls_ort: 1.12e-05  ldet: -1.4e-04
epoch: 999  ls_tot:    0.021  ls_fit:    0.013  ls_ort: 8.38e-06  ldet: -5.8e-05
slogdet (rho)    =   1   5.507e-05
slogdet (param)  =   1  -4.864e-05
fit =   12627.76172, fitT =   12255.37109

This illustrates that both the “fitting” loss (MSELoss) and the orthogonality
penalty train down to low values.  Note, however, that, as expected,
model.weight does not train to match rho (nor rho.T) (fit and fitT,
respectively) because the 30 sample vectors are not enough to fully
reconstruct rho.
Best.
K. Frank