SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We use the chain rule to update the parameters  $\theta$ . This requires a choice of distribution  $q_{\phi}(r|\mathbf{P})$  for which computing the mean  $\mathbb{E}_{q_{\phi}(r|\mathbf{P})}[r]$  is easy. The simplest option is to parameterize  $q_{\phi}$  with a univariate normal distribution. We predict the mean and variance using a neural network with parameters  $\phi$ . For example, a neural network  $m_{\phi}$  would compute  $\mu = m_{\phi}(\mathbf{P})$ . Then, the mean parameter is the expectation on the right-hand side of Equation 26. The loss function for  $f_{\theta}$  with this parameterization is:
$$
\mathcal {L} _ {N N} (\boldsymbol {\theta}) = - m _ {\phi} \left(f _ {\boldsymbol {\theta}} (\mathbf {x})\right), \quad \mathbf {x} \sim \mathcal {D} \tag {27}
$$
Interestingly, like A-NESI, this gives zero-variance gradient estimates! Of course, bias comes from the error in the approximation of  $q_{\phi}$ .
Like A-NESI, we expect the success of this method to rely on the ease of finding a suitable prior over  $\mathbf{P}$  to allow proper training of the inference model. See the discussion in Section 3.2.3. We also expect that, like in A-NESI, it will be easier to train the inference model if the output  $r = g(\mathbf{z})$  is structured instead of a single scalar. We refer to Section 3.2.2 for this idea. Other challenges might be stochastic and noisy output measurements of  $r$  and non-stationarity of  $g$ , for instance, when training a VAE [26].
Figure 6: The tree flow network corresponding to weighted model counting on the formula  $A \lor B$ . Following edges upwards means setting the corresponding binary variable to true (and to false by following edges downwards). We first choose probabilities for the propositions  $A$  and  $B$ , then choose whether we want to sample a world that satisfies the formula  $A \lor B$ .  $y = 1$  is the WMC of  $A \lor B$ , and equals its outgoing flow  $P_A + P_B - P_A P_B$ . Terminal states (with two circles) represent choices of the binary variables  $A$  and  $B$ . These are connected to a final sink node, corresponding to the prior over worlds  $p(\mathbf{w}|\mathbf{P})$ . The total ingoing and outgoing flow to this network is 1, as we deal with normalized probability distributions  $p$  and  $q_{\phi}$ .
D A-NESI and GFlowNets
A-NeSI is heavily inspired by the theory of GFlowNets [9, 10], and we use this theory to derive our loss function. In the current section, we discuss these connections and the potential for future research by taking inspiration from the GFlowNet literature. In this section, we will assume the reader is familiar with the notation introduced in [10] and refer to this paper for the relevant background.
D.1 Tree GFlowNet representation
The main intuition is that we can treat the inference model  $q_{\phi}$  in Equation 3 as a 'trivial' GFlowNet. We refer to Figure 6 for an intuitive example. It shows what a flow network would look like for the formula  $A \lor B$ . We take the reward function  $R(\mathbf{w}, \mathbf{y}) = p(\mathbf{w}, \mathbf{y})$ . We represent states  $s$  by  $s = (\mathbf{P}, \mathbf{y}_{1:i}, \mathbf{w}_{1:j})$ , that is, the belief  $\mathbf{P}$ , a list of some dimensions of the output instantiated with a value and a list of some dimensions of the world assigned to some value. Actions  $a$  set some value to the next output or world variable, i.e.,  $A(s) = Y_{i+1}$  or  $A(s) = W_{j+1}$ .
Note that this corresponds to a flow network that is a tree everywhere but at the sink since the state representation conditions on the whole trajectory observed so far. We demonstrate this in Figure 6. We assume there is some fixed ordering on the different variables in the world, which we generate the value of one by one. Given this setup, Figure 6 shows that the branch going up from the node  $y$  corresponds to the regular weighted model count (WMC) introduced in Equation 1.
The GFlowNet forward distribution  $P_F$  is  $q_{\phi}$  as defined in Equation 3. The backward distribution  $P_B$  is  $p(\mathbf{w}, \mathbf{y} | \mathbf{P})$  at the sink node, which chooses a terminal node. Then, since we have a tree, this determines the complete trajectory from the terminal node to the source node. Thus, at all other states, the backward distribution is deterministic. Since our reward function  $R(\mathbf{w}, \mathbf{y}, \mathbf{P}) = p(\mathbf{w}, \mathbf{y} | \mathbf{P})$  is normalized, we trivially know the partition function  $Z(\mathbf{P}) = \sum_{\mathbf{w}} \sum_{\mathbf{y}} R(\mathbf{w}, \mathbf{y} | \mathbf{P}) = 1$ .
D.2 Lattice GFlowNet representation
Our setup of the generative process assumes we are generating each variable in the world in some order. This is fine for some problems like MNISTAdd, where we can see the generative process as 'reading left to right'. For other problems, such as Selenium, the order in which we would like to generate the symbols is less obvious. Would we generate block by block? Row by row? Column by column? Or is the assumption that it needs to be generated in some fixed order flawed by itself?
Figure 7: The lattice flow network corresponding to weighted model counting on the formula  $A \vee B$ . In this representation, nodes represent sets of included propositions. Terminal states represent sets of random variables such that  $A \vee B$  is true given  $y = 1$ , or false otherwise.
In this section, we consider a second GFlowNet representation for the inference model that represents states using sets instead of lists. We again refer to Figure 7 for the resulting flow network of this representation for  $A \vee B$ . We represent states using  $s = (\mathbf{P}, \{y_i\}_{i \in I_Y}, \{w_i\}_{i \in I_W})$ , where  $I_Y \subseteq \{1, \dots, k_Y\}$  and  $I_W \subseteq \{1, \dots, k_W\}$  denote the set of variables for which a value is chosen. The possible actions from some state correspond to  $A(s) = \bigcup_{i \notin I_W} W_i$  (and analogous for when  $\mathbf{y}$  is not yet completely generated). For each variable in  $W$  for which we do not have the value yet, we add its possible values to the action space.
With this representation, the resulting flow network is no longer a tree but a DAG, as the order in which we generate the different variables is now different for every trajectory. What do we gain from this? When we are primarily dealing with categorical variables, the two gains are 1) we no longer need to impose an ordering on the generative process, and 2) it might be easier to implement parameter sharing in the neural network that predicts the forward distributions, as we only need a single set encoder that can be reused throughout the generative process.
However, the main gain of the set-based approach is when worlds are all (or mostly) binary random variables. We illustrate this in Figure 7. Assume  $W = \{0,1\}^{k_W}$ . Then we can have the following state and action representations:  $s = (\mathbf{P},\mathbf{y},I_W)$ , where  $I_W \subseteq \{1,\dots,k_W\}$  and  $A(s) = \{1,\dots,k_W\} \setminus I_W$ . The intuition is that  $I_W$  contains the set of all binary random variables that are set to 1 (i.e., true), and  $\{1,\dots,k_W\} \setminus I_W$  is the set of variables set to 0 (i.e., false). The resulting flow network represents a partial order over the set of all subsets of  $\{1,\dots,k_W\}$ , which is a lattice, hence the name of this representation.
With this representation, we can significantly reduce the size and computation of the flow network required to express the WMC problem. As an example, compare Figures 6 and 7, which both represent the WMC of the formula  $A \lor B$ . We no longer need two nodes in the branch  $y = 0$  to represent that we generate  $A$  and  $B$  to be false, as the initial empty set {} already implies they are. This will save us two nodes. Similarly, we can immediately stop generating at  $\{A\}$  and  $\{B\}$ , and no longer need to generate the other variable as false, which also saves a computation step.
While this is theoretically appealing, the three main downsides are 1)  $P_B$  is no longer trivial to compute; 2) we have to handle the fact that we no longer have a tree, meaning there is no longer a unique optimal  $P_F$  and  $P_B$ ; and 3) parallelization becomes much trickier. We leave exploring this direction in practice for future work.
E Analyzing the Joint Matching Loss
This section discusses the loss function we use to train the joint variant in Equation 8. We recommend interested readers first read Appendix D.1. Throughout this section, we will refer to  $p \coloneqq p(\mathbf{w},\mathbf{y}|\mathbf{P})$  (Equation 5) and  $q \coloneqq q_{\phi ,\psi}(\mathbf{w},\mathbf{y}|\mathbf{P})$  (Equation 2). We again refer to [10] for notational background.
E.1 Trajectory Balance
We derive our loss function from the recently introduced Trajectory Balance loss for GFlowNets, which is proven to approximate the true Markovian flow when minimized. This means sampling from the GFlowNet allows sampling in proportion to reward  $R(s_{n}) = p$ . The Trajectory Balance loss is given by
$$
\mathcal {L} (\tau) = \left(\log \frac {F \left(s _ {0}\right) \prod_ {t = 1} ^ {n} P _ {F} \left(s _ {t} \mid s _ {t - 1}\right)}{R \left(s _ {n}\right) \prod_ {t = 1} ^ {n} P _ {B} \left(s _ {t - 1} \mid s _ {t}\right)}\right) ^ {2}, \tag {28}
$$
where  $s_0$  is the source state, in our case  $\mathbf{P}$ , and  $s_n$  is some terminal state that represents a full generation of  $\mathbf{y}$  and  $\mathbf{w}$ . In the tree representation of GFlowNets for inference models (see Appendix D.1), this computation becomes quite simple:
1.  $F(s_0) = 1$ , as  $R(s_n) = p$  is normalized;
2.  $\prod_{t=1}^{n} P_F(s_t | s_{t-1}) = q$ : The forward distribution corresponds to the inference model  $q_\phi(\mathbf{w}, \mathbf{y} | \mathbf{P})$ ;
3.  $R(s_{n}) = p$ , as we define the reward to be the true joint probability distribution  $p(\mathbf{w},\mathbf{y}|\mathbf{P})$ ;
4.  $\prod_{t=1}^{n} P_B(s_{t-1}|s_t) = 1$ , since the backward distribution is deterministic in a tree.
Therefore, the trajectory balance loss for (tree) inference models is
$$
\mathcal {L} (\mathbf {P}, \mathbf {y}, \mathbf {w}) = \left(\log \frac {q}{p}\right) ^ {2} = \left(\log q - \log p\right) ^ {2}, \tag {29}
$$
i.e., the term inside the expectation of the joint matching loss in Equation 8. This loss function is stable because we can sum the individual probabilities in log-space.
A second question might then be how we obtain 'trajectories'  $\tau = (\mathbf{P},\mathbf{y},\mathbf{w})$  to minimize this loss over. The paper on trajectory balance [33] picks  $\tau$  on-policy, that is, it samples  $\tau$  from the forward distribution (in our case, the inference model  $q_{\phi ,\psi}$ ). We discussed when this might be favorable in our setting in Appendix B (Equation 16). However, the joint matching loss as defined in Equation 8 is off-policy, as we sample from  $p$  and not from  $q_{\phi ,\psi}$ .
E.2 Relation to common divergences
These questions open quite some design space, as was recently noted when comparing the trajectory balance loss to divergences commonly used in variational inference [34]. Redefining  $P_F = q$  and  $P_B = p$ , the authors compare the trajectory balance loss with the KL-divergence and the reverse KL-divergence and prove that
$$
\nabla_ {\phi} D _ {K L} (q | | p) = \frac {1}{2} \mathbb {E} _ {\tau \sim q} [ \nabla_ {\phi} \mathcal {L} (\tau) ]. \tag {30}
$$
That is, the on-policy objective minimizes the reverse KL-divergence between  $p$  and  $q$ . We do not quite find such a result for the off-policy version we use for the joint matching loss in Equation 8:
$$
\nabla_ {\phi} D _ {K L} (p | | q) = - \mathbb {E} _ {\tau \sim p} [ \nabla_ {\phi} \log q ] \tag {31}
$$
$$
\mathbb {E} _ {\tau \sim p} [ \nabla_ {\phi} \mathcal {L} (\tau) ] = - 2 \mathbb {E} _ {\tau \sim p} [ (\log p - \log q) \nabla_ {\phi} \log q ] \tag {32}
$$
So why do we choose to minimize the joint matching loss rather than the (forward) KL divergence directly? This is because, as is clear from the above equations, it takes into account how far the 'predicted' log-probability  $\log q$  currently is from  $\log p$ . That is, given a sample  $\tau$ , if  $\log p < \log q$ , the joint matching loss will actually decrease  $\log q$ . Instead, the forward KL will increase the probability for every sample it sees, and whether this particular sample will be too likely under  $q$  can only be derived through sampling many trajectories.
Furthermore, we note that the joint matching loss is a 'pseudo' f-divergence with  $f(t) = t\log^2 t$  [34]. It is not a true f-divergence since  $t\log^2 t$  is not convex. A related well-known f-divergence is the Hellinger distance given by
$$
H ^ {2} (p | | q) = \frac {1}{2} \mathbb {E} _ {\tau \sim p} [ (\sqrt {p} - \sqrt {q}) ^ {2} ]. \tag {33}
$$
This divergence similarly takes into account the distance between  $p$  and  $q$  in its derivatives through squaring. However, it is much less stable than the joint matching loss since both  $p$  and  $q$  are computed by taking the product over many small numbers. Computing the square root over this will be much less numerically stable than taking the logarithm of each individual probability and summing.
Finally, we note that we minimize the on-policy joint matching  $\mathbb{E}_q[(\log p - \log q)^2]$  by taking derivatives  $\mathbb{E}_q[\nabla_{\phi ,\psi}(\log p - \log q)^2 ]$ . This is technically not minimizing the joint matching, since it ignores the gradient coming from sampling from  $q$ .
F Dirichlet prior
This section describes how we fit the Dirichlet prior  $p(\mathbf{P})$  used to train the inference model. During training, we keep a dataset of the last 2500 observations of  $\mathbf{P} = f_{\theta}(\mathbf{x})$ . We have to drop observations frequently because  $\theta$  changes during training, meaning that the empirical distribution over  $\mathbf{P}$ s changes as well.
We perform an MLE fit on  $k_{W}$  independent Dirichlet priors to get parameters  $\alpha$  for each. The log-likelihood of the Dirichlet distribution cannot be found in closed form [39]. However, since its log-likelihood is convex, we run ADAM [25] for 50 iterations with a learning rate of 0.01 to minimize the negative log-likelihood. We refer to [39] for details on computing the log-likelihood and alternative options. Since the Dirichlet distribution accepts positive parameters, we apply the softplus function on an unconstrained parameter during training. We initialize all parameters at 0.1.
We added L2 regularization on the parameters. This is needed because at the beginning of training, all observations  $\mathbf{P} = f_{\theta}(\mathbf{x})$  represent uniform beliefs over digits, which will all be nearly equal. Therefore, fitting the Dirichlet on the data will give increasingly higher parameter values, as high parameter values represent low-entropy Dirichlet distributions that produce uniform beliefs. When the Dirichlet is low-entropy, the inference models learn to ignore the input belief  $\mathbf{P}$ , as it never changes. The L2 regularization encourages low parameter values, which correspond to high-entropy Dirichlet distributions.
G Designing symbolic pruners
We next discuss four high-level approaches for designing the optional symbolic pruner, each with differing tradeoffs in accuracy, engineering time and efficiency:
1. Mathematically derive efficient solvers. For simple problems, we could mathematically derive an exact solver. One example of an efficient symbolic pruner, along with a proof for exactness, is given for Multi-digit MNISTAdd in Appendix H. This pruner is linear-time. However, for most problems we expect the pruner to be much more computationally expensive.
2. Use SAT-solvers. Add the sampled symbols  $\mathbf{y}$  and  $\mathbf{w}_{1:i}$  to a CNF-formula, and ask an SAT-solver if there is an extension  $\mathbf{w}_{i+1:k_W}$  that satisfies the CNF-formula. SAT-solvers are a general approach that will work with every function  $c$ , but using them comes at a cost.
The first is that we would require grounding the logical representation of the problem. Furthermore, to do SAT-solving, we have to solve a linear amount of NP-hard problems. However, competitive SAT solvers can deal with substantial problems due to years of advances in their design [7], and a linear amount of NP-hard calls is a lower complexity class than #P hard. Using SAT-solvers will be particularly attractive in problem settings where safety and verifiability are critical.
3. Prune with local constraints. In many structured prediction tasks, we can use local constraints of the symbolic problem to prune paths that are guaranteed to lead to branches that can never create possible worlds. However, local constraints do not guarantee that each non-pruned path contains a possible world, but this does not bias the inference model, as the neural network will (eventually) learn when an expansion would lead to an unsatisfiable state.
One example is the shortest path problem, where we filter out directions that would lead outside the  $N \times N$  grid, or that would create cycles (See Appendix I.3). However, this only ensures we find a path, but does not ensure it is the shortest one.
4. Learn the pruner. Finally, we can learn the pruner, that is, we can train a neural network to learn satisfiability checking. One possible approach is to reuse the inference model trained on the belief  $\mathbf{P}$  that uniformly distributes mass over all worlds.
Learned pruners will be as quick as regular inference models, but are less accurate than symbolic pruners and will not guarantee that constraints are always satisfied during test-time.
We leave experimenting with learning the pruner for future work.
H MNISTAdd Symbolic Pruner
In this section, we describe a symbolic pruner for the Multi-digit MNISTAdd problem, which we compute in time linear to  $N$ . Note that  $\mathbf{w}_{1:N}$  represents the first number and  $\mathbf{w}_{N+1:2N}$  the second. We define  $n_1 = \sum_{i=1}^N w_i \cdot 10^{N-i-1}$  and  $n_2 = \sum_{i=1}^N w_{N+i} \cdot 10^{N-i-1}$  for the integer representations of these numbers, and  $y = \sum_{i=1}^{N+1} y_i \cdot 10^{N-i}$  for the sum label encoded by  $\mathbf{y}$ . We say that partial generation  $\mathbf{w}_{1:k}$  has a completion if there is a  $\mathbf{w}_{k+1:2N} \in \{0, \dots, 9\}^{2N-k}$  such that  $n_1 + n_2 = y$ .
Proposition H.1. For all  $N \in \mathbb{N}$ ,  $\mathbf{y} \in \{0,1\} \times \{0,\dots ,9\} ^N$  and partial generation  $\mathbf{w}_{1:k - 1} \in \{0,\ldots ,9\} ^k$  with  $k \in \{1,\dots ,2N\}$ , the following algorithm rejects all  $w_{k}$  for which  $\mathbf{w}_{1:k}$  has no completions, and accepts all  $w_{k}$  for which there are:
-  $k \leq N$ : Let  $l_k = \sum_{i=1}^{k+1} y_k \cdot 10^{k+1-i}$  and  $p_k = \sum_{i=1}^k w_k \cdot 10^{k-i}$ . Let  $S = 1$  if  $k = N$  or if the  $(k+1)th$  to  $(N+1)th$  digit of  $y$  are all 9, and  $S = 0$  otherwise. We compute two boolean conditions for all  $w_k \in \{0, \ldots, 9\}$ :
$$
0 \leq l _ {k} - p _ {k} \leq 1 0 ^ {k} - S \tag {34}
$$
We reject all  $w_{k}$  for which either condition does not hold.
-  $k > N$ : Let  $n_2 = y - n_1$ . We reject all  $w_k \in \{0, \dots, 9\}$  different from  $w_k = \left\lfloor \frac{n_2}{10^{N-k-1}} \right\rfloor \mod 10$ , and reject all  $w_k$  if  $n_2 < 0$  or  $n_2 \geq 10^N$ .
Proof. For  $k > N$ , we note that  $n_2$  is fixed given  $y$  and  $n_1$  through linearity of summation, and we only consider  $k \leq N$ . We define  $a_k = \sum_{i=k+2}^{N+1} y_i \cdot 10^{N+1-i}$  as the sum of the remaining digits of  $y$ . We note that  $y = l_k \cdot 10^{N-k} + a_k$ .
Algorithm rejects  $w_{k}$  without completions. We first show that our algorithm only rejects  $w_{k}$  for which no completion exists. We start with the constraint  $0 \leq l_{k} - p_{k}$ , and show that whenever this constraint is violated (i.e.,  $p_{k} > l_{k}$ ),  $\mathbf{w}_{1:k}$  has no completion. Consider the smallest possible completion of  $\mathbf{w}_{k+1:N}$ : setting each to 0. Then  $n_{1} = p_{k} \cdot 10^{N-k}$ . First, note that
$$
1 0 ^ {N - k} > 1 0 ^ {N - k} - 1 \geq a _ {k}
$$
Next, add  $l_{k}\cdot 10^{N - k}$  to both sides
$$
\left(l _ {k} + 1\right) \cdot 1 0 ^ {N - k} > l _ {k} \cdot 1 0 ^ {N - k} + a _ {k} = y
$$
By assumption,  $p_k$  is an integer upper bound of  $l_k$  and so  $p_k \geq l_k + 1$ . Therefore,
$$
n _ {1} = p _ {k} \cdot 1 0 ^ {N - k} > y
$$
Since  $n_1$  is to be larger than  $y$ ,  $n_2$  has to be negative, which is impossible.
Next, we show the necessity of the second constraint. Assume the constraint is unnecessary, that is,  $l_{k} > p_{k} + 10^{k} - S$ . Consider the largest possible completion  $\mathbf{w}_{k+1:N}$  by setting each to 9. Then
$$
\begin{array}{l} n _ {1} = p _ {k} \cdot 1 0 ^ {N - k} + 1 0 ^ {N - k} - 1 \\ = \left(p _ {k} + 1\right) \cdot 1 0 ^ {N - k} - 1 \\ \end{array}
$$
We take  $n_2$  to be the maximum value, that is,  $n_2 = 10^N - 1$ . Therefore,
$$
n _ {1} + n _ {2} = 1 0 ^ {N} - \left(p _ {k} + 1\right) \cdot 1 0 ^ {N - k} - 2
$$
We show that  $n_1 + n_2 < y$ . Since we again have an integer upper bound, we know  $l_k \geq p_k + 10^k - S + 1$ . Therefore,
$$
\begin{array}{l} y \geq \left(p _ {k} + 1 + 1 0 ^ {k} - S\right) 1 0 ^ {N - k} + a _ {k} \\ \geq n _ {1} + n _ {2} + 2 - S \cdot 1 0 ^ {N - k} + a _ {k} \\ \end{array}
$$
There are two cases.
-  $S = 0$ . Then  $a_{k} < 10^{N - k} - 1$ , and so
$$
y \geq n _ {1} + n _ {2} + 2 + a _ {k} > n _ {1} + n _ {2}.
$$
-  $S = 1$ . Then  $a_{k} = 10^{N - k} - 1$ , and so
$$
y \geq n _ {1} + n _ {2} + 1 > n _ {1} + n _ {2}.
$$
Algorithm accepts  $w_{k}$  with completions Next, we show that our algorithm only accepts  $w_{k}$  with completions. Assume Equation 34 holds, that is,  $0 \leq l_{k} - p_{k} \leq 10^{k} - S$ . We first consider all possible completions of  $\mathbf{w}_{1:k}$ . Note that  $p_{k} \cdot 10^{N-k} \leq n_{1} \leq p_{k} \cdot 10^{N-k} + 10^{N-k} - 1$  and  $0 \leq n_{2} \leq 10^{N} - 1$ , and so
$$
p _ {k} \cdot 1 0 ^ {N - k} \leq n _ {1} + n _ {2} \leq (p _ {k} + 1) \cdot 1 0 ^ {N - k} + 1 0 ^ {N} - 2.
$$
$$
l _ {k} \cdot 1 0 ^ {N - k} \leq y \leq (l _ {k} + 1) \cdot 1 0 ^ {N - k} - 1.
$$
By assumption,  $p_k \leq l_k$ , so  $p_k \cdot 10^{N - k} \leq l_k \cot 10^{N - k}$ . For the upper bound, we again consider two cases. We use the second condition  $l_k \leq 10^k + p_k - S$ :
-  $S = 0$ . Then (since there are no trailing 9s),
$$
\begin{array}{l} y \leq (l _ {k} + 1) \cdot 1 0 ^ {N - k} - 2 \\ \leq (1 0 ^ {k} + p _ {k} + 1) \cdot 1 0 ^ {N - k} - 1 \\ = (p _ {k} + 1) \cdot 1 0 ^ {N - k} + 1 0 ^ {N} - 2. \\ \end{array}
$$
-  $S = 1$ . Then with trailing 9s,
$$
\begin{array}{l} y = (l _ {k} + 1) \cdot 1 0 ^ {N - k} - 1 \\ \leq \left(1 0 ^ {k} + p _ {k}\right) \cdot 1 0 ^ {N - k} - 1 \\ = p _ {k} \cdot 1 0 ^ {N - k} + 1 0 ^ {N} - 1 \\ \leq \left(p _ {k} + 1\right) \cdot 1 0 ^ {N - k} + 1 0 ^ {N} - 2, \\ \end{array}
$$
since  $10^{N - k}\geq 1$
$$
p _ {k} \cdot 1 0 ^ {N - k} \leq y \leq (p _ {k} + 1) \cdot 1 0 ^ {N - k} + 1 0 ^ {N} - 2
$$
and so there is a valid completion.
I Details of the experiments
I.1 Multi-digit MNISTAdd
Like [36, 37], we take the MNIST [30] dataset and use each digit exactly once to create data. We follow [36] and require more unique digits for increasing  $N$ . Therefore, the training dataset will be of size  $60000 / 2N$  and the test dataset of size  $10000 / 2N$ .
I.1.1 Hyperparameters
We performed hyperparameter tuning on a held-out validation set by splitting the training data into 50.000 and 10.000 digits, and forming the training and validation sets from these digits. We progressively increased  $N$  from  $N = 1$ ,  $N = 3$ ,  $N = 4$  to  $N = 8$  during tuning to get improved insights into what hyperparameters are important. The most important parameter, next to learning rate, is the weight of the L2 regularization on the Dirichlet prior's parameters which should be very high. We used ADAM [25] throughout. We ran each experiment 10 times to estimate average accuracy, where each run computes 100 epochs over the training dataset. We used Nvidia RTX A4000s GPUs and 24-core AMD EPYC-2 (Rome) 7402P CPUs.