Replace a parameter with other variable while keeping the backpropagation

Dronem_Ebidanko: <p>I ran into the same problem as this person, <a href="https://github.com/pytorch/pytorch/issues/40165" class="inline-onebox" rel="noopener nofollow ugc">How to replace a parameter with other variable while keeping the backpropagation? · Issue #40165 · pytorch/pytorch · GitHub</a>, the person was referred to this forum but I did not find the same post.</p>
<p>I don’t want to change the model structure. For every parameter c I want it to be c = a*b, and I want the backprop to be able to go back to a and b, not stopping at c.</p>
<p>I know some laborious (and perhaps inefficient) ways to do this, for example re-writing the model architecture, or manually calculating a or b in each backprop, but is there a nicer way to do this?</p>
<p>Any suggestion is appreciated and welcomed.</p>

J_Johnson: <p>You can just use the <code>nn.Parameter</code> class found <a href="https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html" rel="noopener nofollow ugc">here</a>.</p>
<p>For example:</p>
<pre><code class="lang-auto">class CustomLinearLayer(nn.Module):
    def __init__(self, in_channels, out_channels, bias=True, device = None):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        self.A=nn.Parameter(torch.zeros((self.in_channels,self.out_channels), device = device))
        self.B=nn.Parameter(torch.zeros((self.in_channels, self.out_channels), device = device))

        if bias:
            self.bias = nn.Parameter(torch.zeros((self.out_channels), device = device))
        else:
            self.bias = None
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / (self.A.size(0))**(1/4)
        for param in self.parameters():
            nn.init.uniform_(param, -stdv, stdv)
        if self.bias is not None:
            stdv = 1. / (self.A.size(0))**(1/2)
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x):
        C = self.A*self.B #elementwise multiplication
        if self.bias is not None:
            return x@C+self.bias
        else:
            return x@C

model = CustomLinearLayer(10, 20)

dummy_inputs = torch.rand(100, 10)

print(model(dummy_inputs).size())
</code></pre>
<p>Note the weight initialization method, which is provided, will give you fairly decent results for what you’re trying to do(FYI, improper initialization may make the model produce nans or show little to no improvements).</p>

Dronem_Ebidanko: <p>Thanks for the reply.<br>
This is essentially what I did, I went into the huggingface transformer code and wrote a class wrapper of a module to do pretty much what you did. I guess this is the best solution for now. It is burdensome if the model architecture is not original (i.e., if the model architecture is from a big pile of codes, then doing this modification takes a while). What I was hoping was something like:</p>
<pre><code class="lang-auto">for param in model.parameters(): 
    paramA = nn.Parameter(..., require_grad = False)
    paramB = nn.Parameter()
    param = paramA * paramB

</code></pre>
<p>and this keeps the original model architecture but replaces each leaf parameter with two a multiplication. Specifically in this application, we have a predefined param A and want to train for paramB.</p>
<p>Thanks again for the suggestion and the reply.</p>

J_Johnson: <p>LoRA might be similar to what you’re doing, except that’s elementwise addition of the outer product of two vectors, instead of multiplication of a tensor. You can see here for coding ideas:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py">
  <header class="source">

      <a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">microsoft/LoRA/blob/main/loralib/layers.py</a></h4>


      <pre><code class="lang-py">#  ------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  ------------------------------------------------------------------------------------------
import torch
import torch.nn as nn
import torch.nn.functional as F

import math
from typing import Optional, List

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = r
</code></pre>



  This file has been truncated. <a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Additionally, lucidrains also has a PyTorch implementation of PaLM with RLHF via using LoRA weights, that you might find interesting to consider, as that is implemented with transformers:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py">
  <header class="source">

      <a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">microsoft/LoRA/blob/main/loralib/layers.py</a></h4>


      <pre><code class="lang-py">#  ------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  ------------------------------------------------------------------------------------------
import torch
import torch.nn as nn
import torch.nn.functional as F

import math
from typing import Optional, List

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = r
</code></pre>



  This file has been truncated. <a href="https://github.com/microsoft/LoRA/blob/main/loralib/layers.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>