An MLP is a simple linear layer followed by a non-linearity i.e. each Expert Args: dim (int): The input dimension of the linear layer. dropout (float, optional): The dropout probability. Defaults to 0.1. Attributes: net (nn.Sequential): The sequential network consis
| 9 | |
| 10 | # Expert module |
| 11 | class Expert(nn.Module): |
| 12 | """An MLP is a simple linear layer followed by a non-linearity i.e. each Expert |
| 13 | |
| 14 | Args: |
| 15 | dim (int): The input dimension of the linear layer. |
| 16 | dropout (float, optional): The dropout probability. Defaults to 0.1. |
| 17 | |
| 18 | Attributes: |
| 19 | net (nn.Sequential): The sequential network consisting of linear layers, ReLU activation, and dropout. |
| 20 | |
| 21 | """ |
| 22 | |
| 23 | def __init__(self, dim: int, dropout: int = 0.1): |
| 24 | super().__init__() |
| 25 | self.net = nn.Sequential( |
| 26 | BitLinear(dim, 4 * dim), |
| 27 | nn.ReLU(), |
| 28 | BitLinear(4 * dim, dim), |
| 29 | nn.Dropout(dropout), |
| 30 | ) |
| 31 | |
| 32 | def forward(self, x): |
| 33 | return self.net(x) |
| 34 | |
| 35 | |
| 36 | # Changing the above to accomodate noisy top-k gating |