Benchmarking speed against pytorch nn.GRU

[yenhochen]

Hi, I'm wondering if you have benchmarked your code against the pytorch nn.GRU on GPU. In my own experiments, I'm finding that this implementation of MinGRU is 10x slower than nn.GRU on GPUs. Would love to chat more if you have any insights!

[cheind]

Hi,

Yes, this issue has also been reported in other Mingru repositories. For more details, see this discussion.

The key points are:
a) nn.GRU is highly optimized at the CUDA level, whereas reference implementations are not.
b) The authors of the original paper compared their implementation not to nn.GRU but to a reference implementation of nn.GRU.
c) At very long sequences Mingru reference implementations are improving compared to nn.GRU

[yenhochen]

Thank you for your response! These reasons make sense to me. I have found that making the the following modifications to your functional.py code ends up achieving significantly faster performance so that run time is at least comparable to nn.GRU

def g(x: torch.Tensor) -> torch.Tensor:
    """Activation function for hidden state
    Ensures that g(h) is non-negative and hence avoids
    complex numbers in log-space.
    """
    
    # Use torch.where for conditional computation
    out = torch.where(
        x >= 0,
        x + 0.5,          # Apply linear transformation for non-negative values
        torch.sigmoid(x)  # Apply sigmoid for negative values
    )
    return out

def log_g(x: torch.Tensor) -> torch.Tensor:
    """Log-activation function for hidden state
    Ensures that g(h) is non-negative and hence avoids
    complex numbers in log-space.
    """

    # Use torch.where for conditional computation
    out = torch.where(
        x >= 0,
        (F.relu(x) + 0.5).log(),  # Apply log for non-negative values
        -F.softplus(-x)    # Apply softplus for negative values
    )
    return out

[cheind]

That's strange, I was expecting this to be slower, since where evaluates both branches.