We are very grateful for your kind words, as we worked hard to make our writing and math clear, and are very excited about this new perspective and the resulting insights. We hope that we can sufficiently address your concerns.
> Results, efficacy of the method still remains to be demonstrated (minor)
We agree that our pilot experiment results are largely preliminary and will benefit from further investigation in future work. We take the view that our experimental findings are secondary and meant to supplement our primary theoretical results and method proposal.
> Unclear whether the method will help for larger models or for models where the approximation errors (under-estimation / over-estimation) are small.
One important implication of our results is that larger models do not automatically avoid under/over-estimation errors, since the errors caused by the softmax bottleneck (SMB) depend not on the model size, but on the discrepancy between the hidden size and the vocabulary size. A larger model with a proportionally larger vocabulary size will suffer the same bottleneck effects. By the same argument, a model with a small vocabulary (i.e., close to equal hidden and vocab sizes) could avoid the SMB altogether, no matter the size. Unfortunately, we do not see many models of this type in common use.
One size-related investigation we would like to pursue in future work is the hypothesis that our method will become **more** effective for larger model sizes. We hypothesize that some of our results were weaker than expected since our method primarily addresses SMB-related errors. If the model predicts the next token incorrectly, none of the considered truncation methods (including ours) will be able to correct the error. This issue should be reduced with larger models, since larger models should have fewer modeling errors, leaving only SMB errors. We see some anecdotal evidence for this where, for instance, BA-$\eta$ becomes stronger compared to $\eta$ sampling as model size increases (Figure 5).
> Regarding Eq. 3: what I am going to propose is a bit dirty but would it be possible at each step to minimize (W^T p - W^T \hat p)^2, wrt to p with a sparsity constraint (e.g. l1) and the range constraints, and reject all tokens for which |p| = 0? It might not derive from the theory but it might capture the overall idea? just wondering.
This idea is worth exploring! At first glance, the major difference with this approach is that it would require all truncated tokens to have 0 probability simultaneously, whereas our method only requires a solution with each truncated token having 0 probability independently. The advantage of your proposed method would be a faster runtime, with the drawback that the list of truncated tokens might be incomplete. Perhaps there is a middle ground, where we iteratively find solutions, each time finding multiple tokens to truncate, then removing the sparsity constraint from these tokens for the next iteration so we can find new potentially-0-probability tokens.
> Concerning the results: there isn't much of a pattern in the MAUVE results if I look at the improvements of BA across model scales. Isn't BA expected to help more with smaller scales given that the approximation error might be bigger?
See our response to your concern above: there may be multiple effects here, such as the rate at which the model makes prediction mistakes independent of the SMB, which would confound the trend because these types of mistakes would become less frequent at larger model sizes.
> What happens with bigger models? given that the approximation error will be smaller, would your method still help?
Again, we refer to our size-related responses above, with an additional perspective: we expect that small models will continue to play a role in language generation, given the ease of deployment in edge devices, as well as cheaper inference costs. Even if our method does not provide benefits for larger models, smaller models will continue to suffer from a SMB and our method (or some modification of it) will remain relevant.
> Nitpicks: It might be clearer to re-introduce the \epsilon and \eta baselines in the experiments. I struggled a bit to remember given that they are just introduced in the background section. Eq. 10 in the appendix is missing a parenthesis. Can you include a pseudo-code of the final BA implementation in the main paper?
Thank you for pointing these out, we update accordingly in our revision.