Question: Represent the scientific article to find a short TLDR

Question: State of the art sequence-to-sequence models for large scale tasks perform a fixed number of computations for each input sequence regardless of whether it is easy or hard to process. In this paper, we train Transformer models which can make output predictions at different stages of the network and we investigate different ways to predict how much computation is required for a particular sequence. Unlike dynamic computation in Universal Transformers, which applies the same set of layers iteratively, we apply different layers at every step to adjust both the amount of computation as well as the model capacity. On IWSLT German-English translation our approach matches the accuracy of a well tuned baseline Transformer while using less than a quarter of the decoder layers. The size of modern neural sequence models (; ;) can amount to billions of parameters . For example, the winning entry of the WMT'19 news machine translation task in English-German used an ensemble totaling two billion parameters. While large models are required to do better on hard examples, small models are likely to perform as well on easy ones, e.g., the aforementioned ensemble is probably not required to translate a short phrase such as "Thank you". However, current models apply the same amount of computation regardless of whether the input is easy or hard. In this paper, we propose Transformers which adapt the number of layers to each input in order to achieve a good speed-accuracy trade off at inference time. We extend Graves (2016; ACT) who introduced dynamic computation to recurrent neural networks in several ways: we apply different layers at each stage, we investigate a range of designs and training targets for the halting module and we explicitly supervise through simple oracles to achieve good performance on large-scale tasks. Universal Transformers (UT) rely on ACT for dynamic computation and repeatedly apply the same layer . Our work considers a variety of mechanisms to estimate the network depth and applies a different layer at each step. fix the number of steps for large-scale machine translation whereas we vary the number of steps to demonstrate substantial improvements in speed at no loss in accuracy. UT uses a layer which contains as many weights as an entire standard Transformer and this layer is applied several times which impacts speed. Our approach does not increase the size of individual layers. We also extend the resource efficient object classification work of to structured prediction where dynamic computation decisions impact future computation. Related work from computer vision includes; and who explored the idea of dynamic routing either by exiting early or by skipping layers. We encode the input sequence using a standard Transformer encoder to generate the output sequence with a varying amount of computation in the decoder network. Dynamic computation poses a challenge for self-attention because omitted layers in prior time-steps may be required in the future. We experiment with two approaches to address this and show that a simple approach works well (§2). Next, we investigate different mechanisms to control the amount of computation in the decoder network, either for the entire sequence or on a per-token basis. This includes multinomial and binomial classifiers supervised by the model likelihood or whether the argmax is already correct as well as simply thresholding the model score (§3). Experiments on IWSLT14 German-English Figure 1: Training regimes for decoder networks able to emit outputs at any layer. Aligned training optimizes all output classifiers C n simultaneously assuming all previous hidden states for the current layer are available. Mixed training samples M paths of random exits at which the model is assumed to have exited; missing previous hidden states are copied from below. translation as well as WMT'14 English-French translation show that we can match the performance of well tuned baseline models at up to 76% less computation (§4). We first present a model that can make predictions at different layers. This is known as anytime prediction for computer vision models and we extend it to structured prediction . We base our approach on the Transformer sequence-to-sequence model . Both encoder and decoder networks contain N stacked blocks where each has several sub-blocks surrounded by residual skip-connections. The first sub-block is a multi-head dot-product self-attention and the second a position-wise fully connected feed-forward network. For the decoder, there is an additional sub-block after the self-attention to add source context via another multi-head attention. Given a pair of source-target sequences (x, y), x is processed with the encoder to give representations s = (s 1, . . ., s |x|). Next, the decoder generates y step-by-step. For every new token y t input to the decoder at time t, the N decoder blocks process it to yield hidden states (h where block n is the mapping associated with the n th block and embed is a lookup table. The output distribution for predicting the next token is computed by feeding the activations of the last decoder layer h N t into a softmax normalized output classifier W : Standard Transformers have a single output classifier attached to the top of the decoder network. However, for dynamic computation we need to be able to make predictions at different stages of the network. To achieve this, we attach output classifiers C n parameterized by W n to the output h n t of each of the N decoder blocks: ∀n, p(y t+1 |h n t) = softmax(W n h n t) The classifiers can be parameterized independently or we can share the weights across the N blocks. Dynamic computation enables the model to use any of the N exit classifiers instead of just the final one. Some of our models can choose a different output classifier at each time-step which in an exponential number of possible output classifier combinations in the sequence length. We consider two possible ways to train the decoder network (Figure 1). Aligned training optimizes all classifiers simultaneously and assumes all previous hidden states required by the self-attention are available. However, at test time this is often not the case when we choose a different exit for every token which leads to misaligned states. Instead, mixed training samples several sequences of exits for a given sentence and exposes the model to hidden states from different layers. Generally, for a given output sequence y, we have a sequence of chosen exits (n 1, . . ., n |y|) and we denote the block at which we exit at time t as n t. Aligned training assumes all hidden states h n−1 1,..., h n−1 t are available in order to compute selfattention and it optimizes N loss terms, one for each exit (Figure 1a): The compound loss L dec (x, y) is a weighted average of N terms w.r.t. to (ω 1, . . . ω N). We found that uniform weights achieve better BLEU compared to other weighing schemes (c.f . Appendix A). At inference time, not all time-steps will have hidden states for the current layer since the model exited early. In this case, we simply copy the last computed state to all upper layers, similar to mixed training (§2.2.2). However, we do apply layer-specific key and value projections to the copied state. Aligned training assumes that all hidden states of the previous time-steps are available but this assumption is unrealistic since an early exit may have been chosen previously. This creates a mismatch between training and testing. Mixed training reduces the mismatch by training the model to use hidden states from different blocks of previous time-steps for self-attention. We sample M different exit sequences (n and for each one we evaluate the following loss: When n t < N, we copy the last evaluated hidden state h n t to the subsequent layers so that the self-attention of future time steps can function as usual (see Figure 1b). We present a variety of mechanisms to predict the decoder block at which the model will stop and output the next token, or when it should exit to achieve a good speed-accuracy trade-off. We consider two approaches: sequence-specific depth decodes all output tokens using the same block (§3.1) while token-specific depth determines a separate exit for each individual token (§3.2). We model the distribution of exiting at time-step t with a parametric distribution q t where q t (n) is the probability of computing block 1,..., block n and then emitting a prediction with C n. The parameters of q t are optimized to match an oracle distribution q * t with cross-entropy: The exit loss (L exit) is back-propagated to the encoder-decoder parameters. We simultaneously optimize the decoding loss (Eq.) and the exit loss (Eq.) balanced by a hyper-parameter α to ensure that the model maintains good generation accuracy. The final loss takes the form: In the following we describe for each approach how the exit distribution q t is modeled (illustrated in Figure 2) and how the oracle distribution q * t is inferred. (a) Sequence-specific depth Decoder depth Decoder depth Figure 2: Variants of the adaptive depth prediction classifiers. Sequence-specific depth uses a multinomial classifier to choose an exit for the entire output sequence based on the encoder output s (2a). It then outputs a token at this depth with classifier C n. The token-specific multinomial classifier determines the exit after the first block and proceeds up to the predicted depth before outputting the next token (2b). The token geometric-like classifier (2c) makes a binary decision after every block to dictate whether to continue (C) to the next block or to stop (S) and emit an output distribution. For sequence-specific depth, the exit distribution q and the oracle distribution q * are independent of the time-step so we drop subscript t. We condition the exit on the source sequence by feeding the average s of the encoder outputs to a multinomial classifier: where W h and b h are the weights and biases of the halting mechanism. We consider two oracles to determine which of the N blocks should be chosen. The first is based on the sequence likelihood and the second looks at an aggregate of the correctly predicted tokens at each block. This oracle is based on the likelihood of the entire sequence after each block and we optimize it with the Dirac delta centered around the exit with the highest sequence likelihood. We add a regularization term to encourage lower exits that achieve good likelihood: Correctness-based: Likelihood ignores whether the model already assigns the highest score to the correct target. Instead, this oracle chooses the lowest block that assigns the largest score to the correct prediction. For each block, we count the number of correctly predicted tokens over the sequence and choose the block with the most number of correct tokens. A regularization term controls the trade-off between speed and accuracy. Oracles based on test metrics such as BLEU are feasible but expensive to compute since we would need to decode every training sentence N times. We leave this for future work. The token-specific approach can choose a different exit at every time-step. We consider two options for the exit distribution q t at time-step t: a multinomial with a classifier conditioned on the first decoder hidden state h Multinomial q t: The most probable exit arg max q t (n|x, y <t) is selected at inference. Geometric-like q t: where, d is the dimension of the decoder states, W h ∈ R N ×d and w h ∈ R d are the weights of the halting mechanisms, and b h their biases. During inference the decoder exits when the halting signal χ n t exceeds a threshold τ n which we tune on the valid set to achieve a better accuracy-speed trade-off. If thresholds (τ n) 1≤n<N have not been exceeded, then we default to exiting at block N. The two classifiers are trained to minimize the cross-entropy with respect to either one the following oracle distributions: Likelihood-based: At each time-step t, we choose the block whose exit classifier has the highest likelihood plus a regularization term weighted by λ to encourage lower exits. This oracle ignores the impact of the current decision on the future time-steps and we therefore consider smoothing the likelihoods with an RBF kernel. where we control the size of the surrounding context with σ the kernel width. We refer to this oracle as LL(σ, λ) including the case where we only look at the likelihood of the current token with σ → 0. Correctness-based: Similar to the likelihood-based oracle we can look at the correctness of the prediction at time-step t as well as surrounding positions. We define the target q * t as follows: Confidence thresholding Finally, we consider thresholding the model predictions (§2), i.e., exit when the maximum score of the current output classifier p(y t+1 |h n t) exceeds a hyper-parameter threshold τ n. This does not require training and the thresholds τ = (τ 1, . . ., τ N −1) are simply tuned on the valid set to maximize BLEU. Concretely, for 10k iterations, we sample a sequence of thresholds τ ∼ U N −1, decode the valid set with the sampled thresholds and then evaluate the BLEU score and computational cost achieved with this choice of τ. After 10k evaluations we pick the best performing thresholds, that is τ with the highest BLEU in each cost segment. We evaluate on several benchmarks and measure tokenized BLEU : IWSLT'14 German to English (De-En). We use the setup of and train on 160K sentence pairs. We use N = 6 blocks, a feed-forward network (ffn) of intermediate-dimension Uniform n = 1 n = 2 n = 3 n = 4 n = 5 n = 6 Average WMT'14 English to French (En-Fr). We also experiment on the much larger WMT'14 EnglishFrench task comprising 35.5m training sentence pairs. We develop on 26k held out pairs and test on newstest14. The vocabulary consists of 44k joint BPE types . We use a Transformer big architecture and tie the embeddings of the encoder, the decoder and the output classifiers ((W n) 1≤n≤6; §2.1). We average the last ten checkpoints and use a beam of width 4. Models are implemented in fairseq and are trained with Adam . We train for 50k updates on 128 GPUs with a batch size of 460k tokens for WMT'14 En-Fr and on 2 GPUs with 8k tokens per batch for IWSLT'14 De-En. To stabilize training, we re-normalize the gradients if the norm exceeds g clip = 3. For models with adaptive exits, we first train without exit prediction (α = 0 in Eq.) using the aligned mode (c.f . §2.2.1) for 50k updates and then continue training with α = 0 until convergence. The exit prediction classifiers are parameterized by a single linear layer (Eq.) with the same input dimension as the embedding dimension, e.g., 1024 for a big Transformer; the output dimension is N for a multinomial classifier or one for geometric-like. We exit when χ t,n > 0.5 for geometric-like classifiers. We first compare the two training regimes for our model (§2.2). Aligned training performs selfattention on aligned states (§2.2.1) and mixed training exposes self-attention to hidden states from different blocks (§2.2.2). We compare the two training modes when choosing either a uniformly sampled exit or a fixed exit n = 1,..., 6 at inference time for every time-step. The sampled exit experiment tests the robustness to mixed hidden states and the fixed exit setup simulates an ideal setting where all previous states are available. As baselines we show six separate standard Transformers with N ∈ [1..6] decoder blocks. All models are trained with an equal number of updates and mixed training with M =6 paths is most comparable to aligned training since the number of losses per sample is identical. Table 1 shows that aligned training outperforms mixed training both for fixed exits as well as for randomly sampled exits. The latter is surprising since aligned training never exposes the self-attention mechanism to hidden states from other blocks. We suspect that this is due to the residual connections which copy features from lower blocks to subsequent layers and which are ubiquitous in Transformer models (§2). Aligned training also performs very competitively to the individual baseline models. Aligned training is conceptually simple and fast. We can process a training example with N exits in a single forward/backward pass while M passes are needed for mixed training. In the remaining paper, we use the aligned mode to train our models. Appendix A reports experiments with weighing the various output classifiers differently but we found that a uniform weighting scheme worked well. On our largest setup, WMT'14 English-French, the training time of an aligned model with six output classifiers increases only marginally by about 1% compared to a baseline with a single output classifier keeping everything else equal. Next, we train models with aligned states and compare adaptive depth classifiers in terms of BLEU as well as computational effort. We measure the latter as the average exit per output token (AE). As baselines we use again six separate standard Transformers with N ∈ [1..6] with a single output classifier. We also measure the performance of the aligned mode trained model for fixed exits n ∈ [1..6]. For the adaptive depth token-specific models (Tok), we train four combinations: likelihoodbased oracle (LL) + geometric-like, likelihood-based oracle (LL) + multinomial, correctness based oracle (C) + geometric-like and correctness-based oracle (C) + multinomial. Sequence-specific models (Seq) are trained with the correctness oracle (C) and the likelihood oracle (LL) with different values for the regularization weight λ. All parameters are tuned on the valid set and we report on the test set for a range of average exits. Figure 3 shows that the aligned model (blue line) can match the accuracy of a standard 6-block Transformer (black line) at half the number of layers (n = 3) by always exiting at the third block. The aligned model outperforms the baseline for n = 2,..., 6. For token specific halting mechanisms (Figure 3a) the geometric-like classifiers achieves a better speed-accuracy trade-off than the multinomial classifiers (filled vs. empty triangles). For geometriclike classifiers, the correctness oracle outperforms the likelihood oracle (Tok-C geometric-like vs. Tok-LL geometric-like) but the trend is less clear for multinomial classifiers. At the sequence-level, likelihood is the better oracle (Figure 3b). The rightmost Tok-C geometric-like point (σ = 0, λ = 0.1) achieves 34.73 BLEU at AE = 1.42 which corresponds to similar accuracy as the N = 6 baseline at 76% fewer decoding blocks. Figure 3). The best accuracy of the aligned model is 34.95 BLEU at exit 5 and the best comparable Tok-C geometric-like configuration achieves 34.99 BLEU at AE = 1.97, or 61% fewer decoding blocks. When fixing the budget to two decoder blocks, Tok-C geometric-like with AE = 1.97 achieves BLEU 35, a 0.64 BLEU improvement over the baseline (N = 2) and aligned which both achieve BLEU 34.35. Confidence thresholding (Figure 3c) performs very well but cannot outperform Tok-C geometriclike. In this section, we look at the effect of the two main hyperparameters on IWSLT'14 De-En: λ the regularization scale (c.f . Eq.), and the RBF kernel width σ used to smooth the scores (c.f . Eq.). We train Tok-LL Geometric-like models and evaluate them with their default thresholds (exit if χ n t > 0.5). Figure 4a shows that higher values of λ lead to lower exits. Figure 4b shows the effect of σ for two values of λ. In both curves, we see that wider kernels favor higher exits. Finally, we take the best performing models form the IWSLT benchmark and test them on the large WMT'14 English-French benchmark. Results on the test set (Figure 5a) show that adaptive depth still shows improvements but that they are diminished in this very large-scale setup. Confidence thresholding works very well and sequence-specific depth approaches improve only marginally over the baseline. Tok-LL geometric-like can match the best baseline of BLEU 43.4 (N = 6) by using only AE = 2.96 which corresponds to less than half the decoder blocks; the best aligned of BLEU 43.6 can be outmatched at (AE = 3.51, BLEU = 43.71). In this setup, Tok-LL geometric-like slightly outperforms the Tok-C counterpart. Confidence thresholding matches the accuracy of the N =6 baseline with AE 2.5 or 59% fewer decoding blocks. However, confidence thresholding requires computing the output classifier at each block to determine whether to halt or continue. This is a large overhead since output classifiers predict 44k types for this benchmark (§4.1). To better account for this, we measure the average number of FLOPs per output token (details in Appendix B). Figure 5b shows that the Tok-LL geometric-like approach provides a better trade-off when the overhead of the output classifiers is considered. Figures 7 and 6 show outputs for examples of the IWSLT'14 De-En and the WMT'14 En-Fr test sets, respectively, together with the exit and model probability for each token. Less computation is used at the end of the sentence since periods and end of sentence markers (</s>) are easy to predict. The amount of computation increases when the model is less confident e.g. in Figure 6a, predicting'présent' (meaning 'present') is hard. A straightforward translation is'était là' but the model chooses'present' which is also appropriate. In Figure 6b, the model uses more computation to predict the definite article'les' since the source has omitted the article for'passengers'. We extended anytime prediction to the structured prediction setting and introduced simple but effective methods to equip sequence models to make predictions at different points in the network. We compared a number of different mechanisms to predict the required network depth and find that a simple correctness based geometric-like classifier obtains the best trade-off between speed and accuracy. Results show that the number of decoder layers can be reduced by more than three quarters at no loss in accuracy compared to a well tuned Transformer baseline. In this section we experiment with different weights for scaling the output classifier losses. Instead of uniform weighting, we bias towards specific output classifiers by assigning higher weights to their losses. Table 2 shows that weighing the classifiers equally provides good . Uniform n = 1 n = 2 n = 3 n = 4 n = 5 n = 6 Average Uniform n = 1 n = 2 n = 3 n = 4 n = 5 n = 6 Average Table 2: Aligned training with different weights (ω n) on IWSLT De-En. For each model we report BLEU on the dev set evaluated with a uniformly sampled exit n ∼ U([1..6]) for each token and a fixed exit n ∈ [1..6] throughout the sequence. The average corresponds to the average BLEU over the fixed exits. Gradient scaling Adding intermediate supervision at different levels of the decoder in richer gradients for lower blocks compared to upper blocks. This is because earlier layers affect more loss terms in the compound loss of Eq.. To balance the gradients of each block in the decoder, we scale up the gradients of each loss term (− LL n) when it is updating the parameters of its associated block (block n with parameters θ n) and revert it back to its normal scale before back-propagating it to the previous blocks. Figure 8 and Algorithm 1 illustrate this gradient scaling procedure. The θ n are updated with γ n -amplified gradients from the block's supervision and (N −n) gradients from the subsequent blocks. We choose γ n = γ(N − n) to control the ratio γ:1 as the ratio of the block supervision to the subsequent blocks' supervisions. Table 3 shows that gradient scaling can benefit the lowest layer at the expense of higher layers. However, no scaling generally works very well. Algorithm 1 Pseudo-code for gradient scaling (illustrated for a single step t) 6: end for 7: function SCALE_GRADIENT(Tensor x, scale γ) 8: STOP_GRADIENT in PyTorch with x.detach. 10: end function Uniform n = 1 n = 2 n = 3 n = 4 n = 5 n = 6 Average Uniform n = 1 n = 2 n = 3 n = 4 n = 5 n = 6 Average Table 3: Aligned training with different gradient scaling ratios γ: 1 on IWSLT'14 De-En. For each model we report the BLEU4 score evaluated with a uniformly sampled exit n ∼ U([1..6]) for each token and a fixed exit n ∈ [1..6]. The average corresponds to the average BLEU4 of all fixed exits. This section details the computation of the FLOPS we report. The per token FLOPS are for the decoder network only since we use an encoder of the same size for all models. We breakdown the FLOPS of every operation in Algorithm 2 (blue front of the algorithmic statement). We omit non-linearities, normalizations and residual connections. The main operations we account for are dot-products and by extension matrix-vector products since those represent the vast majority of FLOPS (we assume batch size one to simplify the calculation). Table 4: FLOPS of basic operations, key parameters and variables for the FLOPS estimation. With this breakdown, the total computational cost at time-step t of a decoder block that we actually go through, denoted with FC, is: where the cost of mapping the source' keys and values is incurred the first time the block is called (flagged with FirstCall). This occurs at t = 1 for the baseline model but it is input-dependent with depth adaptive estimation and may never occur if all tokens exit early. If skipped, a block still has to compute the keys and value of the self-attention block so the selfattention of future time-steps can function. We will denote this cost with FS and we have FS = 4d Confidence thresholding: FP(t, q(t)) = 2q(t)V d d For a set of source sequences {x (i) } i∈I and generated hypotheses {y (i) } i∈I, the average flops per token is: Baseline (N blocks): Adaptive depth: q(t)FC(x (i), t) + (N − q(t))FS + FP(t, q(t)) + 2V d d In the case of confidence thresholding the final output prediction cost (2V d d) is already accounted for in the exit prediction cost FP.

represent the tldr for finding the corresponding scientific article

Sequence model that dynamically adjusts the amount of computation for each input.

We use a hypernetwork to predict optimal weights given hyperparameters, and jointly train everything together.


Question: Represent the scientific article to find a short TLDR

Question: Variational Auto-encoders (VAEs) are deep generative latent variable models consisting of two components: a generative model that captures a data distribution p(x) by transforming a distribution p(z) over latent space, and an inference model that infers likely latent codes for each data point . Recent work shows that traditional training methods tend to yield solutions that violate modeling desiderata: the learned generative model captures the observed data distribution but does so while ignoring the latent codes, ing in codes that do not represent the data (e.g. van den ;); the aggregate of the learned latent codes does not match the prior p(z). This mismatch means that the learned generative model will be unable to generate realistic data with samples from p(z)(e.g. ;). In this paper, we demonstrate that both issues stem from the fact that the global optima of the VAE training objective often correspond to undesirable solutions. Our analysis builds on two observations: the generative model is unidentifiable – there exist many generative models that explain the data equally well, each with different (and potentially unwanted) properties and bias in the VAE objective – the VAE objective may prefer generative models that explain the data poorly but have posteriors that are easy to approximate. We present a novel inference method, LiBI, mitigating the problems identified in our analysis. On synthetic datasets, we show that LiBI can learn generative models that capture the data distribution and inference models that better satisfy modeling assumptions when traditional methods struggle to do so. Introduction Variational Auto-encoders (VAEs) are deep generative latent variable models consisting of two components: a generative model that captures a data distribution p(x) by transforming a distribution p(z) over latent space, and an inference model that infers likely latent codes for each data point . Recent work shows that traditional training methods tend to yield solutions that violate modeling desiderata: the learned generative model captures the observed data distribution but does so while ignoring the latent codes, ing in codes that do not represent the data (e.g. van den ;); the aggregate of the learned latent codes does not match the prior p(z). This mismatch means that the learned generative model will be unable to generate realistic data with samples from p(z)(e.g. ;). In this paper, we demonstrate that both issues stem from the fact that the global optima of the VAE training objective often correspond to undesirable solutions. Our analysis builds on two observations: the generative model is unidentifiable -there exist many generative models that explain the data equally well, each with different (and potentially unwanted) properties and bias in the VAE objective -the VAE objective may prefer generative models that explain the data poorly but have posteriors that are easy to approximate. We present a novel inference method, LiBI, mitigating the problems identified in our analysis. On synthetic datasets, we show that LiBI can learn generative models that capture the data distribution and inference models that better satisfy modeling assumptions when traditional methods struggle to do so. Background A VAE is comprised of a generative model and an inference model. Under the generative model, we posit that the observed data and the latent codes are jointly distributed as p θ (x, z) = p θ (x|z)p(z). The likelihood p θ (x|z) is defined by a neural network f with parameters θ and an output noise model ∼ p such that x|z = f θ (z) +. Direct maximization of the expected observed data log-likelihood E p(x) log Z p θ (x, z)dz over θ is intractable. Instead, we maximize the variational lower bound (ELBO), where q η(x) ∈ Q is a variational distribution with parameters η(x). Since the bound is tight when q η(x) (z) = p θ (z|x), we aim to infer p θ (z|x). To speed up finding the variational parameters η for some new input x, we train a neural inference model g with parameters We demonstrate two general ways wherein global optima of the ELBO correspond to undesirable models. In the following, we fix our variational family to be mean-field Gaussian. Case 1: Learning the Inference Model Compromises the Quality of the Generative Model. Suppose that the variational family does not contain the posteriors of the data-generating-model. Then, often, inference must trade-off between learning a generative model that explains the data well and one that has posteriors that are easy for the inference network to approximate. Thus, the global minima of the VAE objective can specify models that both fail to capture the data distribution and whose aggregated posterior fails to match the prior. As demonstration, consider the following model (described fully in Appendix C.2): with σ 2 = 0.01, B = [0.006 0 0 0.006] and A = 0.75 0.25 1.5 −1.0 as the data generating model. Here, we fix B (which also fixes the covariance of the observation noise) and learn the parameter θ = A. In this example, the ground-truth posteriors are non-diagonal Gaussians. Here, the VAE objective can achieve a lower loss by compromising the MLE objective in order to better satisfy the PM objective -i.e. the VAE objective will prefer a model that fails to capture the data distribution but has a diagonal Gaussian posterior over the ground-truth model. Figure 1C shows the data distribution of the ground truth model θ GT, φ GT (with L(θ GT, φ GT) = 0.532) differs from the distribution of the learned model θ *, φ * in Figure 1D (with L(θ *, φ *) = 0.196). Moreover, since the learned model fails to capture the data distribution, its aggregated posterior fails to match the prior (see Figures 1E vs. 1F): Even when we restrict the class of generative models to ones that fit the data well, the posterior matching objective will still select a model with a simple posterior. Unfortunately, the selected generative model may have undesirable properties like uninformative latent codes. As demonstration, consider the model from Equation 3 with the data generating model specified by: σ 2 = 0.01, A = 0.75 0.25 1.5 −1.0, and B is some diagonal matrix with values in [0, σ 2]. In this case, we fix A and and learn the parameter θ = B. Since the observation noise covariance I ·σ 2 −B changes with B, the data marginal is fixed at p θ (x) = N 0, AA + I · σ 2 for every B. Thus, for every θ, the MLE objective is 0. However, although every choice of θ explain the data equally well, the posterior matching objective (and hence the VAE objective) is minimized by θ's whose posteriors have the least amount of correlation. Figure 1A shows that L(θ, φ) prefers high value in the upper diagonal of B and low value in the lower diagonal. Figure 1B shows the informativeness of the latent codes for the corresponding θ. We see that the data to latent code mutual information I(X; Z) corresponding to the θ selected by L(θ, φ) is not optimal. That is, even if the true data generating model produces highly informative latent codes, the VAE objective may select a model that produces uninformative latent codes. Discussion The principles of our analysis extend to non-linear VAEs and complex variational families. In the VAE objective, the posterior matching objective acts like a regularizing term, biasing the learned generative models towards simple models with posteriors that are easy to approximate (with respect to the choice of variational family). Thus, joint training of the inference and generative models introduces unintended and undesirable optima, which would not appear when these models are learned separately. Case 2: Learning the Inference Model Selects an Undesirable Generative Model. Even if the variational family is rich, the inference for the posterior can nonetheless bias the learning for the generative model. It is well known that the generative model is nonidentifiable under the MLE objective -there are many models that minimize the MLE objective. To focus on the effects of non-identifiability, let us assume that the variational family is expressive enough that it contains the posteriors of multiple models that could have generated the data. Then the posterior matching objective is 0 since we can find parameters φ such that q φ (z|x) = p θ (z|x) for any such θ. Consequently, L(θ, φ) has multiple global minima corresponding to the multiple generative models that maximizes the date likelihood. Some of these models may not satisfy our desiderata; e.g., the latent codes have low mutual information with the data. As demonstration, consider the following model (fully described in Appendix C.1): In this case, the mean-field variational family includes the posterior p θ (z|x) for all θ, i.e. the posterior matching objective can be fully minimized. Furthermore, every θ ∈ [0, σ 2] yields the same data marginal, p θ (x) = N 0, σ 2, and thus minimizes the MLE objective. However, not all choice of θ are equivalent. Given θ, the mutual information between the learned latent codes and the data is I θ (X; Z) = Const − 1 2 log(σ 2 − θ 2). Thus, the set of global minima of L(θ, φ) contain many models that produce uninformative latent codes. Discussion We've shown that posterior collapse can happen at global optima of the VAE objective and that, in these cases, collapse cannot always be mitigated by improving the inference model (as in) or by limiting the capacity of the generative model (as in ; ;). In Section 1, we showed that common problems with traditional VAE training stem from the non-identifiability of the likelihood and the bias of the VAE objective towards models with simple posteriors, even if such models cannot capture the data distribution. We propose a novel inference method to specifically target these problems. To avoid the biasing effect of the PM objective on learning the generative model, we decouple the training of the generative and inference models -first we learn a generative model, then we learn an inference model while fixing the learned generative model (note that amortization allows for efficient posterior inference). To avoid undesirable global optima of the likelihood, we learn a generative model constrained by task-specific modeling desiderata. For instance, if informative latent codes are necessary for the task, the likelihood can be constrained so that the mutual information between the data and latent codes under θ is at least δ. While there are a number of works in literature that incorporate task-specific constraints to VAE training (e.g. ; Zhao et al. (2017 Zhao et al. (, 2018 ;), adding these constraints to the VAE objective directly affects both the generative and the inference models, and, consequently, may introduce additional undesirable global optima. In our approach, added constraints only directly affects the generative model -i.e. the quality of inference cannot be compromised by the added constraints. We call our training framework Likelihood Before Inference (LiBI), and propose one possible instantiation of this framework here. Step 1: Learning the Generative Model We compute a tractable approximation to the MLE objective, constrained so that the likelihood satisfies task-specific modeling desiderata (such as high I(X; Z)) as needed.: where each c i is a constraint applied to the likelihood. We do this by computing joint maximum likelihood estimates for θ and z n while additionally constraining the z n's to have come from our assumed model (see Appendix D for a formal derivation of this approximation): where HZ(·) is the Henze-Zirkler test statistic for Gaussianity, µ(·), Σ(·) represent the empirical mean and covariance, and the z n's are amortized using a neural network z n = h(x n ; ϕ) parametrized by ϕ. These constraints encourage the generative model to capture p(x) given p(z), i.e. the aggregated posterior under this model will match the prior p(z). Step 2: Learning the Inference Model Given the θ learned in Step 1, we learn φ to compute approximate posteriors q φ (z|x):. We note that φ, too, will satisfy our modeling assumptions, since with a fixed θ, the model nonidentifiability we describe in Section 1 is no longer present. Step 3: Reinitialize Inference for the Generative Model We repeat the process, initializing h(x n ; ϕ) = µ(x n ; φ), where µ(x n ; φ) is the mean of q φ (z n |x n). This steps provides an intelligent random initialization allowing step 1 to learn a better quality model. In theory, if the generative model and the inference models are learned perfectly in Steps 1 and 2, then Step 3 is obviated. In practice, we find that Step 3 improves the quality of the generative model and only a very small number of iterations is actually needed. Discussion Using LiBI, we can now evaluate the quality of the generative model and the inference models independently. This is in contrast to traditional VAE inference, in which the ELBO entangles issues of modeling and issues of inference. On 4 synthetic data sets for which we know the data generating model, we compare LiBI with existing inference methods: VAE , β-VAE , β-VAE with annealing, Lagging inference networks . Across all datasets, LiBI learns generative models that better capture p(x) (as quantified by loglikelihood and the Smooth k-NN test statistic ) and for which the aggregated posterior better matches the prior (see Appendix B). Conclusion In this paper, we show that commonly noted issues with VAE training are attributable to the fact that global optima of the VAE training objective often includes undesirable solutions. Based on our analysis, we propose a novel training procedure, LiBI, that avoid these undesirable optima while retaining the tractability of traditional VAE inference. On synthetic datasets, we show that LiBI able to learn generative models that capture the data distribution and inference models whose aggregated posterior matches the prior while traditional methods struggle to do so. Two common issues noted in VAE literature are posterior collapse and the mismatch between aggregated posterior and prior. Posterior collapse occurs when the posterior under both the generative model and approximate posterior learned by the inference model are equal the prior p(z) . Surprisingly, under posterior collapse, the model is still able to generate samples from p data (x)(e.g. ;). This is often attributed to the fact the generative model is very powerful and is therefore able to maximize the log data marginal likelihood without the help of the auxiliary latent codes (van den). Existing literature focuses on mitigating model collapse in one of the three ways: 1. modifying the optimization procedure to bias training way from collapse ; 2. choosing variational families that make collapse less likely to occur ; 3. modifying the generative and inference model architecture to encourage more information sharing between the x's and the z's . Although much of existing literature describes issue of posterior collapse and proposes methods to avoid it, less attention has been given to explaining why it occurs. conjecture that it occurs as a of the joint training: since the likelihood changes over the course of training, it is incentivized to ignore the output of the inference network whose output in the early stages of training is not yet meaningful. Mismatch between aggregated posterior and prior refers to the case when q φ (z) = p(z), where One might expect the two distributions to match because for any given likelihood θ, one should be able to recover the prior from the true posterior p(z|x) as follows: An x produced by the generated model from a z that is likely under the prior but unlikely under the aggregate posterior may have "poor sample quality", since the the generative model is unlikely to have encountered such a z during training . Existing literature mitigate this issue by either increasing the flexibility of the prior to better fit the aggregate posterior or developing a method to sample more robustly from the latent space . Examples of the latter include training a second VAE to be able to generate z from u and then sampling from p θ (x, z) using a Gibbs sampler . In this work, we provide a unifying analysis of both posterior collapse and mismatch, showing that both can occur as global optima of the VAE objective. Through our analysis, we also show that at these optima, neither issue can be reliably resolved by existing methods. In Figures 2 and 3, we compare the posteriors learned by traditional VAE inference and by LiBI, respectively, on the synthetic dataset LinearJTEx. Here we demonstrate that traditional inference learns a generative model θ under which it is easy to approximate the corresponding posteriors. However, this comes at the cost of θ being unable to capture the data distribution. Assume the following generative process for the data: For this generative process, p θ (x) = N 0, σ 2 for any value of θ such that 0 ≤ θ ≤ σ 2. Additionally, θ directly controls I(X; Z) -when θ = 0, we have that I(X; Z); when θ = σ 2, we have that I(X; Z) = ∞. To see this, we will compute I θ (X; Z) directly (by computing p θ (x, z) and p(x)p(z)): As such, we can compute the mutual information between x and z as follows: For this model, the posterior p θ (z|x), is: Since this example is univariate, the mean-field Gaussian variational family will include the true posterior for any θ. Assume the following generative process for the data: where B is a diagonal matrix with diagonal elements between 0 and σ 2. For this generative process, p B (x) = N 0, AA + I · σ 2 for all valid values of B. For this model, the complete data likelihood and marginals are, Therefore, I B (X; Z) can be computed as follows: Lastly, the posterior for this model, p B (z|x), is a Gaussian with mean and covariance, For our choice of A, the mean-field Gaussian will not include the true posterior for this model. The best-fitting mean-field approximation to the true posterior can be computed as in Appendix C.3. Let B be a diagonal matrix and let Σ be a full-covariance matrix. where each element in the above sum is independent and is minimized when B ii = 1 Σ, and where Σ −1 ii is the ith diagonal entry of Σ −1. The LiBI Framework The LiBI framework is composed of two steps: learning a high-quality likelihood capable of generating the observed data distribution, and fixing the likelihood learned in Step 1, performing inference to learn the latent codes given the data. We emphasize that our framework is general, so one can use various existing methods for either step. For example, one can use a GAN for Step 1, and MCMC sampling for Step 2. In this section, we derive a tractable approximation to Step 1 that can be easily enhanced to include constraints for task-specific desiderata, and that is amenable to gradient-based optimization methods. wherein Equation 31, we approximate E p(z) [p θ (x n |z)] with a single sample, z n, that makes its corresponding x n most likely (this is analogous to the Empirical Bayes EB MAP Type II estimates often used to tune prior hyper-parameters). This step, however, has a problem: it is biased towards learning z n's close to 0. We will now demonstrate that this issue exists and is a of non-identifiability in the MLE estimate with respect to θ, {z n} N n=1. We then provide a solution to this problem. Characterization of Non-Identifiability in Tractable Approximation Consider the following: let Z = {z n} N n=1 be the true z's and θ used to generate the observed data, X = {x n} N n=1 in the following generative process: Now, consider, an alternative Z = {ẑ n} N n=1 andθ such that, yielding the following alternative generative process: Under these generative processes, both the data marginals and the likelihoods are equal: However, since in our model we assumed the prior is fixed p(z) = N (0, I), the alternate parameters Z,θ are preferred by the joint log-likelihood when c > 1, since log pθ(x n |ẑ n) = log p θ (x n |z n) by construction and log N (ẑ n |0, I) > log N (z n |0, I) since theẑ n's are closer to 0 when c > 1. This will cause our approximation from Equation 31 to prefer the modelθ, which generates a different data distribution that the true data distribution: Identifying the Tractable Approximation using the Henze-Zirkler Test Statistic Returning to our approximation of the MLE objective in Equation 31, we can avoid this issue by constraining the z n's to have come from the prior: We do this by constraining the z n's to be Gaussian using the Henze-Zirkler test for Gaussianity and by constraining the empirical mean and covariance of the z n's to be that of the standard normal: We hypothesize that if the likelihood function, f θ, is "smooth" and well-behaved (that is, that it maps nearby z's to nearby x's), that our approximation of the likelihood will come close to the true one. Using this framework, we first recover a high-quality likelihood (a likelihood that, unlike in the traditional VAE objective, is not compromised to match the approximate posterior). Our framework therefore naturally encourages this likelihood to satisfy modeling assumptions; that is, if we find a θ for which the x's are reconstructed accurately given Gaussian z's, the aggregated posterior under θ, p θ (z), will match the prior p(z). Given this likelihood, we can then learn a posterior that accurately approximates p θ (z|x). We note that φ, too, will satisfy our modeling assumptions, since with a fixed θ, the model non-identifiability we describe is no longer present. The LiBI Inference Method We incorporate the constraints in Equation 43 as smooth penalties into the Lagrangian in Equation 44. We additionally define h(x n ; ϕ) to be a neural network parameterized by ϕ that, given x n, returns the specific z n that generated it. ϕ allows us to amortize Equation 44. We repeat the following steps R times: 1. Step 1: θ t, ϕ t = argmin θ,ϕ − 1 N n log p θ (x n |h(x n ; ϕ)) + HZ exp HZ {h(x n ; ϕ)} N n=1 + exp Σ {h(x n ; ϕ)} N n=1 − I 2. Step 2: = argmin φ 1 N n −ELBO(θ t, φ)

Represent the TLDR for finding the corresponding scientific article

We characterize problematic global optima of the VAE objective and present a novel inference method to avoid such optima.

represent this text to find a much much longer article about it with say hundreds-thousands of words

Examples:
Provided: "End-to-end learning of invariant representations with variables across examples such as if someone went somewhere then they are there." Match: "Human reasoning involves recognising common underlying principles across many examples by utilising variables. The by-products of such reasoning are invariants that capture patterns across examples such as "if someone went somewhere then they are there" without mentioning specific people or places. Humans learn what variables are and how to use them at a young age, and the question this paper addresses is whether machines can also learn and use variables solely from examples without requiring human pre-engineering. We propose Unification Networks that incorporate soft unification into neural networks to learn variables and by doing so lift examples into invariants that can then be used to solve a given task. We evaluate our approach on four datasets to demonstrate that learning invariants captures patterns in the data and can improve performance over baselines. Humans have the ability to process symbolic knowledge and maintain symbolic thought . When reasoning, humans do not require combinatorial enumeration of examples but instead utilise invariant patterns with placeholders replacing specific entities. Symbolic cognitive models embrace this perspective with the human mind seen as an information processing system operating on formal symbols such as reading a stream of tokens in natural language. The language of thought hypothesis frames human thought as a structural construct with varying sub-components such as "X went to Y". By recognising what varies across examples, humans are capable of lifting examples into invariant principles that account for other instances. This symbolic thought with variables is learned at a young age through symbolic play . For instance a child learns that a sword can be substituted with a stick and engage in pretend play. Although variables are inherent in models of computation and symbolic formalisms, as in first-order logic , they are pre-engineered and used to solve specific tasks by means of unification or assignments that bound variables to given values. However, when learning from data only, being able to recognise when and which symbols should take on different values, i.e. symbols that can act as variables, is crucial for lifting examples into general principles that are invariant across multiple instances. Figure 1 shows the invariant learned by our approach: if someone is the same thing as someone else then they have the same colour. With this invariant, our approach can solve all of the training and test examples in task 16 of the bAbI dataset . In this paper we address the question of whether a machine can learn and use the notion of a variable, i.e. a symbol that can take on different values. For instance, given an example of the form "bernhard is a frog" the machine would learn that the token "bernhard" could be someone else and the token "frog" could be something else. If we consider unification a selection of the most appropriate value for a variable given a choice of values, we can reframe it as a form of attention. Attention models (; ;) allow neural networks to focus, attend to certain parts of the input often for the purpose of selecting a relevant portion. Since attention mechanisms are also differentiable they are often jointly learned within a task. This perspective motivates our idea of a unification mechanism that utilises attention and is therefore fully differentiable which we refer to as soft unification. Hence, we propose an end-to-end differentiable neural network approach for learning and utilising the notion of a variable that in return can lift examples into invariants used by the network to perform reasoning tasks. Specifically, we (i) propose a novel architecture capable of learning and using variables by lifting a given example through soft unification, (ii) present the empirical of our approach on four datasets and (iii) analyse the learned invariants that capture the underlying patterns present in the tasks. Our implementation using Chainer is publicly available at [link removed](anonymous link provided with submission). Reasoning with variables involves identifying what variables are, the setting in which they are used as well as the process by which they are assigned values. When the varying components, i.e. variables, of an example are identified, the remaining structure can be lifted into an invariant which then accounts for multiple other instances. Definition 1 (Variable). Given a set of symbols S, a variable X is defined as a pair X (x, s d) where s d ∈ S is the default symbol of the variable and x is a discrete random variable of which the support is S. The representation of a variable φ V (X:s d) is equal to the expected value of the corresponding random variable x given the default symbol s d: where For example, φ could be an embedding and φ V (X:s d) would become a weighted sum of symbol embeddings as in conventional attention models. The default symbol of a variable is intended to capture the variable's bound meaning following the idea of referants by. We denote variables using X, Y, A etc. such as X:bernhard where X is the name of the variable and bernhard the default symbol as shown in Figure 1. Definition 2 (Invariant). Given a structure (e.g. list, grid) G over S, an invariant is a pair I (G, ψ) where G ∈ G is the invariant example such as a tokenised story with tokens as symbols and ψ: S → is a function representing the degree to which the symbol is considered a variable. Thus, the final representation of a symbol s included in G, φ I (s) is: the linear interpolation between its representation φ(s) and its variable bound value with itself as the default symbol φ V (X:s). We adhere to the term invariant and refrain from mentioning rules, unground rules, etc. used in logicbased formalisms, e.g. Muggleton & de , since neither the invariant structure needs to be rule-like nor the variables carry logical semantics. This distinction is clarified in Section 6. Definition 3 (Unification). Given an invariant I and an example K ∈ G, unification binds the variables in I to symbols in K. Defined as a function g: I × G → G, unification binds variables by computing the probability mass functions, P in equation 1, and returns the unified representation using equation 2. The probability mass function of a variable X:s d is: where φ U: S → R d is the unifying feature of a symbol and φ U (K) ∈ R |K|×d is applied element wise to symbols in K. If g is differentiable, it is referred to as soft unification. We distinguish φ from φ U to emphasise that the unifying properties of the symbols might be different from their representations. For example, φ(bernhard) could represent a specific person whereas φ U (bernhard) the notion of someone. Overall soft unification incorporates 3 learnable components: φ, ψ, φ U which denote the base features, variableness and unifying features of a symbol respectively. Given an upstream, potentially task specific, network f: G → S, an invariant I ∈ I and an input example K ∈ G with a corresponding desired output a ∈ S, the following holds: where f now predicts based on the unified representation produced by g. In this work, we focus on g, the invariants it produces together with the interaction of f • g. Since soft unification is end-to-end differentiable, it can be incorporated into existing task-specific upstream architectures. We present 3 architectures that model f • g using multi-layer perceptrons (MLP), convolutional neural networks (CNN) and memory networks to demonstrate the flexibility of our approach. In all cases, the d dimensional representation of symbols are learnable embeddings φ(s) = O[s] T E with E ∈ R |S|×d randomly initialised by N and O[s] the one-hot encoding of the symbol. The variableness of symbols is a learnable weight ψ w (s) = σ(w s) where w ∈ R |S| and σ is the sigmoid function. We consider every symbol independently a variable irrespective of its surrounding context and leave further contextualised formulations as future work. The underlying intuition of this configuration is that a useful symbol for a correct prediction might need to take on other values for different inputs. This usefulness can be viewed as the inbound gradient to the corresponding w s parameter and ψ w (s) acting as a gate. For further model details including the size of the embeddings, please refer to Appendix A. Unification MLP (UMLP) (f : MLP, g: RNN) We combine soft unification into a multi-layer perceptron to process fixed length inputs. In this case, the structure G is a sequence of symbols with a fixed length l, e.g. a sequence of digits 4234. Given an embedded input φ(k) ∈ R l×d, the upstream MLP computes the output symbol based on the flattened representations f (φ(k)) = softmax(hE T) where h ∈ R d is the output of the last layer. However, to compute the unifying features φ U, definition 3, g uses a bi-directional GRU is the hidden state of the GRU at symbol k and W U ∈ R d×d is a learnable parameter. This model emphasises the flexibility around the boundary of f • g and that the unifying features can be computed in any differentiable manner. Unification CNN (UCNN) (f : CNN, g: CNN) Given a grid of embedded symbols φ(K) ∈ R w×h×d where w is the width and h the height, we use a convolutional neural network such that the final prediction is f (φ(K)) = softmax((W h + b)E T ) where h this time is the of global max pooling and W, b are learnable parameters. We also model g using a separate convolutional network with the same architecture as f and set φ U (k) = c 2 (relu(c 1 (k)) where c 1, c 2 are the convolutional layers. The grid is padded with 0s to obtain w × h × d after each convolution such that every symbol has a unifying feature. This model conveys how soft unification can be adapted to the specifics of the domain for example by using a convolution in a spatially structured input. Unification Memory Networks (UMN) (f : MemNN, g: RNN) Soft unification does not need to happen prior to f in a f • g fashion but can also be incorporated at any intermediate stage multiple times. To demonstrate this ability, we unify the symbols at different memory locations at each iteration of a Memory Network. Memory networks can handle a list of lists structure such as a tokenised story as shown in Figure 2. The memory network f uses the final hidden state of a bi-directional GRU (outer squares in Figure 2) as the sentence representations to compute a context attention. At each iteration, we unify the words between the attended sentences using the same approach in UMLP with another bi-directional GRU (inner diamonds in Figure 2) for unifying features φ U (bernhard) = W U Φ(bernhard). Following equation 2, the new unified representation of the memory slot is computed and f uses it to perform the next iteration. Concretely, g produces an unification tensor U ∈ R M ×m×N ×d where M and m is the number of sentences and words in the invariant respectively, and N is the number of sentences in the example such that after the context attentions are applied over M and N, we obtain φ(k) ∈ R m×d as the unified sentence at that iteration. Note that unlike in the UMLP case, the sentences can be of varying length. The prediction is then softmax(W h J I + b) where h J I is the hidden state of the invariant after J iterations. This setup, however, requires pre-training f such that the context attentions match the correct sentences. A task might contain different questions such as "Where is X?" and "Why did X go to Y?". To let the models differentiate between questions and potentially learn different invariants, we extend them with a repository of invariants I ∈ I and aggregate the predictions from each invariant. One simple approach is to sum the predictions of the invariants I∈I f • g(I, K) used in UMLP and UCNN. Another approach could be to use features from the invariants such as memory representations in the case of UMN. For UMN, we weigh the predictions using a bilinear attention η based on the hidden states at the first iteration h To initially form the repository of invariants, we use the bag-of-words representation of the questions and find the most dissimilar ones based on their cosine similarity as a heuristic to obtain varied examples. We use 4 datasets consisting of context, query and an answer (C, q, a): fixed length sequences of symbols, shapes of symbols in a grid, story based natural language reasoning with the bAbI dataset and logical reasoning represented as logic programs, examples shown in Table 1 with further samples in Appendix B. In each case we use an appropriate model: UMLP for fixed length sequences, UCNN for grid and UMN for iterative reasoning. We use synthetic datasets of which the data generating distributions are known to evaluate not only the quantitative performance but also the quality of the invariants learned by our approach. Fixed Length Sequences We generate sequences of length l = 4 with 8 unique symbols represented as digits to predict (i) a constant, (ii) the head of the sequence, (iii) the tail and (iv) the duplicate symbol. We randomly generate 1000 triples and then only take the unique ones to ensure the test split contains unseen examples. The training is then performed over a 5-fold cross-validation. Grid To spatially organise symbols, we generate a grid of size 3×3 with 8 unique symbols organised into 2 × 2 box of identical symbol, a vertical, diagonal or horizontal sequence of length 3, a cross or a plus shape and a triangle. In each task we predict (i) the identical symbol, (ii) the head of the sequence, (iii) the centre of the cross or plus and (iv) the corner of the triangle respectively. We follow the same procedure from sequences and randomly generate 1000 discarding duplicate triples. bAbI The bAbI dataset has become a standard benchmark for evaluating memory based networks. It consists of 20 synthetically generated natural language reasoning tasks (refer to for task details). We take the 1k English set and use 0.1 of the training set as validation. Each token is lower cased and considered a unique symbol. Following previous works , we take multiple word answers also to be a unique symbol in S. Logical Reasoning To demonstrate the flexibility of our approach and distinguish our notion of a variable from that used in logic based formalisms, we generate logical reasoning tasks in the form of logic programs using the procedure by. The tasks involve learning f (C, Q) = True ↔ C Q over 12 classes of logic programs exhibiting varying paradigms of logical reasoning including negation by failure . We generate 1k and 10k logic programs per task for training with 0.1 as validation and another 1k for testing. We set the arity of literals to 1 or 2 using one random character from the English alphabet for predicates and constants, e.g. p(p) and an upper case character for logical variables, e.g. p(X). We probe three aspects of soft unification: the impact of unification on performance over unseen data, the effect of multiple invariants and data efficiency. To that end, we train UMLP and UCNN with and without unification, UMN with pre-training using 1 or 3 invariants over either the entire training set or only 50 examples. Every model is trained 3 times via back-propagation using Adam on an Intel Core i7-6700 CPU using the following objective function: where L nll is the negative log-likelihood with sparsity regularisation over ψ at τ = 0.1 to discourage the models from utilising spurious number of variables. For UMLP and UCNN, we set λ K = 0, λ U = 1 for training just the unified output and the converse for the non-unifying versions. To pre-train the UMN, we start with λ K = 1, λ U = 0 for 40 epochs then set λ U = 1 to jointly train the unified output. For iterative tasks, the mean squared error between hidden states (h at each iteration j and, in the strongly supervised cases, the negative log-likelihood for the context attentions using provided supporting facts are also added to the objective function. Further details such as batch size and total number of epochs are available in Appendix C. Figure 3: Test accuracy over iterations for Unification MLP and Unification CNN models with 1 invariant versus no unification. We observe that with soft unification the models achieve higher accuracy with fewer iterations than their plain counterparts on both per task training sizes. Figure 3 portrays how soft unification generalises better to unseen examples in test sets -the same sequence or grid never appears in both the training and test sets as outlined in Section 4 -over plain models. Despite f • g having more trainable parameters than f alone, the models with unification not only maintain higher accuracy in each iteration and solve the tasks in as few as 250 iterations with ≤ 1000 training examples but also improve accuracy by ≈ 0.3 when trained with only ≤ 50 per task. We believe soft unification architecturally biases the models towards learning structural patterns which in return achieves better on recognising common patterns of symbols across examples. Results with multiple invariants are identical and the models seem to ignore the extra invariants due to the fact that the tasks can be solved with a single invariant and the regularisation applied on ψ zeroing out unnecessary invariants; further in Appendix D. The fluctuations in accuracy around iterations 750 to 1000 in UCNN are also caused by penalising ψ which forces the model to relearn the task with less variables half way through training. Following Tables 2 and 3, we observe a trend of better performance through strong supervision, more data per task and using only 1 invariant. We believe strong supervision aids with selecting the correct sentences to unify and in a weak setting the model attempts to unify arbitrary context sentences often failing to follow the iterative reasoning chain. The increase in performance with more data and strong supervision is consistent with previous work reflecting how f • g can be bounded by the efficacy of f modelled as a memory network. As a , only in the supervised case do we observe a minor improvement over MemNN by 0.7 in Table 2 and no improvement in the weak case over comparable memory based networks in Table 3 . This dependency on f also limits the ability of f • g to learn from 50 examples per task failing 17/20 and 12/12 of bAbI and logical reasoning tasks respectively. The increase in error rate with 3 invariants, we speculate, stems from having more parameters and more pathways in the model rendering training more difficult. 6 ANALYSIS After training, we can extract the learned invariants by applying a threshold on ∀s ∈ S : ψ(s) > t indicating whether a symbol is used as a variable or not. We set t = 0.0 for all datasets except for bAbI, we use t = 0.1. The magnitude of this threshold seems to depend on the amount of regularisation τ, equation 5, and the number of training steps along with batch size all controlling how much ψ is pushed towards 0. Sample invariants shown in Figure 4 describe the common patterns present in the tasks with parts that contribute towards the final answer becoming variables. Extra symbols such as is or travelled do not emerge as variables, as shown in Figure 4a; we attribute this behaviour to the fact that changing the token travelled to went does not influence the prediction but changing the action, the value of Z:left to'picked' does. However, based on random initialisation, our approach can convert an arbitrary symbol into a variable and let f compensate for the unifications it produces. For example, the invariant "X:8 5 2 2" could predict the tail of another example by unifying the head with the tail using φ U, equation 3, of those symbols. Further examples are shown in Appendix D. Pre-training f as done in UMN seems to produce more robust and consistent invariants compared to immediately training f • g since, we speculate, by equation 4 f might encourage g(I, K) ≈ K. Interpretability versus Ability A desired property of interpretable models is transparency . A novel outcome of the learned invariants in our approach is that they provide an approximation of the underlying general principle present in the data such as the structure of multi-hop reasoning shown in Figure 4e. However, certain aspects regarding the ability of the model such as how it performs temporal reasoning, are still hidden inside f. In Figure 4b, although we observe Z:morning as a variable, the overall learned invariant captures nothing about how changing the value of Z:morning alters the behaviour of f. The model might look before or after a certain time point (e) Logical reasoning task 5 with arity 1. The model captures how S:n could entail X:i in a chain. Figure 4: Invariants learned across the four datasets using the three architectures. For iterative reasoning datasets, bAbI and logical reasoning, they are taken from strongly supervised UMN. X:bill went somewhere depending what Z:morning binds to. Without the regularising term on ψ(s), we initially noticed the models using, one might call extra, symbols as variables and binding them to the same value occasionally producing unifications such as "bathroom bathroom to the bathroom" and still f predicting, perhaps unsurprisingly, the correct answer as bathroom. Hence, regularising ψ with the correct amount τ in equation 5 to reduce the capacity of unification seems critical in extracting not just any invariant but one that represents the common structure. Soft unification from equation 3 reveals three main patterns: one-to-one, one-to-many or many-toone bindings as shown in Figure 5; further examples are in Appendix D. Figure 5a captures what one might expect unification to look like where variables unify with their corresponding counterparts, e.g. X:bernhard with brian and Y:frog with lion. However, occasionally the model can optimise to use less variables and squeeze the required information into a single variable, for example by binding Y:bathroom to john and kitchen as shown in Figure 5b. We believe this occurs due to the sparsity constraint on ψ(s) encouraging the model to be as conservative as possible. Finally, if there are more variables than needed as in Figure 5c, we observe a many-to-one binding with Y:w and Z:e mapping to the same constant q. This behaviour begs the question how does the model differentiate between p(q) and p(q, q). We speculate the model uses the magnitude of ψ(w) = 0.037 and ψ(e) = 0.042 to encode the difference despite both variables unifying with the same constant. Learning an underlying general principle in the form of an invariant is often the means for arguing generalisation in neural networks. For example, Neural Turing Machines are tested on previously unseen sequences to support the view that the model might have captured the underlying pattern or algorithm. In fact, claim "MemNNs can discover simple linguistic patterns based on verbal forms such as (X, dropped, Y), (X, took, Y) or (X, journeyed to, Y) and can successfully generalise the meaning of their instantiations." However, this claim is based on the output of f and unfortunately it is unknown whether the model has truly learned such a representation or indeed is utilising it. Our approach sheds light to this ambiguity and presents these linguistic patterns explicitly as invariants ensuring their utility through g without solely analysing the output of f on previously unseen symbols. Although we associate these invariants with our existing understanding of the task to mistakenly anthropomorphise the machine, for example by thinking it has learned X:mary as someone, it is important to acknowledge that these are just symbolic patterns. In these cases, our interpretations do not necessarily correspond to any understanding of the machine, relating to the Chinese room argument made by. Learning invariants by lifting ground examples is related to least common generalisation by which inductive inference is performed on facts such as generalising went(mary,kitchen) and went(john,garden) to went(X,Y). Unlike in a predicate logic setting, our approach allows for soft alignment and therefore generalisation between varying length sequences. Existing neuro-symbolic systems focus on inducing rules that adhere to given logical semantics of what variables and rules are. For example, δILP constructs a network by rigidly following the given semantics of first-order logic. Similarly, Lifted Relational Neural Networks ground first-order logic rules into a neural network while Neural Theorem Provers build neural networks using backwardchaining on a given knowledge base with templates. However, the notion of a variable is pre-defined rather than learned with a focus on presenting a practical approach to solving certain problems, whereas our motivation stems from a cognitive perspective. At first it may seem the learned invariants, Section 6, make the model more interpretable; however, this transparency is not of the model f but of the data. The invariant captures patterns that potentially approximates the data generating distribution but we still do not know how the model f uses them upstream. Thus, from the perspective of explainable artificial intelligence (XAI) , learning invariants or interpreting them do not constitute an explanation of the reasoning model f even though "if someone goes somewhere then they are there" might look like one. Instead, it can be perceived as causal attribution in which someone being somewhere is attributed to them going there. This perspective also relates to gradient based model explanation methods such as Layer-Wise Relevance Propagation and Grad-CAM . Consequently, a possible view on ψ, equation 2, is a gradient based usefulness measure such that a symbol utilised upstream by f to determine the answer becomes a variable similar to how a group of pixels in an image contribute more to its classification. Finally, one can argue that our model maintains a form of counterfactual thinking in which soft unification g creates counterfactuals on the invariant example to alter the output of f towards the desired answer, equation 4. The question where Mary would have been if Mary had gone to the garden instead of the kitchen is the process by which an invariant is learned through multiple examples during training. This view relates to methods of causal inference in which counterfactuals are vital as demonstrated in structured models by. We presented a new approach for learning variables and lifting examples into invariants through the usage of soft unification. Evaluating on four datasets, we analysed how Unification Networks perform comparatively to existing similar architectures while having the benefit of lifting examples into invariants that capture underlying patterns present in the tasks. Since our approach is end-toend differentiable, we plan to apply this technique to multi-modal tasks in order to yield multi-modal invariants for example in visual question answering. A MODEL DETAILS Unification MLP (UMLP) To model f as a multi-layer perceptron, we take symbol embeddings of size d = 16 and flatten sequences of length l = 4 into an input vector of size φ(k) ∈ R 64. The MLP consists of 2 hidden layers with tanh non-linearity of sizes 2d and d respectively. To process the query, we concatenate the one-hot encoding of the task id to φ(k) yielding a final input of size 64 + 4 = 68. For unification features φ U, we use a bi-directional GRU with hidden size d and an initial state of 0. The hidden state at each symbol is taken with a linear transformation to give φ U (s) = W U Φ(s) where Φ(s) is the hidden state of the biGRU. The variable assignment is then computed as an attention over the according to equation 3. Unification CNN (UCNN) We take symbols embeddings of size d = 32 to obtain an input grid φ(K) ∈ R 3×3×32. Similar to UMLP, for each symbol we append the task id as a one-hot vector to get an input of shape 3 × 3 × (32 + 4). Then f consists of 2 convolutional layers with d filters each, kernel size of 3 and stride 1. We use relu non-linearity in between the layers. We pad the grid with 2 columns and 2 rows to a 5 × 5 such that the output of the convolutions yield again a hidden output H ∈ R 3×3×d of the same shape. As the final hidden output h, we take a global max pool to over H to obtain h ∈ R d. Unification function g is modelled identical to f without the max pooling such that φ U (K ij) = H ij where H is the hidden output of the convolutional layers. Unlike previous architectures, with UMN we interleave g into f. We use embedding sizes of d = 32 and model f with an iterative memory network. We take the final hidden state of a bi-directional GRU, with initial state 0, Φ M to represent the sentences of the context C and query q in a ddimensional vector M i = Φ M (C i) and the query m q = Φ M (q). The initial state of the memory network is h 0 = m q. At each iteration j: where Φ A is another d-dimensional bi-directional GRU and ρ(x, y) = [x; y; x y; (x − y) 2 ] with the element-wise multiplication and [;] the concatenation of vectors. Taking β j as the context attention, we obtain the next state of the memory network: and iterate J many times in advance. The final prediction becomes f (C, q) = softmax(W h J + b). All weight matrices W and bias vectors b are independent but are tied across iterations. B GENERATED DATASET SAMPLES Table 5: Training sizes for randomly generated fixed length sequences and grid tasks with 8 unique symbols. The reason for Grid task (i) to be smaller is because there are at most 32 combinations of 2 × 2 boxes in a 3 × 3 grid with 8 unique symbols. Task Sequences Grid i 704.7 ± 12.8 25.6 ± 1.8 ii 709.4 ± 13.8 623.7 ± 14.1 iii 709.7 ± 14.0 768.2 ± 12.5 iv 624.8 ± 12.4 795.2 ± 10.3 C TRAINING DETAILS C.1 UNIFICATION MLP & CNN Both unification models are trained on a 5-fold cross-validation over the generated datasets for 2000 iterations with a batch size of 64. We don't use any weight decay and save the training and test accuracies every 10 iterations, as presented in Figure 3. We again use a batch size of 64 and pre-train f for 40 epochs then f together with g for 260 epochs. We use epochs for UMN since the dataset sizes are fixed. To learn g alongside f, we combine error signals from the unification of the invariant and the example. Following equation 4, the objective function not only incorporates the negative log-likelihood L nll of the answer but also the mean squared error between intermediate states h j I and h j K at each iteration as an auxiliary loss: We pre-train by setting λ U = 0 for 40 epochs and then set λ U = 1. For strong supervision we also compute the negative log-likelihood L nll for the context attention β j, described in Appendix A, at each iteration using the supporting facts of the tasks. We apply a dropout of 0.1 for all recurrent neural networks used and only for the bAbI dataset weight decay with 0.001 as the coefficient. Figure 6: Results of Unification MLP and CNN on increasing number of invariants. There is no impact on performance when more invariants per task are given. Upon closer inspection, we noticed the models ignore the extra invariants and only use 1. We speculate the regularisation ψ encourages the models to use a single 1 invariant. Xsandra went back to the Y:bathroom is X:sandra in the Y:bathroom yes Figure 9: Invariants learned that do not match the data generating distribution from UMLP and UCNN using ≤ 1000 examples to train. In these instances the unification still bind to the the correct symbols in order to predict the desired answer; quantitatively we get the same . Variable default symbols are omitted for clarity. (a) bAbI task 2. When a variable is unused in the next iteration, e.g. Z:football, it unifies with random tokens often biased by position. (b) Logical reasoning task 1. A one-to-one alignment is created between predicates and constants. (c) Logical reasoning task 3. Arity 1 atom forced to bind with arity 2 creates a one-to-many mapping."

Neural network regression should use Dirichlet output distribution when targets are probabilities in order to quantify uncertainty of predictions.