I am trying to implement the a mixture of expert layer, similar to the one described in:

  
      
      arXiv.org
  
  
    

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer 34

The capacity of a neural network to absorb information is limited by its
number of parameters. Conditional computation, where parts of the network are
active on a per-example basis, has been proposed in theory as a way of
dramatically increasing...


  
  
    
    
  
  


Basically this layer have a number of sub-layers F_i(x_i) which process a projected version of the input. There is also a gating layer G_i(x_i) which is basically an attention mechanism over all sub-expert-layers:
sum(G_i(x_i)*F_i(x_i).
My Naive approach is to build a list for the sub-layers:
sublayer_list = nn.ModuleList()
for i in range(num_of_layer):
sublayer_list.append(self.make_layer())
Then when applying this I use another for loop
out_list= []
for i,l in enumerate(sublayer_list):
out_list.appned(l(input(i)))
However the addition of this Mixture-of-Expert layer slows training by almost 7 times (against one with MoE layer swapped for a similar-sized MLP). I am wondering if there are more efficient ways to implement this in pytorch? Many thanks!