<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>

CC BY-SA 4.0
2024-10-15T14:03:43.827Z
2024-10-15T14:12:45.820Z
2024-10-15T14:12:45.820Z
hal9000
<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>

CC BY-SA 4.0
2024-11-06T20:22:59.230Z
2024-11-06T20:22:59.230Z
cheind
<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>

CC BY-SA 4.0
2024-10-15T12:59:41.860Z
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)
hal9000
<recurrent-neural-networks>
minGRU mismatch parallel log space and sequentual logspace
hal9000
# 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?
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=)

 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=)

**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=)


 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=)
```


---

### 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.


---

