BitMoE (Bitwise Mixture of Experts) module. Args: dim (int): The input dimension. num_experts (int): The number of experts in the mixture. top_k (int, optional): The number of experts to select for each input. Defaults to 2.
| 81 | |
| 82 | |
| 83 | class BitMoE(nn.Module): |
| 84 | """ |
| 85 | BitMoE (Bitwise Mixture of Experts) module. |
| 86 | |
| 87 | Args: |
| 88 | dim (int): The input dimension. |
| 89 | num_experts (int): The number of experts in the mixture. |
| 90 | top_k (int, optional): The number of experts to select for each input. Defaults to 2. |
| 91 | """ |
| 92 | |
| 93 | def __init__(self, dim: int, num_experts: int, top_k: int = 2): |
| 94 | super(BitMoE, self).__init__() |
| 95 | self.router = NoisyTopkRouter(dim, num_experts, top_k) |
| 96 | self.experts = nn.ModuleList([Expert(dim) for _ in range(num_experts)]) |
| 97 | self.top_k = top_k |
| 98 | |
| 99 | def forward(self, x): |
| 100 | gating_output, indices = self.router(x) |
| 101 | final_output = torch.zeros_like(x) |
| 102 | |
| 103 | # Reshape inputs for batch processing |
| 104 | flat_x = x.view(-1, x.size(-1)) |
| 105 | flat_gating_output = gating_output.view(-1, gating_output.size(-1)) |
| 106 | |
| 107 | # Process each expert in parallel |
| 108 | for i, expert in enumerate(self.experts): |
| 109 | # Create a mask for the inputs where the current expert is in top-k |
| 110 | expert_mask = (indices == i).any(dim=-1) |
| 111 | flat_mask = expert_mask.view(-1) |
| 112 | |
| 113 | if flat_mask.any(): |
| 114 | expert_input = flat_x[flat_mask] |
| 115 | expert_output = expert(expert_input) |
| 116 | |
| 117 | # Extract and apply gating scores |
| 118 | gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1) |
| 119 | weighted_output = expert_output * gating_scores |
| 120 | |
| 121 | # Update final output additively by indexing and adding |
| 122 | final_output[expert_mask] += weighted_output.squeeze(1) |
| 123 | |
| 124 | return final_output |
| 125 | |
| 126 | |
| 127 | # x = torch.randn(2, 4, 8) |