How to add orthogonal constrain to weight?

zhangpuhan: <p>I’m trying to add orthogonal constrain to the weight of a Linear layer. I find this is similar to the function torch.nn.utils.weight_norm, so I write a similar class WeightOrtho to orthogonalize the weight. Each time the forward pre hook assigns a new orthogonal weight to the layer. But this seems not working. The weight is not updating in the training. Is there a good way to add orthogonal constrain to weight?</p>
<pre><code class="lang-auto">weight_ortho(torch.nn.Linear(6, 2, bias=True), name='weight', dim=0)
</code></pre>
<pre><code class="lang-auto">class WeightOrtho(object):
    def __init__(self, name, dim):
        if dim is None:
            dim = -1
        self.name = name
        self.dim = dim

    def compute_weight(self, module):
        w = getattr(module, self.name)
        return torch.nn.Parameter(torch.transpose(torch.qr(torch.transpose(w, 0, 1))[0], 0, 1))

    @staticmethod
    def apply(module, name, dim):
        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, WeightOrtho) and hook.name == name:
                raise RuntimeError("Cannot register two weight_ortho hooks on "
                                   "the same parameter {}".format(name))

        if dim is None:
            dim = -1

        fn = WeightOrthho(name, dim)

        weight = getattr(module, name)

        # remove w from parameter list
        del module._parameters[name]

        module.register_parameter(name , Parameter(weight.data))
        setattr(module, name, fn.compute_weight(module))

        # recompute weight before every forward()
        module.register_forward_pre_hook(fn)

        return fn

    def remove(self, module):
        weight = self.compute_weight(module)
        delattr(module, self.name)
        del module._parameters[self.name]
        module.register_parameter(self.name, Parameter(weight.data))

    def __call__(self, module, inputs):
        setattr(module, self.name, self.compute_weight(module))

def weight_ortho(module, name='weight', dim=0):
    WeightOrtho.apply(module, name, dim)
    return module
</code></pre>

Bick95: <p>Have a look at this <a href="https://github.com/facebookresearch/MUSE/blob/3159355b93f5c3c4883808ba785ba9d18d7f5e81/src/trainer.py#L181" rel="nofollow noopener">link</a>.</p>
<p>It’s an iterative orthogonalization procedure which you have to call iteratively until an acted upon linear layer converges to orthogonality. If you are wondering about their implementation, <code>self.params.map_beta</code> is just a scalar hyperparameter which defaults to <a href="https://github.com/facebookresearch/MUSE/blob/master/unsupervised.py#L41" rel="nofollow noopener">0.001</a>.</p>
<p>A bit simplified, their procedure looks as follows:</p>
<pre><code class="lang-auto">beta = 0.001

def orthogonalize(self):
        """
        Orthogonalize the mapping.
        """
        W = self.linear_layer.weight.data
        W.copy_((1 + beta) * W - beta * W.mm(W.transpose(0, 1).mm(W)))
</code></pre>

Lezcano: <p>A bit late to the party, but starting in 1.10 this feature is supported natively in PyTorch:<br>
<a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html</a></p>
<p>Leaving this here in case someone else bumps into this post in the future.</p>