From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Submission Withdrawn by the Authors

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_6RhL
title: Score update after discussing with other reviewers and AC
comment: In a video call with other reviewers and the AC, two issues were raised. First, the significance of perplexity score improvement reported in this paper is hard to interpret in the context of related work (I do not work in NLP, so please refer to comments from other reviewers and the AC for their elaborations on this point). Another issue is whether the "unification" contribution is somewhat overstated. Indeed, the only relevant neural memory paper that proposes a selection method seems to be Lample et al [2019]. Due to these two concerns raised in our discussion, I lowered my score to 5.

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_o8N9
title: Thanks for the clarifications.
comment: I appreciate the authors' feedback and clarifications -- the rating is increased to 5.

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_o8N9
title: "Unified View" is still unclear
comment: Thanks the authors for the detailed explanation regarding the load balancing. 

Because of the subtlety in load balancing, it is a bit tricky to draw "equivalence" statement based on a simple rewrite of routing function. For example, the comment says that because of larger number of experts, mode collapsing problem is not severe anymore. Then the question becomes if Switch Transformer also selects larger number of experts, will the load balancing no longer be necessary? If not, why ReLU plays a different role?

Since the paper has a title about "unified view", my impression is that the paper will spend great amount of discussion on the details of the unification. This goes beyond mathematical formulation, and compares in depth on the equivalence e.g. with or without load balancing on both side. However, the paper's main focus in the later section is to propose the memory block, which could be proposed without mentioning the "unified view", as the "unified view" plays little part in the majority of the paper.

While I agree that avg-K is a novel idea, I would like to emphasize again that it is hard to follow what the paper is trying to achieve: is it mainly to demonstrate the "unified view", or is it trying to propose the "block memory"?

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: There are technical reasons that we can not run small-scale experiments 
comment: We apologize for the late reply. As informed in [a, b], **the smallest standard size for base architecture is about 222M** (larger than Roberta-base). When designing our experiment, we chose the base size of 355M (similar to Roberta-large) because we wanted to strengthen the implication of our result for a larger-scale model.

We thank the reviewer for raising the concern of computational equity. However, this was partly done as an intern project, and we no longer have the resources available to run further models. It is difficult for us to set up baselines in currently studied setups: 
1. [c] tried RandHash with a smaller base architecture size (<220M), but the data used --- MassiveText (10.5 TB) [d] --- is not publicly available and beyond storage space available to the first author. 
2. [b] tried RandHash with a small and public dataset (e.g., wikitext-103) but the model size is beyond the compute resources currently available (two 2080 Ti or one A60) to the first author.

Since the setting for **a model of a smaller size (<220M) and a small (and public) dataset** is understudied in the current literature, it is also challenging to develop the proper training configuration with limited time and compute. 

Our **latest supplementary material included the code used to replace the Feedforward module** in an in-place fashion. We used this code for training models except for Switch. This code might be helpful for other practitioners to conduct experiments with smaller sizes on an academic budget. We are sorry that we currently cannot conduct small-scale experiments with an academic budget. 

[a] Roller, Stephen, Sainbayar Sukhbaatar, and Jason Weston. "Hash layers for large sparse models." Advances in Neural Information Processing Systems 34 (2021): 17555-17566.

[b] William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. ArXiv, abs/2101.03961, 2021

[c] Unified Scaling Laws for Routed Language Models. ICML 2022: 4057-4086

[d] Scaling Language Models: Methods, Analysis & Insights from Training Gopher. arXiv:2112.11446, 2021.

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_iswb
title: Dear authors, would you consider or at least comment on the reproducibility problem?
comment: > The scale of experiments (computation measured in ZFLOPs) seems to be out of reach of ordinary labs, hampering the reproducibility. It would be greatly appreciated if some experiments of smaller scale settings are included.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We simplify the PKM paragraph in background section
comment: > The subsection "Product Key Memory" in Section 2.2.2 does not seem to be necessary. It is very dense in technical details, while not being used anywhere later in the paper. A few sentences to summarize the intuition behind Product Key Memory may suffice

Thanks for the suggestion. In the latest draft, we simplify the PKM subsection in Section 2.2.2.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We deleted our interpretability claim and why
comment: > Thanks for pointing out this reference, but I am not fully convinced that this argument supports the interpretability of Avg-K. After all, the memory coefficients are themself learned in an end-to-end way. It is thus not comparable to the discrete input in deterministic hashing, where discrete inputs are just data -- they are fixed and non-learnable.

We think you are making a very good point with Avg-K also being learned end-to-end, in contrast with **a more interpretable deterministic table**, as in [b]. In [b], each expert is responsible for input from a certain domain, e.g., the 1st expert always processes medical text while the 2nd expert always processes sports news. While the selection could be more fine-grained in Avg-K (which we detailed at the end of the comment), **it is still learned end-to-end and less easily interpretable than a domain-specific table**. As such, we deleted the claim about interpretability in our latest draft. 

Your comment also made us realize it is hard to make such a claim because, so far, **there is no study similar to [a] to interpret a model trained with a hash table** (e.g., RandHash in our work). Previously, we thought Avg-K had this advantage because we thought the techniques in [a] could be directly applied to Avg-K but not to RandHash. **We will give our rationale below**:

For an end-to-end trained transformer, the takeaway in [a] is that key vectors in each feedforward layer strongly correlate to some human-interpretable feature in the input text. In Table 1 in [a], the 449th key in 1st feedforward layer always produces high memory coefficients when the word "substitutes" is input ("shallow" feature), or some other key always has high coefficients when the text is about TV shows ("semantic" feature). We meant the final model trained with Avg-K could directly use this interpretability technique to understand the model's behavior. Practitioners can answer interesting questions like whether a memory block now corresponds to some or a meaningful group of human interpretable features in the text. 


[a] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021

[b] Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettlemoyer. Demix layers: Disentangling domains for modular language modeling. CoRR, abs/2108.05036, 2021. URL https://arxiv.org/abs/2108.05036

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Detailed explanation why our model is slower in training and as fast in inference
comment: > Could the authors explain why their model is slower than a Switch MoE in training but as fast as a Switch model in inference?

We will expand our explanation below and clarify.
The fact that our method has slower training is mainly because Avg-K's implementation is not optimized. Our Switch implementation uses expert parallelism. As such, while the FLOPS cost is the same, Avg-K is slower than a standard MoE implementation. 
With the loss added, the training time will likely match Switch's. 

As for inference, assuming our Avg-K implementation is optimized, and the expert size is the same as Switch Transformer. Both of them have the same amount of parameters for the gate function. As such, we have, at inference time:
1. Neither Switch nor Avg-K will calculate the loss. 
2. Both will do a matrix multiplication(same FLOPs) for expert selection. 
3. Both will select the same amount of key-value pairs from different experts and follow the same memory computation.
Therefore, the inference time should be the same.  

**To give a more detailed explanation:** The reason why an imbalanced MoE is slower is from how `all_to_all` operations distribute tokens to different experts. After the gate calculation, the MoE knows which tokens will be routed to which experts (by taking top-k). Conceptually, some experts receive fewer tokens than others, and such allocation is represented by masking matrices containing only 1 or 0. However, the standard implementation will "pad" each allocation matrix to the largest load. This is similar to why shorter sentences are padded to the longest sentence in the batch. Therefore, every expert will receive allocation matrices of the same size, whereas some matrices have "more zeros' ' than necessary. This will lead to slow training overall. A load-balancing loss will decrease the max number of tokens that an expert will receive because some tokens will be routed to other experts instead. Thus, the allocation matrices will be smaller overall, and training will become faster.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Disagreement with Reviewer 4 (o8N9)
comment: We are receiving conflicting reviews from you and Reviewer 4 (o8N9). It will be great to settle it and further improve our paper. 
Reviewer 4 (o8N9) says 
> primary goal of this paper is unclear... 

> experiments themselves are limited...

but you say
> The paper is well-written and insightful.

Can you and reviewer 4 try to settle this disagreement? 

Thank you!

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_6RhL
title: Response to the rebuttal
comment: I'd like to thank the authors for their reply. I have raised my score. The response from the authors has addressed many of my questions. A few remaining ones are below.

> Similarly to the "Sparse is Enough in Scaling Transformers" paper, our architecture (without load balancing), while being FLOPS matched, is slower during training and as fast during inference compared to Switch and other MoE architectures.

Could the authors explain why their model is slower than a Switch MoE in training but as fast as a Switch model in inference? In my experience, the forward pass of a standard MoE in the training phase can be slow due to the usage of auxiliary losses (e.g., those used in Switch). However, no auxiliary loss is used at inference, so a forward pass of Switch at inference can be faster than a forward pass at training. The proposed model does not use any auxiliary loss during training or inference, and therefore does not have this "extra speedup" when one uses the model for inference. It is thus unintuitive to me why the proposed model is slower than a Switch MoE during training but faster than a Switch MoE during inference. Any explanation for this will be appreciated.


> "While deterministic hashing approaches establish a direct connection between the discrete input and the memory units used to process them, the Avg-K method establishes a such connection through the memory coefficients as in Geva et al. 2021 [a]."

Thanks for pointing out this reference, but I am not fully convinced that this argument supports the interpretability of Avg-K. After all, the memory coefficients are themself learned in an end-to-end way. It is thus not comparable to the discrete input in deterministic hashing, where discrete inputs are just data -- they are fixed and non-learnable.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We make full use of PKM for our study.
comment: > the authors did not even try many other direct* selection methods besides PKM.

The focus of our paper is to study existing methods; under the constraints of equal sizes of parameters, PKM is indeed the best choice we could find. In our preliminary study, we did try setting different hyperparameters or designs for PKM, but the improvement is small. To make full use of PKM, we decide to propose PKM-FFN for narrowed analysis (second half of section 5.2 or section 4.1 in the updated draft).

* We think the author meant "indirect"

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We considered other selection mechanisms for sparse neural memory
comment: >  Is the paper [Lample et al, 2019] the only paper proposing selection methods? 

In our early stage, we consider methods under “Approximate Nearest Neighbour(ANN) search” (in section 7) as sparse memory methods, since they apply dimension reduction to input and key tables (similar to ones in [Lample et al, 2019] ). However, as suggested by our preliminary study and Appendix D.2, we found this line of work doesn’t yield practical modeling choices as well as methods in [Lample et al., 2019]. Therefore, we excluded them in our main paper and put them in related work.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: LoRKM is considered
comment: > why they do not include LoRKM in their experiment results, but only use PKM's results? 

In Figure 5, we indeed include the results for LoRKM. This is a more detailed plot for Figure 3 in the main paper.
We didn’t include LoRKM in our narrowed analysis (section 5.1 and the second half of section 5.2) because it’s easier for us to use PKM to align with other methods in terms of parameters. To make full use of PKM, we decided to propose PKM-FFN for narrowed analysis (second half of section 5.2 or section 4.1 in the updated draft).
We considered running LoRKM with a small enough low-rank dimension (equivalent to LoRKM with low-rank-dim = 2.46), but we believe it will under-perform current PKM.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Considerations for not adding load-balancing loss
comment: > load balancing is not considered.

In our first submitted version, we pointed out in the “limitation” paragraph (section 6) that load-balancing is important for future work to incorporate for better unification and requires more advanced theoretical justification for its connection with memory block and block selection method.  

The requirement of a load balancing loss is determined by the sparsity of gradients. If one “expert” gets updated with gradients while other “experts” are starved then the one expert will be selected all of the time leading to a form of mode collapse. In our best case, since our gradients are both sparse but also include a larger number of experts — Switch has #expert=1 (row 4 in Table 4), ours have #expert= 64 (row 4 in Table 4) --- we do not have a severe mode collapse problem. As such, probabilistically, there will always be a gradient flow for most parts of the network to ensure good performance. From the empirical data there seems to be for example more gradient flow than when using ReLU activation functions in a standard transformer (block size = 1) where about 50-90% of hidden units are dead [c]. As such, our setup is more comparable to using a ReLU function than the load balancing required for a mixture of experts. Since ReLU does not need a load balancing loss, it is also not required in our case where most of the network is probabilistically selected over time.

During our experiment, we didn't use the load-balancing loss for a more empirical consideration. In our study with *a block version* of VanillaM (section 5.1), we found its competitive performance without load balancing. VanillaM differs from Avg-K in that VanillaM takes the average after applying GeLU to the hidden states, and Avg-K averages before GeLU. Due to the similarity in form, it is a natural follow-up for us to try Avg-K without load balancing. It is also to our surprise that Avg-K is the first MoE to perform well without load balancing while the Switch Transformer may degenerate [a, b]. 

[a] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations, ICLR 2017
[b] David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Workshop Track Proceedings, 2014
[c] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: The difference between methods is significant
comment: > The authors draw conclusions with noisy metrics.

We thank the reviewer for raising this clarification question. The difference between methods is significant because a model **consistently** outperforms another one across different Out-of-domain(OOD) datasets. The average OOD perplexity generally reflects the relative performance of each domain (with 1 or 2 exceptions out of 22). Generally, people assume the same distribution when talking about statistical significance, whereas we have 22 different ones in our OOD evaluation. It’s unfair to ask a model to achieve, on a very hard or distant OOD distribution, the same perplexity on a very easy or similar OOD distribution. Therefore, it’s a norm that a model will have very different perplexity on different OOD datasets.

In the updated draft, we added Tables 9 and 10 in the Appendix and removed the variance in Tables 3 and 4 to avoid confusion.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We do have an outline of the paper at the end of the introduction sections.
comment: > The primary goal of this paper is unclear

Thanks for raising this clarification question for us. We do have an outline of the paper at the end of the introduction sections. We will clarify the outline of the paper here. We are the first paper to exemplify not only that there exists a mathematical connection between different methods, but we could indeed use the connection to compare different choices of model and rationalize at a better design --- in our case, Avg-K. 
The equivalence is indeed hard to control in general. In our analysis, we try to align equivalence by *using the number of learnable parameters* in the model when we do a narrow analysis (section 5.1 and the second half of section 5.2). So that the model theoretically has the same expressiveness. We will expand on how we do it below:

1. In section 5.1, RandHash uses a learning-free gating, and the key and value tables between RandHash and VanillaM are of the same size. 
2. In the first half of section 5.2, we notice the FLOPs are hard to control when sparse memory like LoRKM has more choices of hyper-parameters. Therefore, we study the overall trend by choosing their default/recommended values from the original paper and changing the memory size. We do have a more detailed plot in Figure 5 (See Appendix). 
3. In the second half of section 5.2, when we study the effect of the gating function, we again align the equivalence by the number of total parameters, especially rows 3-7 in Table 3. We do want to note that our Switch Transformer also has the same amount of parameters. However, it doesn’t perform as well as RandHash (in Table 4) and has different memory block sizes, so we exclude it from our analysis for controlled experiments.

In the updated draft, we clarify how we control the equivalence (Section 4.1). Thanks for the suggestion.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: While being FLOPS matched, Avg-K is slower during training and as fast during inference compared to Switch and other MoE architectures
comment: > models with similar FLOPS may have drastically different run times.

Similarly to the "Sparse is Enough in Scaling Transformers" paper, our architecture (without load balancing), while being FLOPS matched, is slower during training and as fast during inference compared to Switch and other MoE architectures. We currently need optimized implementations for our sparse approach while our MoE implementation is optimized. This makes a fair comparison hard. However, during inference, when computing token-by-token generations, our selection method is cheap because the average can be pre-computed and thus reused. In comparison with VanillaM, the speedups in this case are close to the FLOPS savings.
In general, comparing different models' training times and inference times is hard. Usually, one needs to have an optimal (or near optimal) implementation to make a fair claim about the advantage of one over another. Preferably, different models should use the same codebase for controlled experiments. However, such a codebase or implementation doesn't exist, and building such a codebase/implementation is beyond the scope of our paper.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Why we think interpretability is an advantage of our paper
comment: > do not find interpretability is a strength of Avg-K, since the key table is end-to-end learnable and does not conform to any prior knowledge or human intuition.

We acknowledge the feedback regarding the interpretability of Avg-K and the comparison to hash-based models in this respect. While deterministic hashing approaches establish direct connection between the discrete input and the memory units used to process them, the Avg-K method establishes such connection through the memory coefficients as in Geva et al. 2021 [a]. 

[a] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Choice of Hash Function is informed
comment: > While this may be true for random hashing, it should be noted that it is also possible to perform deterministic hashing, which is fully interpretable and typically yields slightly better results than random hashing [3].

We did notice [3] tried many hash functions (Table 3 in [3]). Our choice of random hashing is from the first author's recommendation since they found this as their simplest and most elegant design choice. Table 3 in [3] shows that the best hash functions are "Dispersed Hash (within token clusters)" and "Balanced assignment"; these two methods also assign tokens to expert **in a (semi-)random fashion**. We quote [3] as below:
> **Dispersed Hash (within token clusters)** …. use the same k-means [clustering with a fixed number of clusters using token embeddings from a baseline Transformer model], but distribute all tokens within each cluster equally across all buckets.

> **Balanced assignment** …. build the lookup table before training the model using the training data distribution by greedily assigning the most frequent tokens to the emptiest buckets.

Therefore, the semantics of different experts are still unclear.


[3] Roller, Stephen, Sainbayar Sukhbaatar, and Jason Weston. "Hash layers for large sparse models." Advances in Neural Information Processing Systems 34 (2021): 17555-17566.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Yes, we expect load-balancing loss to further help Avg-K.
comment: > Would the authors expect standard load balancing losses such as those in [1] and [2] to work in the Avg-K router?

Yes, we expect load-balancing loss to further help Avg-K. We will state our rationale as follows. Comparing VanillaM and Avg-K (in Appendix C.1 or updated draft’s Appendix C.2), one would expect Avg-K to be greatly affected by extremely negative hidden states (before GeLU). Yet, the final model with Avg-K could even outperform VanillaM with the same block size (Figure 6). This means the model will accommodate small design changes. To Avg-K, load balancing is considered “small” because it is an **engineering** fix for the training time and not a necessary part of **learning**. Therefore, we believe that with the loss added, the model will accommodate and could still perform competitively. 

The requirement of a load balancing loss is determined by the sparsity of gradients. If one “expert” gets updated with gradients while other “experts” are starved then the one expert will be selected all of the time leading to a form of mode collapse. In fact, with the same memory block size ($g  = 4096$), we are surprised to observe Avg-K (**w/o load balancing** loss; row 6 in Table 4) could still perform on par with Switch (**w/ load balancing**; row 4 in Table 4). As our load balancing analysis suggests in Appendix C.3, when the number of experts is small, the mode collapse issue in Avg-K is severe. This makes us more confident that Avg-K will perform better with standard load balancing loss added.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Considerations for not adding load-balancing loss
comment: > why do they decide to not add any load-balancing loss? 

The requirement of a load balancing loss is determined by the sparsity of gradients. If one “expert” gets updated with gradients while other “experts” are starved then the one expert will be selected all of the time leading to a form of mode collapse. In our best case, since our gradients are both sparse but also include a larger number of experts — Switch has #expert=1 (row 4 in Table 4), ours have #expert= 64 (row 4 in Table 4) --- we do not have a severe mode collapse problem. As such, probabilistically, there will always be a gradient flow for most parts of the network to ensure good performance. From the empirical data there seems to be for example more gradient flow than when using ReLU activation functions in a standard transformer (block size = 1) where about 50-90% of hidden units are dead [c]. As such, our setup is more comparable to using a ReLU function than the load balancing required for a mixture of experts. Since ReLU does not need a load balancing loss, it is also not required in our case where most of the network is probabilistically selected over time.

During our experiment, we didn't use the load-balancing loss for a more empirical consideration. In our study with *a block version* of VanillaM (section 5.1), we found its competitive performance without load balancing. VanillaM differs from Avg-K in that VanillaM takes the average after applying GeLU to the hidden states, and Avg-K averages before GeLU. Due to the similarity in form, it is a natural follow-up for us to try Avg-K without load balancing. It is also to our surprise that Avg-K is the first MoE to perform well without load balancing while the Switch Transformer may degenerate [a, b]. 


[a] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations, ICLR 2017

[b] David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Workshop Track Proceedings, 2014

[c] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: Load balancing constraints were **not** turned off in our SMoE baselines
comment: > I am wondering if the load balancing constraints were turned off in certain SMoE experiments

No. Experiments with Switch transformers always have load balancing enforced. Although RandHash doesn’t have a load balancing loss, its load balancing is enforced by uniformly routing tokens to experts (“Static gate” in Section 2.2.1).

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We empirically verified the aggregator choice
comment: > the choice of averaging operation is rather arbitrary, and it is desirable to have proper ablations against alternatives like "maximum dot-product of the block" (the classic max-pooling vs. average-pooling).

In our first submitted draft, we empirically verified the aggregator choice. In our experiments with the block version of VanillaM, in Table 7, we tried using the max pooling over hidden states to select the max hidden states in a block (after GeLU). However, this produces worse results than using the average function and shows that a memory block hardly has a single key-value cell that dominates over others. We hypothesize this is because lots of hidden states are overlooked by the max operation, yet the overlooked hidden states still contribute to the computation. In contrast, the performance increases when we consider the average (or average of the absolute values) where every hidden state contributes to the decision. Although the situation is slightly different in Avg-K, the “max-pooled” version of Avg-K will only overestimate the hidden states information even more, and the aggregated value won’t be indicative of the hidden states used for computation.

In the updated draft, we update our rationale for using Average in Appendix C.1

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We do have some bridging explanations at the start of section 6 and we improve it in the updated draft.
comment: > However, a major weakness of the paper is its organization ....  Consequently, section 6 seems to be rather a disjointed part of the paper, though possessing the best results.

Thanks for the suggestions on better organizing the paper. We do have some bridging explanations at the start of section 6. In the updated draft, we improve the transition between sections and better motivate the proposal of Avg-K by summarizing our analysis results and adding pointers to supporting materials. To iterate our observation in section 5.2 that motivates the design choice for Avg-K: 
* **Recap: why is Avg-K a gating function rather than based on sparse memory?**  We found having a gating design helps improve the FLOPs-Perplexity trade-off (Figure 3). 
* **Recap: why doesn’t Avg-K have any factorization in the key table?** PKM$^{E=32}$ and PKM-FFN$^{E=16}$ have the same amount of parameters, but PKM greatly underperforms PKM-FFN (row 4 and 6 in Table 3). This means having a full-parameter key table (PKM-FFN) is important and the difference can’t be compensated by doubling the value table (PKM). 
* **Recap: why does Avg-K use an average to aggregate the information from the key table?** With the same memory block size, we observe that RandHash, PKM-FFN, and VanillaM have the same number of parameters, but they lead to different performances. We consider the procedure the input is processed: `input -> (gate) -> key table`

    * RandHash uses a separate gate, which uses zero information from the key table. 
    * PKM-FFN uses a separate gate, which uses some information from the key table when the weights in the gate are jointly learned with the key table. 
    * VanillaM uses the information solely in the key table. 
The performance is ranked (from best to worst): VanillaM, PKM-FFN, and RandHash. **Therefore**, we hypothesize that the more information is used in the key table, the better performance the trained model will have. And, averaging is a linear function and properly maintains information from the key table. In the updated draft, we detailed our rationale in Appendix C.1.

From: ICLR.cc/2023/Conference/Paper5150/Authors
title: We have verified the choice of aggregator empirically.
comment: > In particular, it seems that an average of the absolute values is more reasonable than an average of the signed values.

We have verified the choice of aggregator empirically. With VanillaM (in Table 7 or footnote 1 in section 5.1), we tried the average absolute value (after GeLU), but it didn’t yield promising results. We hypothesized that a positive feature is good at predicting the value of a label/token against all others. In contrast, a negative value is good at negating the prediction of a single token. As such, positive features are more predictive than negative ones. Although the situation might be different for Avg-K (before GeLU), we expect the selection will only be affected more because of the larger impact from the negative value.
Another consideration we have is that the average function is linear. When we select experts, we use the dot product between input and averaged keys. Due to the linearity, this value is equivalent to 1. taking the dot product between input and every key and 2. taking the average (See Appendix C.1 or C.2 in the updated draft). Thus, by using this design choice, we save a great amount of computation compared with VanillaM, while keeping the neural memory analogy. 

In the updated draft, we update our rationale for using Average in Appendix C.1

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_S3p4
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary_of_the_paper: The paper unifies two frameworks (sparse memory and a mixture of experts) to construct sparse feed-forward networks in Transformers. The paper suggests a few valuable suggestions to improve the state-of-the-art.
strength_and_weaknesses: Strength:
1. The paper unifies two popular approaches (sparse memory/mixture of experts) to construct sparse feed-forward networks in Transformers. The differences are 1) block size, 2) importance weights (shared vs. individual), and 3) selection method (direct vs. indirect).
2. The authors empirically study these differences and guide how to choose them.
3. The final version established the new state-of-the-art.

Weakness:
The routing algorithm is not carefully studied --- the authors did not compare the proposed method to other possible candidates. In particular, it seems that an average of the absolute values is more reasonable than an average of the signed values. For example, if two memory cells have a large amplitude, say +A and -A, these two cells are omitted as a group (since their average is 0). However, according to the rationale in sparse memory, they are strongly activated and thus should be chosen.
clarity,_quality,_novelty_and_reproducibility: The paper is well-written and insightful.
summary_of_the_review: Overall, I believe the submission is a high-quality empirical paper that will be used as a benchmark in future research.
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
empirical_novelty_and_significance: 3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.
flag_for_ethics_review: NO.
recommendation: 6: marginally above the acceptance threshold

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_iswb
confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
summary_of_the_paper: The paper studies the sparse scaling problem of Transformer models. A unified formulation is presented which includes Mixture of Experts and Sparse Neural Memory as special cases. The unified formulation enables mixing traits from previous disconnected methods for more powerful hybrids, and allows more thorough ablation studies. Based on observations of effects of the hybrids, a novel sparse activation method Avg-K is proposed, and is empirically shown to be superior in reducing perplexity in NLP tasks like in domain and out-of-domain completion.
strength_and_weaknesses: The strength of the paper lies in its clear line of development till section~5. The unified formulation allows fair comparison between multiple prior arts, and the figures and tables throw multi-facet lights into the design choices.

However, a major weakness of the paper is its organization. I fail to see strong connections between the proposed Avg-K method, which is superior to all other competitors in experiments, and the unified view. Consequently, section 6 seems to be rather a disjointed part of the paper, though possessing the best results. Having more "bridge explanations" between section 5 and section 6 can ease this feeling of abrupt change. Also, the choice of averaging operation is rather arbitrary, and it is desirable to have proper ablations against alternatives like "maximum dot-product of the block" (the classic max-pooling vs. average-pooling). Also, 
clarity,_quality,_novelty_and_reproducibility: Overall the paper is well organized and well written, except for some structure issues I raise above.

The scale of experiments (computation measured in ZFLOPs) seems to be out of reach of ordinary labs, hampering the reproducibility. It would be greatly appreciated if some experiments of smaller scale settings are included.

Nit: "ml" in "and ml = mi·dm+j"  should be m_l
summary_of_the_review: The paper addresses an important topic and demonstrates powerful results. However, the disjointedness structure makes one wonder how it can be reorganized to improve coherency. Overall I would recommend accepting the paper but would like to hear from peer reviewers on this.

---------- Update after a discussion -----
I would lower my score to 3, after a discussion with peer reviewers organized by AC. 
The reasons are:
* We all acknowledge to some extent that the "Unified view" part is quite disjoint from the Avg-K part. In fact, the "Unified View" part is more about figuring out the relationship the less cited PKM approach with the more cited sparse Mixture of Expert approach. As avg-k is directly based on this unification, it's not clear how this unification can benefit the practice of scaling up neural networks.
* In our discussion, we find another work that has better perplexity result than this paper. Even without that paper as reference, the difference of perplexity in Table 4 seems incremental.
* Withdrawing of "interpretability claim", and scarcity of details helping reproducing results, weakens the contribution of the paper.
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
empirical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
flag_for_ethics_review: NO.
recommendation: 5: marginally below the acceptance threshold

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_6RhL
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary_of_the_paper: This paper proposes a unified view of the sparse mixture-of-experts model and sparse neural memory model. Based on this unified perspective, the authors proposed a new routing method termed Avg-K, which outperforms standard routing methods in a sparse mixture of expert models. 

strength_and_weaknesses: # Strength
The proposed unified view (Section 3) is novel and revealing. The newly formulated Avg-K yields strong empirical performance while being simpler -- it does not require auxiliary losses as in standard sparse MoE.

# Weakness

- Throughout the paper, the authors use ZFLOPs as an indicator of computational complexity. Most methods considered have similar ZFLOPs. Could the authors report train time/inference time as well, e.g., in GPU hours, perhaps in a representative subset of experiments? In my experience, models with similar FLOPS may have drastically different run times.

- It was mentioned that no load balancing loss was used in the Avg-K router. This, empirically, underuses certain memory blocks (as demonstrated in Figure 7 in the appendix) and undermines training efficiency. Could the authors explain why they decided to not add any load-balancing loss? Would the authors expect standard load balancing losses such as those in [1] and [2] to work in the Avg-K router?

- The authors compare the Avg-K approach with several existing sparse MoE models such as the Switch transformer. Unlike Avg-K, these sparse MoE models do have load-balancing losses. Were these load-balancing losses used in the baseline comparison? On page 8, the authors mentioned "Without load balancing constraints, SMoE with learned gate often degenerates to using only a handful of memory blocks", so I am wondering if the load balancing constraints were turned off in certain SMoE experiments.

- In Section 6, the authors mentioned that Avg-K is more interpretable than the Hash layer [3]. While this may be true for random hashing, it should be noted that it is also possible to perform deterministic hashing, which is fully interpretable and typically yields slightly better results than random hashing [3]. I personally do not find interpretability is a strength of Avg-K, since the key table is end-to-end learnable and does not conform to any prior knowledge or human intuition.

[1] Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." ICLR 2017
[2] Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." JMLR, 2021.
[3] Roller, Stephen, Sainbayar Sukhbaatar, and Jason Weston. "Hash layers for large sparse models." Advances in Neural Information Processing Systems 34 (2021): 17555-17566.



clarity,_quality,_novelty_and_reproducibility: The paper is well-presented overall. To further improve the clarity, I list a few suggestions from my side -- it is entirely up to the author to decide whether they are relevant.
- The subsection "Product Key Memory" in Section 2.2.2 does not seem to be necessary. It is very dense in technical details, while not being used anywhere later in the paper. A few sentences to summarize the intuition behind Product Key Memory may suffice.
- The transition from Equation (4) to Equation (5) is confusing. Would be great to have a figure to illustrate how the re-indexing works. For example, the authors can draw two versions of column-wise concatenated value tables: one table, with the original $(i,j)$ index and another with the $l$ index.
- A typo in the line below Equation (5): $ml = ...$ should be $m_l = ...$

I suggest the authors release their code for reproducibility.
summary_of_the_review: The paper provides a new perspective in viewing sparse feedforward layers. In my opinion, the link between the sparse mixture-of-experts layer and the sparse neural memory layer is not something profound -- it is based on the linearity of MoE gating and the final layer of a neural network. Nonetheless, it is good to see that this link leads to actionable improvement in sparse MoE routers. 
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
empirical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
flag_for_ethics_review: NO.
recommendation: 5: marginally below the acceptance threshold

From: ICLR.cc/2023/Conference/Paper5150/Reviewer_o8N9
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary_of_the_paper: The authors provide a unified view towards Mixture-of-Expert models and sparse neural memory, by rewriting the MoE routing as key/value matrix multiplication (plus activation functions). With this correspondence established, the paper then proposes to use block-sized key (weights of FFN before activation) and block sized value (weight matrix after activation) to study the sparse models.

The authors then performed experiments on what block size is optimal, and whether one needs to use direct routing or indirect routing. After that, the authors then propose a new routing method called "Avg-K" and show that it outperforms existing sparse models.

The main contributions of the paper is the study of sparse model parameters and routing mechanisms.
strength_and_weaknesses: The strength of the paper are as follows:
1. From my perspective it is great to study existing sparse models from a unified view; it is important to show that many existing sparse architecture design is in fact special cases of a general design. The idea itself is one of the major strength of the paper.
2. The Avg-K routing is new idea.

There are several weakness of the paper:
1. The primary goal of this paper is unclear to me. Is the paper trying mainly to study the unification, or is the paper trying mainly to propose  the selection method? These two goals would result in very different flavors -- a "unified" view would need first to establish theoretical and empirical equivalence between two architectures; while the "avg-K" selection can be focused on sparse neural memory only.
2. For the unification point of view, the study is limited: when claiming MoE is basically equivalent to sparse neural memory, the author write down some simple derivation, without taking into account other important factor such as load balancing, etc. Theoretical derivation is not entirely indication on practical equivalence. It is better for the authors to empirically show that MoE performs similarly compared to sparse neural memory with indirect routing. 
3. The experiments themselves are limited: for example, can the author explain why they do not include LoRKM in their experiment results, but only use PKM's results? Are there any other selection mechanisms for sparse neural memory? Is the paper [Lample et al, 2019] the only paper proposing selection methods?
4. Finally, many experiment results have large variances. The comparison (in table 3) between direct (VanillaM) and indirect (PKM-FFN) their difference in mean is 8x smaller than their own standard deviation. It is statistically insignificant to claim one method is better than the other. The author also gives the conclusion "direct" selection method is better. But this conclusion is rather weak because the authors did not even try many other indirect selection methods beside PKM.
clarity,_quality,_novelty_and_reproducibility: The structure of the paper is in general easy to follow and retain good clarity of the topic. However, the authors support their claim with limited (and noisy) results, and draw conclusions without a clear picture to quantify them. The average-K selection method is original. 
summary_of_the_review: My recommendation of this paper is rejection. The major issues, as listed above, contain several aspects:
1) The equivalence of MoE and sparse neural memory lacks empirical support, nor did the authors show experimental results of MoE itself to demonstrate the equivalence (or as a benchmark compared with other models, e.g. PKM or vanillaM).
2) The authors draw conclusions with noisy metric.
3) The goal of the paper is also unclear. Does the author mainly try to convey the equivalence of sparse models, or does the author try to study a new selection method for sparse neural memory? These two are separate directions, while the title of the paper indicates the former, the content after 3.2 indicates the latter. In my opinion it is hard to disentangle what is the main focus of the paper.
correctness: 3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.
technical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
empirical_novelty_and_significance: 2: The contributions are only marginally significant or novel.
flag_for_ethics_review: NO.
details_of_ethics_concerns: N/A
recommendation: 5: marginally below the acceptance threshold