InsNeXt: Training Scalable Insertion-based Language Models from Scratch

We thank the reviewer for the thorough reading and comprehensive feedback. We hereby try to respond to each point of them one by one:

• Re Weakness 1) "Representation quality falling short against BERT-style works" We acknowledge that the proposed InsNeXt model falls short against encoder-only models like RoBERTa-large on NLU tasks. We hypothesize that this is due to multiple reasons. First, as all existing generative models including the more recent ones like Qwen falling short against RoBERTa, we actually outperform all autoregressive LMs. We argue that there might be some inherent limitations of the representation produced by generative models, possibly due to how information shall be differently compressed for generation/understanding. Second, RoBERTa is trained for 40Epochs on 160GB text data, which is approximately 1720B tokens, whereas the total number of trained tokens in InsNeXt can barely reach 180B tokens. In fact, we still observe steady performance gains in InsNeXt with continued training (either with more epochs on the same data we already sampled or sampling more from SlimPajama). We are unable to proceed due to the limitation of our computational resources. We respectfully ask for the reviewer’s understanding.

• Re Weakness 2) and Question 1) "Direct comparison between position encodings, Training and inference efficiency" We respectfully ask the reviewer to refer to Table 4 in our appendix for a direct comparison between position encodings in representation learning performance and computational efficiency. In fact, due to the huge scalability differences between the possible choices, all other position encodings than the proposed insertion-oriented ALiBi are practically too slow to be trained on the longer sequences from SlimPajama.

  For inference efficiency, we hereby try to both theoretically and practically analyze the performance w/ the proposed Efficient Bidirectional Recontextualization. As is cited in our paper, we refer to [1] for the proof that, the amortised cost of each insertion operation w/ EBR is O(1 + T(n) / n), where T(n) is the cost for the transformer model to bidirectionally encode an n-length sequence. According to reports from FlashAttention [2, Figure 3], with modern computational infrastructures and efficient attention kernels, T(n) is empirically sublinear to at most linear w.r.t. n when n < 8K. This reflects the fact that the major bottleneck in this case is the tensor copying overhead, which matches the original assumptions in C++ Vector container’s proof. To effectively generalize this assumption to longer context, in practice we can and should always partially bidirectionally re-contextualize within the most recent 8K token window. The following is a more self-contained proof for EBR’s amortised per-insertion cost being O(1):

Assumption: 1) Each newly processed token causes a uniformly distributed copying cost $\eta$; 2) Within the concerned range of context length, the cost for the transformer model to bidirectionally encode an n-length sequence T(n) is O(n);

Proof. Without loss of generality, we consider a growing sequence from 1 to $n = 2^m$ tokens.

The total number of copying cost induced by these $n$ insertions is $O(n\cdot\eta)=O(n)$.

Amongst these insertions, at the steps $i = 2, 4, 8, …, 2^m$, the spiking bidirectional recontextualization causes a cost of $O(2 + 4 + 8 +... +2^m) = O(2^{m+1}-2) = O(2n) = O(n)$

The total overhead cost is thus O(n). Divided by the number of steps $n$, we get O(1) as the amortised cost per-step.

It is trivial to prove the general case, since a growing sequence ending with a non-spiking step simply adds linear copying cost which does not violate the O(1) complexity.

Empirically, in our experiments where the generation procedures do not exceed 128 tokens, we do not observe a statistically significant overhead from the proposed EBR (<30ms/per step). A single decoding step with QK cache for InsNeXt also resembles that of a plain autoregressive language model.

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2022. Dynamic tables, 4 edition, chapter 17. MIT Press, Cambridge, MA. Section 17.4.

[2] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135. Spotlight paper at ICLR 2023.

We appreciate the reviewer’s many constructive comments and feedback, and look forward to further discussion. We respectfully ask the reviewer if you can further acknowledge our work, should the concerns be addressed here or in our extended conversation.

We thank the reviewer for feedback. We hereby try to respond to each point of them one by one:

• Re Weakness 1) and Question 2) ”Representation quality falling short against BERT-style works” We acknowledge that the proposed InsNeXt model falls short against encoder-only models like RoBERTa-large on NLU tasks. We hypothesize that this is due to multiple reasons. First, as all existing generative models including the more recent ones like Qwen falling short against RoBERTa, we actually outperform all autoregressive LMs. We argue that there might be some inherent limitations of the representation produced by generative models, possibly due to how information shall be differently compressed for generation/understanding. Second, RoBERTa is trained for 40Epochs on 160GB text data, which is approximately 1720B tokens, whereas the total number of trained tokens in InsNeXt can barely reach 180B tokens. In fact, we still observe steady performance gains in InsNeXt with continued training (either with more epochs on the same data we already sampled or sampling more from SlimPajama). We are unable to proceed due to the limitation of our computational resources. We respectfully ask for the reviewer’s understanding.

• Re Weakness 2) and Question 1) ”Inference efficiency concerns with Efficient Bidirectional Recontextualization” We hereby try to both theoretically and practically analyze the performance w/ the proposed Efficient Bidirectional Recontextualization. As is cited in our paper, we refer to [1] for the proof that, the amortised cost of each insertion operation w/ EBR is O(1 + T(n) / n), where T(n) is the cost for the transformer model to bidirectionally encode an n-length sequence. According to reports from FlashAttention [2, Figure 3], with modern computational infrastructures and efficient attention kernels, T(n) is empirically sublinear to at most linear w.r.t. n when n < 8K. This reflects the fact that the major bottleneck in this case is the tensor copying overhead, which matches the original assumptions in C++ Vector container’s proof. To effectively generalize this assumption to longer context, in practice we can and should always partially bidirectionally re-contextualize within the most recent 8K token window. The following is a more self-contained proof for EBR’s amortised per-insertion cost being O(1):

Assumption: 1) Each newly processed token causes a uniformly distributed copying cost $\eta$; 2) Within the concerned range of context length, the cost for the transformer model to bidirectionally encode an n-length sequence T(n) is O(n);

Proof. Without loss of generality, we consider a growing sequence from 1 to $n = 2^m$ tokens.

The total number of copying cost induced by these $n$ insertions is $O(n\cdot\eta)=O(n)$.

Amongst these insertions, at the steps $i = 2, 4, 8, …, 2^m$, the spiking bidirectional recontextualization causes a cost of $O(2 + 4 + 8 +... +2^m) = O(2^{m+1}-2) = O(2n) = O(n)$

The total overhead cost is thus O(n). Divided by the number of steps $n$, we get O(1) as the amortised cost per-step.

It is trivial to prove the general case, since a growing sequence ending with a non-spiking step simply adds linear copying cost which does not violate the O(1) complexity.

Empirically, in our experiments where the generation procedures do not exceed 128 tokens, we do not observe a statistically significant overhead from the proposed EBR. On a single RTX4090 GPU hosted on WSL 2, Ubuntu 22.04 and the huggingface transformers library, with our 0.6B model, a trivial step that does not trigger the recontextualization takes approximately 13.3ms on average. The spiking step that triggers it takes approximately 21.3ms, including all time cost to update the recontextualized QK caches. For longer context, this additional overhead can always be contained in the current magnitude by using a local window, such that we only re-contextualize the nearest, say 128 or 1024 tokens. The user also get to customize if they want to resize the window or trigger the recontextualization less/more often for better efficiency/quality.

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2022. Dynamic tables, 4 edition, chapter 17. MIT Press, Cambridge, MA. Section 17.4.

[2] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flash Attention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135. Spotlight paper at ICLR 2023.

We appreciate the reviewer’s suggestions, and look forward to further discussion. We respectfully ask the reviewer if you can re-evaluate our work, should the major concerns be addressed here or in our extended conversation.

We thank the reviewer for the suggestions. We hereby try to respond to each point of them one by one:

• Re Weakness 1a) ”Limited novelty” Just to clarify, the main contribution of this paper focuses on tackling the training-time efficiency problem of insertion-based language models. While we acknowledge that InsNet has addressed many major efficiency issues of insertion-based language models, it only considers the case where training sequences are short, thus no efficient attention mechanism is needed. The original sinusoidal position encoding in InsNet falls obsolete and is incompatible with more recently studied efficient attention kernels like FlashAttn and FlexAttn. We’ve tackled these problems with InsNeXt. To the best of our knowledge, we are the first to scale up fully insertion-based language models. By showing the community some promising results on scaling up insertion-based models, we want to attract more attention and resources to further boost the scaling-up of such a new paradigm of language models.

• Re Weakness 1b) and Question 1 ”Proof of correctness of EBR”: As is cited in our paper, we refer to [1] for the proof that, the amortised cost of each insertion operation w/ EBR is O(1 + T(n) / n), where T(n) is the cost for the transformer model to bidirectionally encode an n-length sequence. According to reports from FlashAttention [2, Figure 3], with modern computational infrastructures and efficient attention kernels, T(n) is empirically sublinear to at most linear w.r.t. n when n < 8K. This reflects the fact that the major bottleneck in this case is the tensor copying overhead, which matches the original assumptions in C++ Vector container’s proof. To effectively generalize this assumption to longer context, in practice we can and should always partially bidirectionally re-contextualize within the most recent 8K token window. The following is a more self-contained proof for EBR’s amortised per-insertion cost being O(1) that we plan to add to our revision:

Assumption: 1) Each newly processed token causes a uniformly distributed copying cost $\eta$; 2) Within the concerned range of context length, the cost for the transformer model to bidirectionally encode an n-length sequence T(n) is O(n);

Proof. Without loss of generality, we consider a growing sequence from 1 to $n = 2^m$ tokens.

The total number of copying cost induced by these $n$ insertions is $O(n\cdot\eta)=O(n)$.

Amongst these insertions, at the steps $i = 2, 4, 8, …, 2^m$, the spiking bidirectional recontextualization causes a cost of $O(2 + 4 + 8 +... +2^m) = O(2^{m+1}-2) = O(2n) = O(n)$

The total overhead cost is thus O(n). Divided by the number of steps $n$, we get O(1) as the amortised cost per-step.

It is trivial to further prove the general case, since a growing sequence ending with a non-spiking step simply adds linear copying cost which does not violate the O(1) complexity.

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2022. Dynamic tables, 4 edition, chapter 17. MIT Press, Cambridge, MA. Section 17.4.

[2] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135. Spotlight paper at ICLR 2023.

• Re Weakness 2) ”Writing Restructuring Suggestions” We thank the reviewer for the writing suggestions. We will update the draft accordingly.

• Re Weakness 3) ”Need Training Efficiency Proof” We respectfully ask the reviewer to refer to Table 4 in our appendix. In this table, we show that by adopting the proposed insertion-oriented ALiBi, the insertion-based language model can be trained $6.68x$ faster compared to InsNet, resulting in 116K tokens/second/Chip for the $base$-sized model at the context length of 1024 tokens. This advantage can only grow with a longer context/larger model as insertion-oriented ALiBi is the only one practically compatible with FlexAttn (a customizable FlashAttn alternative). Please let us know if you have further concerns.

• Re Weakness 4) ”Code editing experiments” We agree that code editing is one interesting task that insertion-based language models should have a good advantage over traditional autoregressive models. However, in this paper we want to establish the foundation of scalability for fully insertion-based language models, following how autoregressive models like GPT were originally developed. We thank the reviewer again for this suggestion as we will definitely put high priority on this direction in our future work.

We appreciate the reviewer’s many constructive comments and feedback, and look forward to further discussion. We respectfully ask the reviewer if you can re-evaluate our work, should the major concerns be addressed here or in our extended conversation.

”Proof of correctness and efficiency of EBR”: As is cited in our paper, we refer to [1] for the proof that, the amortised cost of each insertion operation w/ EBR is O(1 + T(n) / n), where T(n) is the cost for the transformer model to bidirectionally encode an n-length sequence. According to reports from FlashAttention [2, Figure 3], with modern computational infrastructures and efficient attention kernels, T(n) is empirically sublinear to at most linear w.r.t. n when n < 8K. This reflects the fact that the major bottleneck in this case is the tensor copying overhead, which matches the original assumptions in C++ Vector container’s proof. To effectively generalize this assumption to longer context, in practice we can and should always partially bidirectionally re-contextualize within the most recent 8K token window. The following is a more self-contained proof for EBR’s amortised per-insertion cost being O(1) that we plan to add to our revision:

Assumption: 1) Each newly processed token causes a uniformly distributed copying cost $\eta$; 2) Within the concerned range of context length, the cost for the transformer model to bidirectionally encode an n-length sequence T(n) is O(n);

Proof. Without loss of generality, we consider a growing sequence from 1 to $n = 2^m$ tokens.

The total number of copying cost induced by these $n$ insertions is $O(n\cdot\eta)=O(n)$.

Amongst these insertions, at the steps $i = 2, 4, 8, …, 2^m$, the spiking bidirectional recontextualization causes a cost of $O(2 + 4 + 8 +... +2^m) = O(2^{m+1}-2) = O(2n) = O(n)$

The total overhead cost is thus O(n). Divided by the number of steps $n$, we get O(1) as the amortised cost per-step.

It is trivial to further prove the general case, since a growing sequence ending with a non-spiking step simply adds linear copying cost which does not violate the O(1) complexity.

In comparison, trivially bidirectional contextualization at each step will take the average cost of T(n), which starts with constant value in shorter context, growing into O(n) within in medium length of context and asymptotically into $O(n^2)$ when no effective total parallelization shall be possible. In comparison, EBR is indeed theoretically much more efficient, especially as model size/context length grows large.

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2022. Dynamic tables, 4 edition, chapter 17. MIT Press, Cambridge, MA. Section 17.4.

[2] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135. Spotlight paper at ICLR 2023.

We thank the reviewer for the suggestions on clarity, training setup and robustness analysis. We hereby try to respond to each point of them one by one:

• Re Weakness 1) ”Efficiency concerns” Just to clarify, the main contribution of this paper focuses on tackling the training-time efficiency problem of insertion-based language models. If the reviewer is referring to the inference-time efficiency concerns about the proposed Efficient Bidirectional Re-contextualization algorithm, we acknowledge that it introduces additional overhead compared to plain QK-cached insertion-based/autoregressive decoding. However, for most decoding steps, the proposed decoding does the same thing as QK-cached decoding. Only at the steps where the re-contextualization is triggered (i.e. uni-directionally cached context exceeds the length of bi-directionally encoded ones), an additional spiking computation of bidirectional re-encoding and QK-cache update shall be needed. Note that as is referred in our paper (Line 255, or see [1]), the amortized cost of this proposed algorithm is the same as plain QK-cached decoding. According to the statistics we currently have, there isn’t an observable computational performance overhead from the spiking bidirectional recontextualization in EBR for sequences at the length of around or less than 512 tokens (we don’t have data for longer sequences). Note that we can always make EBR windowed in longer sequences (i.e. only impact the most recent 1024 or less QK caches instead of recontextualizing the entire history) to keep the spiking overhead low. We can add a plot of the computation overhead in the camera-ready version. Thank you again for your insightful comments on this.

[1] Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. 2022. Dynamic tables, 4 edition, chapter 17. MIT Press, Cambridge, MA. Section 17.4.

• Re Weakness 2a) ”Insufficient scalability” While we agree that 0.6B is a limited scale compared to most recently published LLMs, as an academic group with limited resources, in this work, we aim to explore and deliver the best practices to train insertion-based LMs that are at similar computational level of scalability of modern decoder-only autoregressive LLMs. We choose the scale of 0.6B to align with the smallest scale of more recently published open-source LLMs like Qwen2.5-0.5B. By showing the community some promising results on an affordably scaled up model, we want to attract more attention and resources to further boost the scaling-up of such a new paradigm of language models. In our experiments, we’ve shown that at a small-to-medium scale, insertion-based language models are indeed showing promises to be more controllable, capable and performant language models that can facilitate the next generation of AI. To the best of our knowledge, we are the first to scale up fully insertion-based language models to this level.

• Re Weakness 2b) ”Ablation of the two-staged training” We’ve already presented the results of a model that is trained only on sentence-level data in our appendix (please refer to Table 4, the InsNet w/ insertion oriented ALiBi roughly shows the performance of models trained on the sentence-level only data), and we will re-compile a new section that specifically discusses the ablation and necessity of the two-staged training in the camera-ready version, emphasizing its significance in accelerating convergence and stabilizing training given our limited computation. We appreciate the reviewer’s constructive suggestion on this.

• Re Question 1) ”Performance with dense insertion operations and longer context” To clarify, by default InsNeXt always decodes with dense, sequential (which means one token inserted per step) insertion operations that are QK-cached, with or without the previously proposed EBR. If we understand the reviewer correctly, we want to respectfully point out that all current generative results are already produced in a fashion that in a short to medium length of span, fully-dense insertion operations are performed to form a sentence or document. For generation of longer contents, due to the limited scale of the model that we eventually obtain, we admit that it is unlikely that the models can still produce reasonable output without a performance loss. In fact, according to our follow-up experiments, after the model approaches the length limit during pretraining, it tends more and more to insert into and expand previously concluded sentences, rather than moving on to the next sentences. This can be partially alleviated if we only allow the model to insert into the most recent 32 positions window, but still eventually leads to degenerated performance beyond approximately 4K tokens.

• Re Question 2) ”Insertion-induced artifacts” We thank the reviewer for bringing this up! As is stated in our response to Q1, forcibly generating contents with the windowed insertion trick beyond the model’s designed capability can cause incoherence and artifacts, including but not limited to non-commonsensical (illogical subject-action-object events) and non-grammatical (unpaired parenthesis etc.) but less repetitive output. This observation of out-of-domain misbehaviors is interestingly different from that of autoregressive models which in a similar situation usually starts to produce highly repetitive outputs. We will add this to the limitation section, and leave detailed study of this for future work.

We appreciate the reviewer’s many constructive comments and feedback, and look forward to further discussion. We respectfully ask the reviewer if they can re-evaluate our work, should the major concerns be addressed in our extended conversation.