Recurrent Network

[MiaoBBBY]

class MambaModel(nn.Module):
config = {}

def __init__(self, positional_embedding):
    super().__init__()
    mamba_config = {
        "d_model": self.config["d_model"],
        "d_state": self.config["d_state"],
        "d_conv": self.config["d_conv"],
        "expand": self.config["expand"],
    }
    self.mamba_forward = nn.Sequential(*[Mamba(**mamba_config) for _ in range(self.config["num_layers"])])
    pe = positional_embedding[None, :, :]
    if self.config.get("trainable_pe"):
        self.pe = nn.Parameter(pe)
    else:  # fixed positional embedding
        self.register_buffer("pe", pe)

def forward(self, output_shape, condition=None):
    assert len(condition.shape) == 3
    x = self.mamba_forward(self.pe.repeat(output_shape[0], 1, 1) + condition)
    return x

I noticed that in the actual recurrent network, the input of Mamba only includes positional encoding and condition input. However, the description in the paper states: "After obtaining the parameter tokens K[i], permutation states S, and position embeddings e[i], we feed them into a recurrent network f(·) that learns tokenwise representations while capturing cross-token dependencies." This seems to be inconsistent with the actual input to the recurrent network.

[MTDoven]

Sorry for the confusion caused by the code implementation. Please refer to line 52 of https://github.com/NUS-HPC-AI-Lab/Recurrent-Parameter-Generation/blob/main/model/__init__.py: although the permutation state and positional embedding are passed in from different sections in the actual engineering implementation, they are consistent with the description in the paper.

[MiaoBBBY]

Sorry for the confusion caused by the code implementation. Please refer to line 52 of https://github.com/NUS-HPC-AI-Lab/Recurrent-Parameter-Generation/blob/main/model/__init__.py: although the permutation state and positional embedding are passed in from different sections in the actual engineering implementation, they are consistent with the description in the paper.

Thanks for your clarification. So, to confirm, the recurrent model mentioned in the paper refers to 'RPG', not Mamba. And Mamba's role is to process the model's architecture and condition information, which isn't closely tied to the model's weights.