SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alessandro Sordoni $^{ab*}$  Xingdi Yuan $^{a}$  Marc-Alexandre Côté $^{a}$  Matheus Pereira $^{a}$
Adam Trischler<sup>a</sup> Ziang Xiao<sup>a</sup> Arian Hosseini<sup>b</sup> Friederike Niedtner<sup>a</sup> Nicolas Le Roux<sup>ab</sup>
Microsoft Research Montréal<sup>a</sup> MILA<sup>b</sup>
Large language models (LLMs) can be seen as atomic units of computation mapping sequences to a distribution over sequences. Thus, they can be seen as stochastic language layers in a language network, where the learnable parameters are the natural language prompts at each layer. By stacking two such layers and feeding the output of one layer to the next, we obtain a Deep Language Network (DLN). We first show how to effectively perform prompt optimization for a 1-Layer language network (DLN-1). Then, we present an extension that applies to 2-layer DLNs (DLN-2), where two prompts must be learned. The key idea is to consider the output of the first layer as a latent variable, which requires inference, and prompts to be learned as the parameters of the generative distribution. We first test the effectiveness of DLN-1 in multiple reasoning and natural language understanding tasks. Then, we show that DLN-2 can reach higher performance than a single layer, showing promise that we might reach comparable performance to GPT-4, even when each LLM in the network is smaller and less powerful. The DLN code is open source. $^{1}$
The size of large language models (LLMs) has grown significantly over the last few years, mainly because of emerging capabilities [5, 31], but at considerable technical and societal costs [49, 2, 4]. Recent efforts have focused either on learning smaller models matching the abilities of larger ones on some tasks using distillation [43, 36, 29, 13], or offloading part of the computation to other dedicated components [28, 22, 25, 18]. In the latter case, this is done through carefully crafted instructions to retrieve the necessary information from these additional modules [48, 41, 6, 54, 24].
Instruction-tuned LLMs map an input sequence to a distribution over output sequences conditioned on an instruction, or prompt. In this paper, we view such LLMs as stochastic language layers, whose learnable parameters are the prompts. Multiple layers can be stacked to form a Deep Language Network (DLN) whose learnable parameters are the prompts associated to each layer. Specifically, each layer uses a template to organize both its prompt and the inputs coming from the layer below into a single sequence before producing the output (see Figure 1). This layering induces a learnable decomposition of the task into a series of smaller sub-tasks, each of which might be more easily solvable by an LLM. This view shares similarities to recent works that chain LLM calls [41, 6, 54]. In this work, we move towards integrating learnable components in the pipeline: each prompt can be learned to maximize the final objective.
*Corresponding author: alsordon@microsoft.com
$^{1}$ Code: https://github.com/microsoft/deep-language-networks.
Figure 1: Left: An illustration of a DLN-1 performing a sentiment analysis task: input and the trainable prompt are merged using a template and fed to the LM for answer generation. Right: a DLN-2 with a residual connection, performing the date understanding task: two prompts need to be learned. In this example, the hidden template extends Chain-Of-Thought [48] with a learnable prefix; we consider the output of the first layer, hidden, as a latent variable  $h$ . We use variational inference to learn  $\pi_0, \pi_1$ . Templates can be considered as an hyperparameter of the network.
We first show how to perform prompt optimization in a shallow 1-layer language network (DLN-1) which parametrizes a distribution  $p_{\mathrm{LM}}(y|x,\pi)$ , where  $x$  and  $y$  are string input and output respectively, and  $\pi$  is the learnable prompt (Figure 1, left). Our prompt optimization techniques extend the Automatic Prompt Engineer (APE) procedure from Zhou et al. [57]. We show how our prompts can include a verbalization of difficult examples from the task: the final prompts are a combination of instruction directives, akin to zero-shot learning [17], and task examples, akin to in-context learning [20]. This significantly improves downstream performance, surpassing APE on several tasks.
Then, we show how to train a 2-layer DLN (DLN-2), which parametrizes a probability distribution:
$$
p _ {\mathrm {D L N} - 2} (y | x) = \sum_ {h} p _ {\mathrm {L M}} (y | h, x, \pi_ {1}) p _ {\mathrm {L M}} (h | x, \pi_ {0}),
$$
where  $h$  is the string output of the first LLM layer (Figure 1, right). We consider  $h$  as a latent variable: to maximize the marginal log-likelihood, we formulate a variational inference algorithm that uses an approximate posterior over  $h$ . Note that this formalism easily encompasses more than two layers.
Considering outputs of the hidden language layers as latent variables allows us to encompass various established prompting methods, such as Chain-Of-Thought (CoT) [48] and self-consistency (SC-CoT) [46]. Particularly, CoT can be seen as a particular DLN-2 with the first layer prompt set to "Let's think step by step" and the second layer prompt set to "The answer is"; we can either learn such prompts or learn a supplement to those as in Figure 1. SC-CoT can be seen as marginalizing over CoT strings sampled from a task-agnostic prior: when using the template in Figure 1 (right), our method generalizes this perspective by learning a task-specific prior distribution over successful CoTs.
The rest of the paper is organized as follows. First, we provide an interpretation of LLMs as shallow networks, drawing a number of analogies with standard parametric and non-parametric models and explaining how best to train them. After exploring their limitations, we propose to stack two such
LLMs to form a DLN-2. We show how they can be trained using a form of variational inference, then demonstrate their performance on a series of reasoning and language understanding tasks.
2 One-Layer Language Networks
A pre-trained LLM with frozen weights might be thought of as a complete function class indexed by prompts. The output  $y$  of an LLM for an input  $x$  can be modulated through a prompt  $\pi$  by feeding a combination of  $\pi$  and  $x$  to the LLM. Hence, from a low-level perspective, the function class of an LLM is defined by its architecture, i.e., its depth, number of heads, context size, etc., and training happens at the parameter level. It is data and compute intensive and should be done rarely. From a high-level perspective, the function class of an LLM is defined by the pre-trained model chosen (LLAMA [44], text-davinci-003, GPT-4, etc.), and training happens by fine-tuning the model or by choosing the prompt.
There are two ways of optimizing an LLM at the prompt level. The first one is prompt engineering, a parametric optimization, where the optimization space is independent of the size of the dataset. Because this optimization usually happens in discrete space, gradient-based techniques do not apply and most efforts rely on a combination of random or local search and human heuristics [21, 55]. The second one is in-context learning (ICL), a non-parametric optimization technique where the solution is a direct function of a subset of examples [5, 20]. This approach works well for few-shot learning but scaling it to larger datasets has both performance and computational issues. We shall now generalize previous work in discrete prompt optimization [57, 21] with the ultimate goal of learning a set of prompts in a language network.
2.1 Language Layers
We use language layer to refer to a (stochastic) computation that takes as input a string  $x$  and outputs a string  $y$ . This computation is modulated by another string,  $\pi$ , generally called a prompt or instruction [57]. The string transduction is performed by an operator LM, by feeding  $x$  and  $\pi$  as context and generating a continuation  $y$ . Templates describe the way  $x$  and  $\pi$  are combined prior to being fed to the LM operator. These are functions that accept strings as variables and output a string. We will denote such templates with this font T. A simple forward template F is the concatenation, i.e.  $\mathsf{F}(x,\pi) = "\{\pi\}\{x\}"$ . We also explore more complex ones, examples of which can be seen in Figure 1.
Given an input  $x$ , a prompt  $\pi$ , and a template  $F$ , a language layer defines a probability distribution  $p_{\mathrm{LM}}(y|\mathbf{F}(x,\pi))$  over output strings  $y$  as computed by the LM. In the next section, we describe a generic framework for optimizing the weights  $\pi$  for a language layer.
2.2 Prompt Optimization: Improved APE
Because the search for the best prompt happens over a discrete space, we will rely on a local search method, using an LLM to implement a distance measure between prompts. The procedure can be seen as an extension of Automatic Prompt Engineer (APE), recently proposed by Zhou et al. [57], and will serve as a stepping stone towards introducing our algorithm for training deep language networks. Our prompt optimization algorithm can be structured as follows:
1. Given the current prompt  $\pi$  and a current batch of examples  $\{x,y\}$ , generate  $N$  "local" candidates  $\pi^1, \ldots, \pi^N$  using a prompt proposal distribution;
2. Score each candidate using a (potentially stochastic) scoring function  $s$ , then choose  $\pi = \arg \max_{\pi^n} s(\pi^n)$ .
Prompt Proposal Local search algorithms assume a distance measure between inputs to crawl the search space. In this setting, we rely on LLMs to generate local modifications to the prompts. Our prompt proposal distribution takes as conditioning information i) the batch given as input to the layer, ii) its corresponding output  $\{x,y,\hat{y}\}$ , and iii) the current prompt  $\pi$ . The proposal distribution  $p_{\mathrm{LM}}(\pi^n|\mathsf{B}_\pi (\{x,y,\hat{y}\},\pi))$  wraps this information using a particular "backward" template  $\mathsf{B}_{\pi}$ , which can be found in Appendix D. This approach is similar to the instruction template used by Zhang et al. [55], with the exception that we also integrate information about the model's own predictions, which
Algorithm 1 One-Layer Language Network (DLN-1) Training Algorithm
Require:	ˆy ~ pLH(t)(y|c)	▷ generates a completion of prefix c with temperature t
Require:	log pLM(h|c)	▷ return log-prob of h following c
Require:	N: prompt samples, I: iterations, D: dataset
Require:	F: template for the inference/forward pass
Require:	Bπ: template for prompt proposal/backward pass.
1: Initialize π with a task description or empty
2: for i in [1, I] do
3:	x, y ~ D	▷ Sample minibatch
4:	ˆy ← p0LM(y|F(x, π))	▷ Do inference pass
5:	π1, ..., πN ~ p0LM(π|Bπ{x, y,ˆy}, π))	▷ Sample N candidate prompts
6:	s1, ..., sN ← log pLM(y|F(x, πn))	▷ Score all prompts
7:	π ← arg maxπn{s1, ..., sN}	▷ Select prompt with best score
8: end for
we found to empirically help performance given that the model tends to propose prompts that correct its own errors. We sample from the prompt proposal distribution to generate a set of  $N$  prompts. A particularly important aspect is ensuring the diversity of the candidate pool  $\pi^1, \ldots, \pi^N$ . We devise several strategies to improve the diversity and the usefulness of the candidate samples in Section 4.
Prompt Selection Once a set of  $N$  prompts has been generated, we use a scoring function to select the updated prompt. We assume access to the log-likelihoods of the LM operator and we rank the candidate prompts to maximize data log-likelihood  $\pi = \arg \max_{\pi^n}\log p_{\mathrm{LM}}(y|\mathbf{F}(x;\pi^n))$ . In practice, we normalize this log-probability by the length of the output string. While we focus on that metric in this work, there is no restriction on the scoring function that can be used. We use backtracking to increase the robustness of our selection mechanism, as well as a memory of well-performing prompts for efficiency. We present both strategies in Section 4. The sketch of a 1-layer prompt optimization algorithm is described in Algorithm 1, ignoring backtracking and memory for simplicity.
The results of our prompt optimization may be found in Table 1 and will be discussed in detail in Section 5.2. We now turn to extending prompt optimization to architectures with two layers.
3 Two-Layer Deep Language Networks (DLN-2)
The natural extension of DLN-1 is DLN-2, in which language layers are stacked, i.e. the output of the first language layer is the input to the second one. A 2-layer network induces a distribution over outputs of the form:
$$
p _ {\mathrm {D L N} - 2} (y | x) = \sum_ {h} p _ {\mathrm {L M}} \left(y \mid \mathbf {F} _ {r} (h, x, \pi_ {1})\right) p _ {\mathrm {L M}} \left(h \mid \mathbf {F} \left(x, \pi_ {0}\right)\right) \tag {1}
$$
where  $h$  is a latent variable that potentially makes it easier to explain the target  $y$ . The output layer is also conditioned on  $x$  through  $\mathbb{F}_r$ , forming a residual connection (Figure 1). This formulation is reminiscent of past work using latent language representations to guide document summarization [26]. In our case, however, the encoding/decoding distributions are parameterized by natural language prompts  $\Pi = \{\pi_0, \pi_1\}$ , and we do not assume access to the LLM parameters.
While this architecture has more expressive power than a shallow language network, the prompt optimization problem becomes harder now that we have to jointly search over both  $\pi_0$  and  $\pi_{1}$ . Doing random search in this space is impractical [8] and manual tuning of the weights is also exponentially harder than with a single prompt. We turn to variational inference to address this issue.
3.1 Variational Inference Objective
The layerwise decomposition of our system allows us to leverage tools from approximate inference in probabilistic models to learn  $\Pi$ . In particular, we propose to use variational inference to learn  $\Pi$  [3, 16]. We posit an approximate posterior  $q(h)$  over the latent variable  $h$ , and bound the marginal
log-likelihood of  $y$  given  $x$  by computing the ELBO:
$$
\log p _ {\mathrm {D L N} - 2} (y | x) \geq \sum_ {h} q (h) \left[ \log p _ {\mathrm {L M}} \left(y \mid \mathbf {F} _ {r} (h, x, \pi_ {1})\right) p _ {\mathrm {L M}} \left(h \mid \mathbf {F} (x, \pi_ {0})\right) \right] + H [ q (h) ], \tag {2}
$$
which allows us to decompose the optimization over  $\Pi$  in two independent optimization problems, over both  $\pi_0$  and  $\pi_1$ :
$$
\pi_ {0} ^ {*} = \arg \max  _ {\pi_ {0}} \sum_ {x, h} w _ {h} \log p (h | \mathbf {F} (x, \pi_ {0})), \quad \pi_ {1} ^ {*} = \arg \max  _ {\pi_ {1}} \sum_ {(x, y), h} w _ {h} \log p (y | \mathbf {F} (h, x, \pi_ {1})). \tag {3}
$$
The search over  $\pi_1$  is identical to the prompt optimization described in Section 2, with the difference that the inputs now depend on the approximate posterior samples  $h$  in addition to the inputs  $x$ . The search over  $\pi_0$  uses a similar strategy but uses the posterior samples  $h$  as targets, instead of  $y$ .
Although this bound allows us to decompose the optimization w.r.t.  $\Pi$ , it is only useful if it is close to the true value. Since its looseness is the KL divergence between  $q(h)$  and the true posterior,  $\mathrm{KL}(q(h)||p(h|y,x))$ : we need to find an approximate posterior  $q$  closely matching the true posterior. In what follows, we specify how we parametrize the approximate posterior and how we tighten the approximation via posterior sharpening.
Hidden Proposal We will also be using an LLM to sample candidate hidden states from  $q(h)$ . Unless specified, for simplicity, we use the same LM operator used in our language layers. The approximate posterior can condition on arbitrary amount of information but especially useful might be to condition on the true target  $y$ . If it conditions on the hidden state coming from the "forward pass",  $\hat{h} \sim p_{\mathrm{LM}}(h|\mathbf{F}(x,\pi_0))$ , then  $q_{\mathrm{edit}}(h) = p_{\mathrm{LM}}(h|\mathsf{B}_h(\hat{h},y,\pi_1))$ .  $\mathsf{B}_h$  is a specifically tailored hidden proposal template (Appendix D).  $q_{\mathrm{edit}}$  performs a sort of edit operation, where the LM is tasked to rewrite the hidden variable  $\hat{h}$  given some extra knowledge of the ground-truth label  $y$  and of  $\pi_1$ . Alternatively, we can set the posterior to be equal to the prior, i.e.  $q_{\mathrm{pri}}(h) = p_{\mathrm{LM}}(h|\mathsf{F}(x,\pi_0))$ , or the prior with additional information about the label  $y$ ,  $q_{\mathrm{pri + }}(h) = p_{\mathrm{LM}}(h|\mathsf{B}_y(x,\pi_0,y))$  (Appendix D). This amounts to re-computing the hidden state knowing privileged information about the label. We found most effective to sample hidden states from a mixture of  $q_{\mathrm{pri}}$  and  $q_{\mathrm{pri + }}$ .
Posterior Sharpening Given the absence of learnable parameters in  $q(h)$ , the induced approximate posterior might still be far from the true posterior. To bridge this gap, we reweigh each sample  $h^i$  based on its probability under the true posterior distribution. More precisely, we compute  $\tilde{w}^i = \log p_{\mathrm{LM}}(y|\mathbf{F}_r(h^i,x,\pi_1)) + \log p_{\mathrm{LM}}(h^i|\mathbf{F}(x,\pi_0))$ , then assign to each  $h_i$  the probability  $w_i = \exp (\alpha \tilde{w}_i) / \sum_j\exp (\alpha \tilde{w}_j)$ , where  $\alpha$  is a tunable temperature parameter that controls the entropy of the posterior weights. The full algorithm for training a DLN-2 is presented in Algorithm 2.
4 Practical Instantiation
Although our method aim to learning prompts in stacked LLM architectures, we do rely on a good amount of prompt engineering for our templates. Hereafter, we detail some choices that were fundamental to make our approach work in practice.
**Proposal Diversity** To ensure a diversity of the samples for both the prompt proposal distribution, we found helpful to use two strategies. The first is to modify the backward templates  $\mathsf{B}_{\pi}$  before drawing a sample from the proposal distribution  $p_{\mathrm{LM}}(\pi)$ . To achieve so, we parametrize the basic templates with a “{\{message\}'} variable that we instantiate from a pool of hand-written instructions, that describe different behaviors the model should follow to propose new  $\pi$ , e.g. “shorten the previous instruction”, “give useful examples”, etc. These can be interpreted as meta-instructions, i.e. high-level directives that inform the model on how to create a better instruction for the task, and extend instruction-induction templates used in [12, 55]. These can be found in Appendix D. In the future, we could envision to extend learning to these instructions. In the case of  $\mathsf{B}_{\pi}$ , they could function as parameters for a prior over the weights of the DLN. The second strategy to ensure more diversity is that we instantiate  $\mathsf{B}_{\pi}$  with a different random subset of examples in the current batch, before drawing each sample  $\pi^n$ . This effectively modifies the generation context for each sample  $\pi^n$ .
Learning In-Context Learning One strategy we found particularly effective is to integrate in the pool of meta-instructions an additional instruction that asks the LM to give useful examples to improve
Algorithm 2 Two-Layer Deep Language Network (DLN-2) Training Algorithm
1:  $\hat{y}\sim p_{\mathrm{LM}}^t (c)$  generates a completion of prefix  $c$  with temperature  $t$    
2:  $\log p_{\mathrm{LM}}(h|c)$  return log-prob of  $h$  following  $c$    
3:  $N$  : prompt samples,  $K$  : posterior samples,  $I$  : iterations,  $\mathcal{D}$  dataset   
4: F,  $\mathsf{F}_r$  : templates for the inference/forward pass without and with residual connection   
5:  $\mathtt{B}_{\pi},\mathtt{B}_h$  : templates for prompt and hidden proposal/backward pass.   
6: Initialize  $\pi_0,\pi_1$  with a generic sentence or task description   
7: for i in [1, I] do   
8:  $x,y\sim \mathcal{D}$  Sample minibatch   
9:  $\hat{h}\gets p_{\mathrm{LM}}^{0}(\mathbb{F}(x,\pi_{0}))$  Do inference pass, first layer   
10:  $\hat{y}\gets p_{\mathrm{LM}}^{0}(\mathsf{F}_{r}(\hat{h},x,\pi_{1}))$  Do inference pass, second layer   
11:  $h^1,\ldots ,h^K\sim p_{\mathrm{LM}}^{0.7}(\mathsf{B}_h(\hat{h},x,y,\pi_1,\pi_0))$  Sample  $K$  posterior proposals for  $h$    
12:  $\alpha^1,\dots ,\alpha^K\gets \log p_{\mathrm{LM}}(h^k |\mathbb{F}(x,\pi_0))$  Compute prior log-probs for all  $h^k$  samples   
13:  $\beta^1,\dots ,\beta^K\gets \log p_{\mathrm{LM}}(y|\mathsf{F}_r(h^k,x,\pi_1))$  Compute log-likelihoods for all  $h^k$  samples   
14:  $q^{k}\gets \exp (\alpha^{k} + \beta^{k}) / (\sum_{k}\exp (\alpha^{k} + \beta^{k}))$  Compute normalized posterior weights   
15:  $h^{*}\gets \arg \max_{h}\{q^{1},\dots ,q^{K}\}$  Select best posterior sample   
16:  $\pi_0^1,\dots ,\pi_0^N\sim p_{\mathrm{LM}}^{0.7}(\mathtt{B}_{\pi}(\{x,\hat{h},h^*\} ,\pi_0))$  Sample  $N$  candidate prompts for  $\pi_0$    
17:  $\pi_1^1,\dots ,\pi_1^N\sim p_{\mathrm{LM}}^{0.7}(\mathtt{B}_{\pi}(\{\hat{h},\hat{y},y\} ,\pi_1))$  Sample  $N$  candidate prompts for  $\pi_1$    
18:  $s_0^1,\dots ,s_0^N\gets \sum_k q^k\log p_{\mathrm{LM}}(h^k |\mathbb{F}(x,\pi_0^n))$  Compute ELBO for all prompts  $\pi_0^n$    
19:  $s_1^1,\dots ,s_1^N\gets \sum_k q^k\log p_{\mathrm{LM}}(y|\mathsf{F}_r(h^k,x,\pi_1^n))$  Compute ELBO for all prompts  $\pi_1^n$    
20:  $\pi_0\gets \arg \max_{\pi_0^n}\{s_0^1,\dots ,s_0^N\}$  Select prompt  $\pi_0$  with best score   
21:  $\pi_1\gets \arg \max_{\pi_1^n}\{s_1^1,\dots ,s_1^N\}$  Select prompt  $\pi_1$  with best score   
22: end for
its current prompt  $\pi$ . Empirically, we observed that this allows the model to sample candidate prompts  $\pi^n$  that contain synthetic examples for the task, embedded in natural language. Examples of this interesting behavior can be found in Appendix F. We found that this behavior is particularly interesting as the resulting prompts often perform better than standard ICL. We hypothesize this is due to both  $i$ ) the "verbalization" of the example in the prompt, which modifies the dataset syntax into a more suitable one, and  $ii$ ) the fact that the model can dynamically select which examples are most important to integrate in the prompts, given the errors made during training. Therefore, we suspect that DLN achieves a similar effect to recent techniques that select important examples for ICL [20, 35, 40, 18, 47] with the improvement of naturally conditioning the selection on the end task performance via end-to-end training.
Backtracking and Memory Optimization of both DLN-1 and DLN-2 is challenging due to the fact that we do not have gradient information and we sample a restricted set of candidates  $\pi^n$  at each optimization step due to computational reasons. We deploy multiple strategies to allow the network to be robust to sampling/selection errors. First, we include the current prompt  $\pi$  into the set of candidate prompts to be scored at the current iteration  $\pi^n$ . This allows the model to not take the step if the previous prompt performed better. Second, we keep a memory of  $M = 5$  best prompts found by tracking validation set performance.
Exploration Reward When training a DLN-2, we empirically observed that the first layer prompt  $\pi_0$  was updating very slowly. Due to the fact that the approximate posterior shares templates with the prior used in the forward pass, the posterior samples  $h^i$  are close to  $\hat{h}$  and the maximizer of Equation (3) remains  $\pi_0$ . To address this issue, we add to the scores of each candidate prompt an exploration reward that is proportional to the negative log-probability of those  $\hat{h}$  that led to an incorrect prediction:  $r = -\lambda \log p_{\mathrm{LM}}(\hat{h}|\mathbf{F}(x,\pi^n))$ , if  $\hat{y} \neq y$ . This encourages the model to both find prompts that maximize the log-probability of high-probability posterior samples and at the same time minimize the log-probability of prior samples that led to incorrect predictions. We anneal  $\lambda$  to 0 during training with a constant schedule and we select the initial  $\lambda$  by monitoring validation performance for each task.