Answers.Body: <p>In the sequential forward</p>
  <pre><code>h_prev = torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device)
  </code></pre>
  <p>to</p>
  <pre><code>h_prev = self.g(torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device))
  </code></pre>
  <p>This results in</p>
  <pre><code>import torch 
  import torch.nn as nn 
  import torch.nn.functional as F
  
  
  class minGRU(nn.Module):
      def __init__(self,input_size:int,hidden_size:int):
          super().__init__()
          self.input_size = input_size 
          self.hidden_size = hidden_size 
  
          self.linear_h = nn.Linear(input_size,hidden_size)
          self.linear_z = nn.Linear(input_size,hidden_size)
      
      @staticmethod
      def g(x:torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          The continuous function g ensures that ht ← g(Lineardh (xt)) is positive
          &quot;&quot;&quot;
          return torch.where(x &gt;= 0, x+0.5, torch.sigmoid(x))
      
      @staticmethod
      def log_g(x:torch.Tensor) -&gt; torch.Tensor:
          return torch.where(x &gt;= 0, (F.relu(x)+0.5).log(),-F.softplus(-x))
  
      @staticmethod
      def parallel_scan_log(log_coeffs:torch.Tensor, log_values:torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          log_coeffs: [batch_size, seq_len, input_size]
          log_values: [batch_size, seq_len + 1, input_size]
  
          &quot;&quot;&quot;
          a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
          log_h0_plus_b_star = torch.logcumsumexp(log_values - a_star, dim=1)
          log_h = a_star + log_h0_plus_b_star
          return torch.exp(log_h)[:, 1:]
  
      def sequential_forward(self, x_t:torch.Tensor, h_prev:torch.Tensor=None) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          x_t     : [batch_size, input_size]
          h_prev  : [batch_size, hidden_size]
          &quot;&quot;&quot;
          if(h_prev is None):
              h_prev = self.g(torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device))
  
          z = torch.sigmoid(self.linear_z(x_t)) 
          h_tilde = self.g(self.linear_h(x_t)) 
          h_t = (1 - z) * h_prev + z * h_tilde
  
          return h_t
  
      def forward(self, x:torch.Tensor, h_0:torch.Tensor=None) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          x: [batch_size, seq_len, input_size]
          h_0: [batch_size, 1, hidden_size]
          &quot;&quot;&quot;
          if(h_0 is None):
              h_0 = torch.zeros((x.size(0),1,self.hidden_size),device=x.device)
  
  
          k = self.linear_z(x) 
          log_coeffs = -F.softplus(k)
          log_z = -F.softplus(-k)
          log_tilde_h = self.log_g(self.linear_h(x)) 
  
          log_h_0 = self.log_g(h_0)
          h = self.parallel_scan_log(log_coeffs,torch.cat([log_h_0, log_z + log_tilde_h], dim=1))
          return h #(batch_size,seq_size,hidden_size)
  
  
  
  
  if __name__ == &quot;__main__&quot;:
      batch_size = 3
      seq_size = 2
      input_size =5 
      x = torch.randn(batch_size,seq_size,input_size)
  
      model = minGRU(input_size,6)
      print(model(x))
  
      output_list = []
      ht = None
      for i in range(x.size(1)):
          ht = model.sequential_forward(x[:,i,:],ht)
          output_list.append(ht.unsqueeze(dim=1))
      
      print(torch.cat(output_list,dim=1))
  
  </code></pre>
  <p><b> Parallel output </b></p>
  <pre><code>tensor([[[0.6535, 0.4925, 0.3262, 0.1168, 0.1602, 0.1158],
           [1.0672, 0.7427, 0.3846, 0.1483, 0.2766, 0.1792]],
  
          [[0.6006, 0.1788, 0.2839, 0.1379, 0.1121, 0.1104],
           [0.4730, 0.4698, 0.3884, 0.2961, 0.2522, 0.2705]],
  
          [[0.2452, 0.1759, 0.1381, 0.1435, 0.9452, 1.0829],
           [0.3459, 0.3899, 0.3189, 1.0957, 0.5292, 0.8135]]],
         grad_fn=&lt;SliceBackward0&gt;)
  
  </code></pre>
  <p><b> Sequential output </b></p>
  <pre><code>tensor([[[0.6535, 0.4925, 0.3262, 0.1168, 0.1602, 0.1158],
           [1.0672, 0.7427, 0.3846, 0.1483, 0.2766, 0.1792]],
  
          [[0.6006, 0.1788, 0.2839, 0.1379, 0.1121, 0.1104],
           [0.4730, 0.4698, 0.3884, 0.2961, 0.2522, 0.2705]],
  
          [[0.2452, 0.1759, 0.1381, 0.1435, 0.9452, 1.0829],
           [0.3459, 0.3899, 0.3189, 1.0957, 0.5292, 0.8135]]],
         grad_fn=&lt;CatBackward0&gt;)
  <span class="math-container">```</span>
  </code></pre>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-10-15T14:03:43.827Z
Answers.LastActivityDate: 2024-10-15T14:12:45.820Z
Answers.LastEditDate: 2024-10-15T14:12:45.820Z
Answers.UserDisplayName: hal9000
Answers.Body: <p>While I can't directly identify the specific issue with your code, I might recommend checking out the</p>
  <p><a href="https://github.com/cheind/mingru" rel="nofollow noreferrer">mingru</a> repository I've recently created.</p>
  <p>The code addresses minor paper inconsistencies and ensures consistent behavior between sequential and log-space parallel modes (validated up to 3 stacked layers with 1-e5 precision). The repository includes tests for reliability and features a selective copying example for your reference if you're interested.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-11-06T20:22:59.230Z
Answers.LastActivityDate: 2024-11-06T20:22:59.230Z
Answers.UserDisplayName: cheind
Body: <p>My implementation is as in <a href="https://arxiv.org/pdf/2410.01201" rel="nofollow noreferrer">Were RNNs All We Needed?</a>
  page 17-18 minGRU. However I seem to get a mismatch between the parallel mode and the sequentual mode for minGRU.</p>
  <pre><code>import torch 
  import torch.nn as nn 
  import torch.nn.functional as F
  
  
  class minGRU(nn.Module):
      def __init__(self,input_size:int,hidden_size:int):
          super().__init__()
          self.input_size = input_size 
          self.hidden_size = hidden_size 
  
          self.linear_h = nn.Linear(input_size,hidden_size)
          self.linear_z = nn.Linear(input_size,hidden_size)
      
      @staticmethod
      def g(x:torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          The continuous function g ensures that ht ← g(Lineardh (xt)) is positive
          &quot;&quot;&quot;
          return torch.where(x &gt;= 0, x+0.5, torch.sigmoid(x))
      
      @staticmethod
      def log_g(x:torch.Tensor) -&gt; torch.Tensor:
          return torch.where(x &gt;= 0, (F.relu(x)+0.5).log(),-F.softplus(-x))
  
      @staticmethod
      def parallel_scan_log(log_coeffs:torch.Tensor, log_values:torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          log_coeffs: [batch_size, seq_len, input_size]
          log_values: [batch_size, seq_len + 1, input_size]
  
          &quot;&quot;&quot;
          a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
          log_h0_plus_b_star = torch.logcumsumexp(log_values - a_star, dim=1)
          log_h = a_star + log_h0_plus_b_star
          return torch.exp(log_h)[:, 1:]
  
      def sequential_forward(self, x_t:torch.Tensor, h_prev:torch.Tensor=None) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          x_t     : [batch_size, input_size]
          h_prev  : [batch_size, hidden_size]
          &quot;&quot;&quot;
          if(h_prev is None):
              h_prev = torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device)
  
          z = torch.sigmoid(self.linear_z(x_t))
          h_tilde = self.g(self.linear_h(x_t))
          h_t = (1 - z) * h_prev + z * h_tilde
          return h_t
  
      def forward(self, x:torch.Tensor, h_0:torch.Tensor=None) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          x: [batch_size, seq_len, input_size]
          h_0: [batch_size, 1, hidden_size]
          &quot;&quot;&quot;
          if(h_0 is None):
              h_0 = torch.zeros((x.size(0),1,self.hidden_size),device=x.device)
  
  
          k = self.linear_z(x)
          log_z = -F.softplus(-k)
          log_coeffs = -F.softplus(k)
          log_h_0 = self.log_g(h_0)
          log_tilde_h = self.log_g(self.linear_h(x))
          h = self.parallel_scan_log(log_coeffs,torch.cat([log_h_0, log_z + log_tilde_h], dim=1))
          return h #(batch_size,seq_size,hidden_size)
  
  
  
  
  if __name__ == &quot;__main__&quot;:
      batch_size = 3
      seq_size = 2
      input_size =5 
      x = torch.randn(batch_size,seq_size,input_size)
  
      model = minGRU(input_size,6)
      print(model(x))
  
      output_list = []
      ht = None
      for i in range(x.size(1)):
          ht = model.sequential_forward(x[:,i,:],ht)
          output_list.append(ht.unsqueeze(dim=1))
      
      print(torch.cat(output_list,dim=1))
  
  </code></pre>
  <p>I would expect the output of the sequential and the parallel to match, however for me
  this doesn't seem to be the case.
  <br>
  <b>Parallel output </b></p>
  <pre><code>tensor([[[0.4516, 0.4545, 0.4466, 0.3603, 0.8277, 0.9080],
           [0.4142, 0.3942, 0.3565, 0.3890, 0.8032, 0.6870]],
  
          [[0.8782, 0.4790, 0.5668, 0.3825, 1.6276, 0.5832],
           [0.6244, 1.0245, 1.1185, 0.5435, 0.9712, 0.7073]],
  
          [[0.4372, 0.4457, 0.3962, 0.4179, 0.6144, 1.0675],
           [0.5930, 0.4539, 0.4642, 0.3201, 1.2294, 1.0219]]],
         grad_fn=&lt;SliceBackward0&gt;)
  </code></pre>
  <p><b> Sequential output</b></p>
  <pre><code>tensor([[[0.2932, 0.1517, 0.1377, 0.2691, 0.4856, 0.7588],
           [0.3508, 0.2253, 0.1886, 0.3250, 0.6194, 0.6443]],
  
          [[0.7231, 0.1060, 0.1977, 0.2872, 1.4267, 0.2392],
           [0.5323, 0.8793, 0.9940, 0.4883, 0.8674, 0.4427]],
  
          [[0.2369, 0.1406, 0.1586, 0.2885, 0.2937, 0.9327],
           [0.5108, 0.2299, 0.2871, 0.2936, 1.0596, 0.9325]]],
         grad_fn=&lt;CatBackward0&gt;)
  </code></pre>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-10-15T12:59:41.860Z
Comments.Text: It doesn't seem to be the case that I use linear_z and linear_h in the wrong positions, as when I fix the parameters to be the same by adding         self.linear_h.load_state_dict(self.linear_z.state_dict()) at the end of the constructor I still get the same problem. (so even when I set the parameters of linear_z to be equal to linear_h the problem persists)
Comments.UserDisplayName: hal9000
Tags: <recurrent-neural-networks>
Title: minGRU mismatch parallel log space and sequentual logspace
UserDisplayName: hal9000
ThreadText: # minGRU mismatch parallel log space and sequentual logspace
  
  (Asked by: username_0 on 2024-10-15)
  
  My implementation is as in [Were RNNs All We Needed?](https://arxiv.org/pdf/2410.01201)
  page 17-18 minGRU. However I seem to get a mismatch between the parallel mode and the sequentual mode for minGRU.
  
  ```
  import torch 
  import torch.nn as nn 
  import torch.nn.functional as F
  
  class minGRU(nn.Module):
      def __init__(self,input_size:int,hidden_size:int):
          super().__init__()
          self.input_size = input_size 
          self.hidden_size = hidden_size 
  
          self.linear_h = nn.Linear(input_size,hidden_size)
          self.linear_z = nn.Linear(input_size,hidden_size)
      
      @staticmethod
      def g(x:torch.Tensor) -> torch.Tensor:
          """
          The continuous function g ensures that ht ← g(Lineardh (xt)) is positive
          """
          return torch.where(x >= 0, x+0.5, torch.sigmoid(x))
      
      @staticmethod
      def log_g(x:torch.Tensor) -> torch.Tensor:
          return torch.where(x >= 0, (F.relu(x)+0.5).log(),-F.softplus(-x))
  
      @staticmethod
      def parallel_scan_log(log_coeffs:torch.Tensor, log_values:torch.Tensor) -> torch.Tensor:
          """
          log_coeffs: [batch_size, seq_len, input_size]
          log_values: [batch_size, seq_len + 1, input_size]
  
          """
          a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
          log_h0_plus_b_star = torch.logcumsumexp(log_values - a_star, dim=1)
          log_h = a_star + log_h0_plus_b_star
          return torch.exp(log_h)[:, 1:]
  
      def sequential_forward(self, x_t:torch.Tensor, h_prev:torch.Tensor=None) -> torch.Tensor:
          """
          x_t     : [batch_size, input_size]
          h_prev  : [batch_size, hidden_size]
          """
          if(h_prev is None):
              h_prev = torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device)
  
          z = torch.sigmoid(self.linear_z(x_t))
          h_tilde = self.g(self.linear_h(x_t))
          h_t = (1 - z) * h_prev + z * h_tilde
          return h_t
  
      def forward(self, x:torch.Tensor, h_0:torch.Tensor=None) -> torch.Tensor:
          """
          x: [batch_size, seq_len, input_size]
          h_0: [batch_size, 1, hidden_size]
          """
          if(h_0 is None):
              h_0 = torch.zeros((x.size(0),1,self.hidden_size),device=x.device)
  
          k = self.linear_z(x)
          log_z = -F.softplus(-k)
          log_coeffs = -F.softplus(k)
          log_h_0 = self.log_g(h_0)
          log_tilde_h = self.log_g(self.linear_h(x))
          h = self.parallel_scan_log(log_coeffs,torch.cat([log_h_0, log_z + log_tilde_h], dim=1))
          return h #(batch_size,seq_size,hidden_size)
  
  if __name__ == "__main__":
      batch_size = 3
      seq_size = 2
      input_size =5 
      x = torch.randn(batch_size,seq_size,input_size)
  
      model = minGRU(input_size,6)
      print(model(x))
  
      output_list = []
      ht = None
      for i in range(x.size(1)):
          ht = model.sequential_forward(x[:,i,:],ht)
          output_list.append(ht.unsqueeze(dim=1))
      
      print(torch.cat(output_list,dim=1))
  
  ```
  
  I would expect the output of the sequential and the parallel to match, however for me
  this doesn't seem to be the case.
    
  **Parallel output**
  
  ```
  tensor([[[0.4516, 0.4545, 0.4466, 0.3603, 0.8277, 0.9080],
           [0.4142, 0.3942, 0.3565, 0.3890, 0.8032, 0.6870]],
  
          [[0.8782, 0.4790, 0.5668, 0.3825, 1.6276, 0.5832],
           [0.6244, 1.0245, 1.1185, 0.5435, 0.9712, 0.7073]],
  
          [[0.4372, 0.4457, 0.3962, 0.4179, 0.6144, 1.0675],
           [0.5930, 0.4539, 0.4642, 0.3201, 1.2294, 1.0219]]],
         grad_fn=<SliceBackward0>)
  
  ```
  
  **Sequential output**
  
  ```
  tensor([[[0.2932, 0.1517, 0.1377, 0.2691, 0.4856, 0.7588],
           [0.3508, 0.2253, 0.1886, 0.3250, 0.6194, 0.6443]],
  
          [[0.7231, 0.1060, 0.1977, 0.2872, 1.4267, 0.2392],
           [0.5323, 0.8793, 0.9940, 0.4883, 0.8674, 0.4427]],
  
          [[0.2369, 0.1406, 0.1586, 0.2885, 0.2937, 0.9327],
           [0.5108, 0.2299, 0.2871, 0.2936, 1.0596, 0.9325]]],
         grad_fn=<CatBackward0>)
  
  ```
  
  **Tags:** recurrent-neural-networks
  
  > **Comments:**
  > **username_0**: It doesn't seem to be the case that I use linear_z and linear_h in the wrong positions, as when I fix the parameters to be the same by adding         self.linear_h.load_state_dict(self.linear_z.state_dict()) at the end of the constructor I still get the same problem. (so even when I set the parameters of linear_z to be equal to linear_h the problem persists)  
  
  ---
  
  ## 2 answers
  
  ### Answer by username_0 on 2024-10-15. Score: 0 | Accepted answer
  
  In the sequential forward
  
  ```
  h_prev = torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device)
  
  ```
  
  to
  
  ```
  h_prev = self.g(torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device))
  
  ```
  
  This results in
  
  ```
  import torch 
  import torch.nn as nn 
  import torch.nn.functional as F
  
  class minGRU(nn.Module):
      def __init__(self,input_size:int,hidden_size:int):
          super().__init__()
          self.input_size = input_size 
          self.hidden_size = hidden_size 
  
          self.linear_h = nn.Linear(input_size,hidden_size)
          self.linear_z = nn.Linear(input_size,hidden_size)
      
      @staticmethod
      def g(x:torch.Tensor) -> torch.Tensor:
          """
          The continuous function g ensures that ht ← g(Lineardh (xt)) is positive
          """
          return torch.where(x >= 0, x+0.5, torch.sigmoid(x))
      
      @staticmethod
      def log_g(x:torch.Tensor) -> torch.Tensor:
          return torch.where(x >= 0, (F.relu(x)+0.5).log(),-F.softplus(-x))
  
      @staticmethod
      def parallel_scan_log(log_coeffs:torch.Tensor, log_values:torch.Tensor) -> torch.Tensor:
          """
          log_coeffs: [batch_size, seq_len, input_size]
          log_values: [batch_size, seq_len + 1, input_size]
  
          """
          a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
          log_h0_plus_b_star = torch.logcumsumexp(log_values - a_star, dim=1)
          log_h = a_star + log_h0_plus_b_star
          return torch.exp(log_h)[:, 1:]
  
      def sequential_forward(self, x_t:torch.Tensor, h_prev:torch.Tensor=None) -> torch.Tensor:
          """
          x_t     : [batch_size, input_size]
          h_prev  : [batch_size, hidden_size]
          """
          if(h_prev is None):
              h_prev = self.g(torch.zeros((x_t.size(0),self.hidden_size),device=x_t.device))
  
          z = torch.sigmoid(self.linear_z(x_t)) 
          h_tilde = self.g(self.linear_h(x_t)) 
          h_t = (1 - z) * h_prev + z * h_tilde
  
          return h_t
  
      def forward(self, x:torch.Tensor, h_0:torch.Tensor=None) -> torch.Tensor:
          """
          x: [batch_size, seq_len, input_size]
          h_0: [batch_size, 1, hidden_size]
          """
          if(h_0 is None):
              h_0 = torch.zeros((x.size(0),1,self.hidden_size),device=x.device)
  
          k = self.linear_z(x) 
          log_coeffs = -F.softplus(k)
          log_z = -F.softplus(-k)
          log_tilde_h = self.log_g(self.linear_h(x)) 
  
          log_h_0 = self.log_g(h_0)
          h = self.parallel_scan_log(log_coeffs,torch.cat([log_h_0, log_z + log_tilde_h], dim=1))
          return h #(batch_size,seq_size,hidden_size)
  
  if __name__ == "__main__":
      batch_size = 3
      seq_size = 2
      input_size =5 
      x = torch.randn(batch_size,seq_size,input_size)
  
      model = minGRU(input_size,6)
      print(model(x))
  
      output_list = []
      ht = None
      for i in range(x.size(1)):
          ht = model.sequential_forward(x[:,i,:],ht)
          output_list.append(ht.unsqueeze(dim=1))
      
      print(torch.cat(output_list,dim=1))
  
  ```
  
  **Parallel output**
  
  ```
  tensor([[[0.6535, 0.4925, 0.3262, 0.1168, 0.1602, 0.1158],
           [1.0672, 0.7427, 0.3846, 0.1483, 0.2766, 0.1792]],
  
          [[0.6006, 0.1788, 0.2839, 0.1379, 0.1121, 0.1104],
           [0.4730, 0.4698, 0.3884, 0.2961, 0.2522, 0.2705]],
  
          [[0.2452, 0.1759, 0.1381, 0.1435, 0.9452, 1.0829],
           [0.3459, 0.3899, 0.3189, 1.0957, 0.5292, 0.8135]]],
         grad_fn=<SliceBackward0>)
  
  ```
  
  **Sequential output**
  
  ```
  tensor([[[0.6535, 0.4925, 0.3262, 0.1168, 0.1602, 0.1158],
           [1.0672, 0.7427, 0.3846, 0.1483, 0.2766, 0.1792]],
  
          [[0.6006, 0.1788, 0.2839, 0.1379, 0.1121, 0.1104],
           [0.4730, 0.4698, 0.3884, 0.2961, 0.2522, 0.2705]],
  
          [[0.2452, 0.1759, 0.1381, 0.1435, 0.9452, 1.0829],
           [0.3459, 0.3899, 0.3189, 1.0957, 0.5292, 0.8135]]],
         grad_fn=<CatBackward0>)
  ```
  
  ```
  
  
  ---
  
  ### Answer by username_1 on 2024-11-06. Score: 0
  
  While I can't directly identify the specific issue with your code, I might recommend checking out the
  
  [mingru](https://github.com/username_1/mingru) repository I've recently created.
  
  The code addresses minor paper inconsistencies and ensures consistent behavior between sequential and log-space parallel modes (validated up to 3 stacked layers with 1-e5 precision). The repository includes tests for reliability and features a selective copying example for your reference if you're interested.
  
  
  ---
  
  