SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anh Nguyen* Nikos Karampatziakis* Weizhu Chen
Microsoft Azure AI
Most language models (LMs) are trained and applied in an autoregressive left-to-right fashion, predicting the next token from the preceding ones. However, this ignores that the full sequence is available during training. In this paper, we introduce "Meet in the Middle" (MIM) a new pre-training paradigm that improves data efficiency by training in two directions, left-to-right and right-to-left, and encouraging the respective models to agree on their token distribution for each position. While the primary outcome is an improved left-to-right LM, we also obtain secondary benefits in the infilling task. There, we leverage the two pre-trained directions to propose an infilling procedure that builds the completion simultaneously from both sides. We conduct extensive experiments on both programming and natural languages and show that MIM significantly surpasses existing pre-training paradigms, in left-to-right generation as well as infilling.
Language models (LMs) are becoming essential for various assisted authoring tasks, including text summarization and code completion. In order to be usable in many different applications, most LMs must be able to generate the next token from the sequence of previous tokens. Thus, pre-training has focused on optimizing the model's ability to predict the next token given the previous tokens (the prefix), as measured by perplexity. However, this neglects that at pre-training time we also have access to the subsequent tokens (the suffix). While the suffix cannot be used as an input to the model, there are other ways to incorporate it into pre-training which have not received much attention in the literature. In this work we propose a more data-efficient approach, preserving the autoregressive nature of LMs while fully utilizing pre-training data.
Why should we be interested in alternative pre-training procedures? After all, the main artifact produced during pre-training is an autoregressive left-to-right LM and the pre-training objective closely matches both how the LM is applied and the desire to predict the next token well. Nonetheless, there are still compelling reasons to consider alternative training objectives. First, data efficiency: current LMs are trained with sparse supervision, only learning from the actual next token. By exploring denser supervision during training, we aim to improve performance with less data. Second, related tasks: in real-world scenarios, users may need to infill or modify existing text. Current left-to-right LMs are limited in this regard, as they cannot use both sides of the insertion position, often leading to suboptimal results.
In this work, we introduce a novel pre-training paradigm named "Meet in the Middle" (MIM). The MIM approach is founded on two central principles. Primarily, it introduces an additional right-to-left LM to complement the standard left-to-right model, enabling both models to co-regularize and benefit from each other's context. This improves data efficiency and consistency, with models metaphorically meeting in the middle by adjusting their output probabilities to agree with one another. Secondarily,
$^{2}$ Code and models available at https://github.com/microsoft/Meet-in-the-Middle
MIM enables an efficient infilling inference procedure that capitalizes on the agreement between the two LMs developed during the pre-training phase. For infilling, the models literally meet in the middle, constructing the completion from both ends. This aspect of MIM, while secondary to the pre-training objective, offers crucial benefits. Co-regularization not only reinforces consistency between the two LMs but also aids in early termination of the generation process during infilling.
To train MIM, we employ two decoding flows within a single shared decoder-only architecture  $\mathrm{[BMR^{+}20]}$ ,  $\mathrm{[CND^{+}22]}$ . The forward and backward LMs generate tokens in opposite directions, predicting the next token and previous token, respectively. We pre-train both models jointly on a large text corpus using a combination of standard LM loss and an agreement regularizer. Upon completion, the forward model serves as a drop-in replacement for existing autoregressive LMs, while the backward model can be discarded or utilized for tasks such as infilling.
In our experiments, we assess MIM's effectiveness for pre-training LMs across various domains and tasks. Using public code and language data, we pre-train LMs of different sizes and evaluate their performance based on perplexity and code completion tasks. We compare MIM with FIM  $\left[\mathrm{BJT}^{+}22\right]$  and other baselines, demonstrating its superior performance in terms of perplexity and task-specific evaluation metrics. Furthermore, we conduct ablation studies to validate the effectiveness of our primary proposals during training and inference. In summary, our main contributions include:
- Introducing a novel pre-training paradigm for LMs that efficiently utilizes training data by incorporating both prefix and suffix context while retaining the autoregressive nature of LMs. This is achieved by training forward and backward models and promoting their agreement.
- Proposing a simple, efficient infilling inference procedure that leverages context from both sides and the tendency of forward and backward models to agree. Our approach employs parallelism more effectively than existing procedures, resulting in superior quality and latency compared to the state of the art.
- Pre-training LMs of various sizes on public code and language data using MIM, evaluating their performance on both human and programming languages, and demonstrating that MIM outperforms numerous baselines in standard evaluation metrics. Additionally, we make some models and code publicly available.
We introduce notation used throughout the paper. For a token sequence  $x_{1}, x_{2}, \ldots, x_{N}$ , we denote  $x_{<i}$  as the prefix  $x_{1}, x_{2}, \ldots, x_{i-1}$  and  $x_{>i}$  as the suffix  $x_{i+1}, x_{i+2}, \ldots, x_{N}$ . Definitions for  $x_{\leq i}$  and  $x_{\geq i}$  are analogous. We suppress model dependencies on learnable parameters. Arrows distinguish the two models and their outputs; for example,  $\overrightarrow{p}$  represents the forward model and  $\overleftarrow{p}$  the backward model, while  $\overrightarrow{H}$  and  $\overleftarrow{H}$  are hidden representations from the respective models.
2.1 The Infilling task
Given a token sequence  $x_{1}, x_{2}, \ldots, x_{N}$ , an insertion position  $i$ , the infilling task involves generating a plausible  $M$ -token sequence  $y_{1}, \ldots, y_{M}$  to fill the gap between  $x_{\leq i}$  and  $x_{\geq i + 1}$ . In real-world applications,  $M$  is unknown. The infilling task is harder than the left-to-right generation task as the latter is a special case of infilling with an empty suffix. Since generating the best  $M$ -token left-to-right generation is a search problem that is tackled by heuristics such as greedy search or beam search, we will also follow heuristic approaches for the strictly harder infilling problem.
A technique for infilling that utilizes context from both sides is "Fill in the Middle" (FIM) [BJT+22]. FIM forms the context by concatenating the suffix and prefix (in that order), ensuring coherence near the completion point. FIM is easily applied to pre-trained LMs with minimal modification and is computationally efficient. However, FIM's drawbacks include training on unnatural concatenations of prefix and suffix, giving the prefix more influence on the completion over the suffix, and an ad-hoc training procedure that only considers a few of the  $O(N^2)$  possible splits in an  $N$ -token document.
2.2 Bidirectional Language Modeling
Bidirectional language modeling has been used in the literature in two different ways: The first is to describe non-autoregressive LMs with objectives like Masked Language Modeling, which produce better representations but face difficulties in in-context learning  $\mathrm{[PLR^{+}22]}$ .
Earlier work [HR18,  $\mathrm{PNI^{+}18}$ ] used the term bidirectional language modeling to describe models with left-to-right and right-to-left autoregressive flows. Our approach falls in this category, remaining autoregressive and using future tokens only for regularization during training.
3 Meet in the Middle
Our proposal involves pre-training two models to predict their next token from different input views (prefix vs. suffix), while maintaining agreement on the probability distributions for these tokens. This provides a dense supervision signal compared to simply predicting the next token. It also encourages the models to "meet in the middle" i.e. reach a compromise prediction informed by different views.
We train two decoder-only LMs, a forward model  $\overrightarrow{p}$  and a backward model  $\overleftarrow{p}$ , with shared parameters.  $\overrightarrow{p}$  predicts next tokens in the forward direction, while  $\overleftarrow{p}$  predicts previous tokens in the backward direction. To improve data efficiency, we use a co-regularization term that encourages the models to agree on the predicted probability distribution over the vocabulary for each token. We use total variation distance to measure disagreement between the models. Define the disagreement on  $x_{i}$
$$
D _ {i, x} ^ {T V} (\overleftarrow {p} | | \overleftarrow {p}) = \frac {1}{2} \sum_ {z \in \mathcal {V}} | \overrightarrow {p} (z | x _ {<   i}) - \overleftarrow {p} (z | x _ {> i}) |
$$
where  $\mathcal{V}$  is the vocabulary. The agreement regularizer is the sum of these terms over all sentences and tokens. It provides the models with a denser supervision signal which improves data efficiency and helps us train a better autoregressive LM. This is because two probability distributions over all tokens are compared while in the traditional LM objective only the probability of the observed token matters. Moreover it encourages the models to agree on their predictions which helps with the efficiency of infilling as we will see later. Then, the full training loss for dataset  $S$  of sequences is given by:
$$
\sum_ {x \in S} \sum_ {i = 1} ^ {| x |} - \log (\overrightarrow {p} (x _ {i} | x _ {<   i})) - \log (\overleftarrow {p} (x _ {i} | x _ {> i})) + \beta \cdot D _ {i, x} ^ {T V} (\overrightarrow {p} | | \overleftarrow {p}). \tag {1}
$$
The hyperparameter  $\beta$  is set to 0.1 in our experiments. After pre-training,  $\overrightarrow{p}$  can be used as a left-to-right autoregressive LM.
3.2 Infilling
3.2.1 Inference
Our inference procedure for infilling aims to be efficient and low-latency. Our approach is shown in Figure 1. At a high level, the two models start building a completion each from their own side and they try to literally meet in the middle. Initially, the prefix and suffix are consumed by  $\overrightarrow{p}$  and  $\overleftarrow{p}$  respectively. Then  $\overrightarrow{p}$  and  $\overleftarrow{p}$  generate tokens synchronously one at a time. For each generated token from  $\overrightarrow{p}$  (resp.  $\overleftarrow{p}$ ) we check whether it matches a generated token from  $\overleftarrow{p}$  (resp. from  $\overrightarrow{p}$ ). If there is a match, we have a "meet in the middle" candidate position for joining the two generated sequences.
In an ideal scenario,  $\overleftarrow{p}$  and  $\overleftarrow{p}$  generate identical sequences, albeit in reverse order, resulting in "meeting in the middle" after half of the tokens are generated by each model. The agreement regularizer is thus crucial; if  $\overrightarrow{p}$  and  $\overleftarrow{p}$  yield disparate sequences, our method's speed is inferior to FIM since early termination isn't possible. However, if there's complete agreement, with sufficient parallelism, our method can outperform FIM in generation latency, as each model only needs to generate half the completion autoregressively. Our experiments demonstrate our method's superior latency to FIM, indicating that in most instances, sequences meet close to the middle.
Figure 1: Inference procedure for infilling. Given the prefix "The quick" and the suffix "dog", the models  $\overleftarrow{p}$  and  $\overleftarrow{p}$  generate tokens until a candidate meet-in-the-middle token is detected (shaded area). We use a single token for illustration purposes, although the method can use more tokens. Given the candidate MIM token, the  $\overrightarrow{p}$  (respectively  $\overleftarrow{p}$ ) model can in parallel verify that the tokens generated by  $\overleftarrow{p}$  (resp.  $\overrightarrow{p}$ ) are acceptable completions (only the  $\overrightarrow{p}$  verification is shown to reduce clutter). The numbers in the top right of each box show the order of operations. Two boxes with the same number can be executed in parallel. Similarly, thick lines show that information (embeddings, tokens) can flow in parallel, while thin lines denote sequential steps.
However, the procedure may yield false positives. For instance, matching a token such as "the" from both sides may not guarantee a coherent infilling upon joining. To mitigate this, we employ 4-gram matching in our infilling experiments. Additionally, we run a parallel verification procedure to ensure that the resultant sequence could have been generated by either  $\overrightarrow{p}$  or  $\overleftarrow{p}$  running to completion.
Our parallel verification procedure is adapted from  $\mathrm{[GX S^{+}22]}$ . Concretely, assume that the latest token generated from  $\overrightarrow{p}$  matches with a token in position  $s$  generated from  $\overleftarrow{p}$ , as is the case in Figure 1 for the token "over". Then we can use the tokens that were generated by  $\overleftarrow{p}$  before  $s$  as inputs to  $\overrightarrow{p}$  in parallel. In the context of Figure 1 we provide the tokens "over", "the", "lazy" as inputs to  $\overrightarrow{p}$  in parallel. If each output of  $\overrightarrow{p}$  matches the corresponding input, we have verified that  $\overrightarrow{p}$  would have autoregressively generated the tokens it copied from  $\overleftarrow{p}$ 's output. Grounding the discussion back to Figure 1, if  $\overrightarrow{p}$  generates (in parallel) "the" and "lazy" we have verified that  $\overrightarrow{p}$  would have autoregressively generated the same tokens. In cases of partial or no match, the generation from  $\overrightarrow{p}$  either fast-forwards to the first disagreement point or resumes autoregressive generation from  $\overrightarrow{p}$  and  $\overleftarrow{p}$ . While our implementation uses greedy top-1 sampling and comparison of input and output at each position, we could adopt less stringent acceptance criteria, like a sufficiently high probability of the candidate token in the previous output. If  $\overrightarrow{p}$  and  $\overleftarrow{p}$  finish generating without meeting or passing verification, we return the higher-probability sequence according to its generating model.
3.2.2 Infilling Specific Enhancements
Our model's infilling performance can be boosted by adopting the Synchronous Bidirectional Attention mechanism [ZZZ19]. Though this reduces compatibility with autoregressive language models (LMs), it enhances generation through bidirectional conditioning. The Synchronous Bidirectional Attention layer modifies the regular attention layer activations as follows: Let  $\vec{H}$  and  $\vec{H}$  be the outputs of an attention layer in  $\vec{p}$  and  $\vec{p}$  respectively. For a hyperparameter  $\lambda$ , the fused attention hidden representation becomes  $\vec{H}_f = \vec{H} + \lambda \vec{H}$  and  $\vec{H}_f = \vec{H} + \lambda \vec{H}$ . Instead of  $\vec{H}$  and  $\vec{H}$ , the models  $\vec{p}$  and  $\vec{p}$  use  $\vec{H}_f$  and  $\vec{H}_f$ . If  $\lambda = 0$ ,  $\vec{p}$  and  $\vec{p}$  become classic autoregressive transformers. However, careful training is necessary to avoid information leakage with this layer. To mitigate this, we adopt a two-stage training process, see [ZZZ19] for details.
4.1 Data and Models
Our code models are pre-trained on a large and diverse corpus of public code with permissive licenses, covering multiple programming languages. The dominant languages in our corpus are Python, Java, and C++. After filtering and dedduplication, the corpus contains about 300B tokens. Details are provided in Table 9. Our models train on a single pass over this data. This is about six times larger than the dataset used to train the Incoder model  $\mathrm{[FAL^{+}22]}$ , which included 50B tokens comprising code and Stack Overflow data.
We train our natural language models on data from the following datasets: CC-News, OpenWebText, CC-Stories, and CC-100 to assess the models language modeling capability. In particular:
- CC-News: 63M English news articles crawled from Sept 2016 to Feb 2019 (76GB).
- OpenWebText: An open source recreation of the GPT-2 WebText dataset (38GB).
- CC-Stories: A subset of CommonCrawl data filtered to match the story-like style of Winograd schemas (31GB).
- CC-100: A dataset extracted from CommonCrawl snapshots between Jan 2018 and Dec 2018, filtered to match the style of Wikipedia (292GB).
Apart from the natural language data utilized in our language modeling experiments, we also train our models on a filtered version of the Falcon RefinedWeb corpus  $\mathrm{[PMH^{+}23]}$  used in the previous work of  $\mathrm{[LBE^{+}23]}$ , that has a total of 88B tokens. We then present the performance results on common natural language benchmarks in section 4.3.3
We use a decoder-only transformer LM  $\left[\mathrm{VSP}^{+}17, \mathrm{BMR}^{+}20\right]$  for pre-training. The model is trained to predict the next token in both left-to-right and right-to-left directions, using the same model and parameters. To explore the impact of model size, we pre-train models with different capacities: 350M, 1.3B, 2.7B. For the hyperparameters and training setup of each model size, please refer to the Appendix.
As baselines, we use well-known published results but also pre-train three FIM models with context-level FIM  $\left[\mathrm{BJT}^{+}22\right]$ . Character-level transformations are applied, with a FIM-rate of 0.5 and SPM+PSM joint training. Both the MIM and FIM models are pre-trained using the Megatron-LM framework  $\left[\mathrm{SPP}^{+}19\right]$ . See the Appendix for model configurations and architectural details.
4.2 Benchmarks and metrics
4.2.1 Code generation and infilling
We assess MIM's performance both in autoregressive left-to-right generation and infilling tasks. For the autoregressive generation, the model generates the code body from a function signature, docstring, and test cases. Evaluation datasets include HumanEval [CTJ+21], MBPP [AON+21], and APPS [HBK+21] for Python problems. We use the multilingual HumanEval-X dataset for multiple programming languages. For this task, the left-to-right model alone is used for inference without any provided suffix.
In the infilling task, the model fills blank lines in an incomplete program. This is evaluated using the HumanEval Infilling [BJT+22] and MBXP [AGW+22] datasets. MIM is compared with Incoder, three FIM variants we trained on the same data as MIM, and code-davinci-002, as these datasets provide suffixes that not all models can manage effectively.
We use the pass@k metrics  $\mathrm{[CTJ^{+}21]}$  to measure the percentage of generated codes that pass all tests within the top-k candidates. We report pass@1, pass@10, and pass@100. For generation, top-p sampling with  $p = 0.95$  is employed with temperatures of 0.2 for pass@1 and 0.8 for pass@10 and pass@100. We also include single-line exact match (EM) metrics to indicate the percentage of completed lines matching the reference solution's masked lines  $\mathrm{[FAL^{+}22,BJT^{+}22]}$ .
4.2.2 Language Modeling
In addition to code generation and infilling, we evaluate MIM's language modeling capabilities, particularly its next-token prediction ability measured by perplexity. Perplexity is evaluated in both in-domain and out-of-domain settings. The former involves a held-out subset of the combined training data, while the latter uses the Pile dataset  $\mathrm{[GBB^{+}21]}$ , a multi-domain public language modeling dataset. Average perplexity across all Pile subsets is reported for our models and baselines.
4.2.3 Natural Language Benchmarks
Finally, we evaluate MIM on established benchmarks for natural language, encompassing both commonsense reasoning and language understanding. For commonsense reasoning, we have tested our models on Winogrande, ARC-Easy, ARC-Challenge, BoolQ, and SIQA. In terms of knowledge and language understanding, we have assessed MIM's performance on PIQA, HellaSwag, and OpenbookQA. We compare against phi-1.5-web-only  $\left[\mathrm{LBE}^{+}23\right]$  and present zero-shot results for these tasks using the LM-Eval Harness, with the exception of MMLU where we use 2-shot performance.
4.3 Main results
4.3.1 Code generation and infilling
First we focus on the superior performance of our from-scratch pre-trained FIM model over various baselines, in Table 1. For instance, our 2.7B model yields a  $28.5\%$  pass@1 in HumanEval, significantly outperforming the  $15.2\%$  pass@1 of Incoder's 6.7B model [BJT $^{+}$ 22]. FIM-2.7B also beats other strong baselines such as Codex 2.5B and CodeGen-Multi-6.1B, achieving  $67.8\%$  in pass@100 compared to their respective  $44.9\%$  and  $59.5\%$ . This performance is attributable to larger, cleaner training data and implementation details (discussed in the Appendix).
Having established a strong baseline, we now benchmark MIM's autoregressive generation in a prefix-only setting using both the HumanEval and MBPP datasets. While FIM vs. MIM is the apples-to-apples comparison, we include many other models for reference and note that all these could have leveraged our MIM objective during their training. Table 1 shows MIM consistently outperforming FIM across all metrics for all model sizes (350M, 1.3B, 2.7B). For example, MIM-2.7B enhances HumanEval pass@1 to  $30.7\%$ , a  $2.2\%$  gain over FIM-2.7B. The improvement on MBPP is similar, in terms of pass@1. We note that the improvements increase with the model size. This is evidence that the larger models benefit more from the agreement regularizer. We further evaluate this setting on the APPS dataset, which has three different difficulty levels. We report the results in Table 3. The trend and improvement are similar to the other two datasets, in which MIM consistently outperforms FIM across all the metrics and difficulty levels. In the multilingual setting, we compare MIM with FIM baselines in the HumanEval-X dataset. We observe consistent improvement across all metrics and all the programming languages that we evaluated on. Results are reported in Table 2.
These findings illustrate MIM's superiority over FIM for left-to-right autoregressive generation. The main claim in  $\left[\mathrm{BJT}^{+}22\right]$  is that FIM does not harm the original left-to-right generative capability and can be learned for free. We contend that MIM's pre-training, enhanced by the agreement regularizer, results in a superior generative model, thereby proposing a new pre-training paradigm.
Finally, we test MIM in an infilling setting, as motivated by real-world applications like GitHub Copilot, where developers frequently edit in the middle of source files. MIM outperforms FIM across all metrics and model sizes on both the HumanEval Infilling and the MBXP Infilling benchmarks (Table 1). Notably, the MIM-2.7B model achieves pass@1 of  $26.3\%$  and an exact match metric of  $57.8\%$ , improvements of  $3.5\%$  and  $6.1\%$ , respectively, over the FIM-2.7B model. This underscores the consistent improvement that MIM offers across different model sizes and tasks.
4.3.2 Language Modeling
Here we compare FIM and MIM models pre-trained on language modeling tasks. We use perplexity as our evaluation metric and look at two different settings, namely in-domain and out-of-domain.
The perplexity results of both settings are summarized in Table 4. Across all experiments, MIM consistently outperform FIM baselines in terms of perplexity. The MIM-2.7B model has the best