How to share weights with multple encoders

username: Ahmad_Pouramini
cooked: <p>I have multiple LSTM encoders like this:</p>
<pre><code class="lang-auto">class PromptEncoder(torch.nn.Module):
    def __init__(self,name, length,embedding_dim,id_offset, init_embs, prompt_ids,**kwargs) -&gt; None:
        super().__init__()
        self.length = length
        self.name = name
        self.prompt_ids = prompt_ids
        self.input_ids = torch.nn.parameter.Parameter(torch.tensor(prompt_ids),
             requires_grad=False)
        self.embedding_dim = embedding_dim
        self.id_offset = id_offset
        self.embedding = torch.nn.Embedding(length,embedding_dim)
        self.net_inps = torch.nn.parameter.Parameter(torch.arange(length),
            requires_grad=False)
        self.lstm = torch.nn.LSTM(
            input_size=embedding_dim,
            hidden_size=embedding_dim //2, #my code
            num_layers=2,
            dropout=0,
            bidirectional=True,
            batch_first=True
        )
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(embedding_dim, embedding_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(embedding_dim, embedding_dim)
        )
    def forward(self,prompt_token_ids,pids=None):
        # create embedding vectors for input ids
        embeds = self.embedding(self.net_inps)
        # do forward calculations
        x = self.lstm(embeds.unsqueeze(0))
        emblog.info("lstml embeds: %s",embeds)

        running_weight = self.mlp(x[0]).squeeze(0)

        prompt_token_ids = (prompt_token_ids.view(-1,1) == self.input_ids).int().argmax(dim=1)
        # return weights for prompt_token_ids 
        return F.embedding(prompt_token_ids,running_weight)


</code></pre>
<p>some of <code>prompt_ids</code> or <code>input_ids</code> could be shared among them. However each encoder has its own <code>embedding</code> matrix. They are all part of a container Module and are learned together. I want the shared ids point to a shared embedding so that if one changes, the change reflects to the embedding of all.<br>
This is <code>forward</code> wrapper:</p>
<pre><code class="lang-auto">    def forward(self,input_ids, labels, decoder_input_ids=None,pids=None,**kwargs):
      prompt_masks = self.prompt_token_fn(input_ids)
        if prompt_masks.any():
            input_ids_ = input_ids.clone()
            if self.replacing_token_id is not None:
                # replace prompt ids in input_ids with replacing token
                input_ids_[prompt_masks]=self.replacing_token_id
            # find the model embeddings of input ids except for prompt tokens
            inputs_embeds = self.model_embeddings(input_ids_)
            device=inputs_embeds.device
            for encoder in self.prompt_encoders:
                #encoder = self.prompt_encoders[0]
                prompt_token_fn = encoder.get_prompt_token_fn()
                encoder_masks = prompt_token_fn(input_ids)
                if encoder_masks.any():
                    #find input ids for prompt tokens
                    prompt_input_ids = input_ids[encoder_masks]
                    # call forwards on prompt encoder whose outputs are prompt embeddings
                    prompt_embeds = encoder(prompt_input_ids,\
                        pids).to(device)
                    # replace prompt_embeddings calculated by prompt encoder in input embeddings
                    # in input embeds replace embeddings for prompt token with output of encoder
                    inputs_embeds[encoder_masks]=prompt_embeds
        else:
            inputs_embeds = self.model_embeddings(input_ids)
        
            return self.underlying_model(inputs_embeds=inputs_embeds,**kwargs)

</code></pre>
<pre><code class="lang-auto">

How can I implement such architecture?</code></pre>

username: anantguptadbl
cooked: <p>If the encoders are trained separately why would you want the changes to one encoder to get reflected in another<br>
If you want that to happen, you have to train all of them as part of the same model.</p>
<p>You can use <a href="https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html" rel="noopener nofollow ugc">ModuleList</a> and <a href="https://pytorch.org/docs/stable/generated/torch.nn.ParameterList.html" rel="noopener nofollow ugc">ParameterList</a> to hold the models and the embeddings respectively</p>

username: Ahmad_Pouramini
cooked: <p><a class="mention" href="/u/andreasceid">@andreasceid</a> Thank you very much! The encoder are in a ModuleList. I put more of my code in the question including how they are called in the forward of the container Module. The container module actually wrap a transformer model (T5) which is freezed and the result of forward pass on encoders are fed into it. I am someway beginner with Pytorch and Transformer. For example some parts of my LSTM enoders such as having both input_id parameters and net_ids might be redundant?</p>
<p>Suppose input ids for whole model are <code>[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, ....]</code> in batch. the ids of encoder 1 are <code>[3, 4]</code> and the ids for encoder 2 are <code>[1,3, 5]</code>. where <code>3</code> is common between them. Each encoder must update embedding for its corresponding input. However, <code>3</code> exists in both of them. Maybe I should merge the results for <code>3</code> in the loop in <code>forward</code>, I don’t know. I thought maybe they must refer to a shared embedding space. Please guide me for this problem.</p>

username: anantguptadbl
cooked: <p><a class="mention" href="/u/ahmad_pouramini">@Ahmad_Pouramini</a> this will not be an easy problem to solve. The best way forward would be to have no intersecting groups.</p>
<p>However, if you want to go ahead, you can use some sort of weight averaging</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/">
  <header class="source">
      <img src="https://pytorch.org/favicon.ico?" class="site-icon" width="32" height="32">

      <a href="https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>

  <article class="onebox-body">
    <img src="https://pytorch.org/assets/images/pytorch-logo.png" class="thumbnail onebox-avatar" width="500" height="500">

<h3><a href="https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/" target="_blank" rel="noopener nofollow ugc">PyTorch</a></h3>

  <p>An open source machine learning framework that accelerates the path from research prototyping to production deployment.</p>


  </article>

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

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


username: Ahmad_Pouramini
cooked: <p><a class="mention" href="/u/anantguptadbl">@anantguptadbl</a>  thanks, can’t I connect the common cells of two lstm to a feed forward and get one outupt and use that one?</p>