RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Overcoming VT limitations with UT
Dziri et al. ( 2023 ) and Liu et al. ( 2022 ) find that Vanilla Transformers learn shortcuts for tasks that require multi-step compositional operations, and fail to generalize on larger instances of the problem that require more steps.
Theoretical results have also shown that Vanilla Transformers have limitations in what they can compute that support these findings (Hahn, 2020 ; Hao et al., 2022 ; Merrill et al., 2022 ) .
Universal Transformers (Dehghani et al., 2018 ) are Transformers with tied weights across all layers, and an additional halting mechanism to decide when to stop.
In an ideal scenario of infinite layers (now that all layers have the same parameters) the UT, like the Neural GPU (Kaiser and Sutskever, 2015 ) , is Turing-complete, which overcomes many of the abovementioned issues.

In practice, even with limited depth, UTs have exhibited properties that afford them better performance in compositional generalization tasks (Csordás et al., 2021 ) .
UTs allow operations learned in the Transformer during training to be depth-order invariant.
If some operations during training are learned to be performed in a certain order, during test time, the UT could generalize to an unseen order of operations.
Challenges with Scaling the UT
Despite these compositional abilities, performance tends to decrease on real-world tasks when using UTs.
ALBERT Lan et al. ( 2019 ) improved parameter efficiency by sharing parameters across layers.
This was motivated by an observation that Transformers tend to learn to perform similar operations in the layers, and that sharing these parameters would reduce this redundancy 1 1 1 https://ai.googleblog.com/2019/12/albert-lite-bert-for-self-supervised.html .
However, the authors observe a dip in performance when sharing parameters, contrary to Dehghani et al. ( 2018 ) .

Could the issue be one of model capacity?
Experiments with ALBERT show that scaling up ALBERT can outperform the BERT baseline, even on real-world tasks (Lan et al., 2019 ) . Kaplan et al. ( 2020 ) also show that a shared-parameter Transformer has better scaling properties in terms of parameter-to-performance, but poorer properties in terms of computation-to-performance, since parameter count causes the computation to increase. Tay et al. ( 2022 ) scale up different sequence models, and remark on difficulties with scaling up UTs, limiting the experiments they can perform on UT. Takase and Kiyono ( 2021 ) outline several strategies of scaling up shared-parameter Transformers to deal with these issues by using different parameter-sharing schemes.

Our experiments show that SMoE techniques can be applied successfully to the UT to scale it up, achieving the UT’s parameter efficiency while not incurring the same computation costs.
We also perform experiments that support the compositional generalization claims of prior work, and provide better baselines for those tasks.
Method
Like UT, we reuse the same SUT block for every layer of the Transformer.
Within each SUT block, we use SMoEs to achieve sparsity for the feed-forward network (FFD) and attention heads separately.
We use the Mutual Information Maximization loss proposed in Chen et al. ( 2022 ) and modified for unsupervised tasks in Shen et al. ( 2023 ) .
Finally, we propose a stick-breaking process formulation of dynamic halting, which affects how the attention mechanism works in the SUT, and the Adaptive Computation Time (ACT) auxiliary loss we use to minimize the number of layers used.
Sparse Mixture of Experts
A Mixture of Experts module consists of E sub-modules f_{1},\ldots,f_{E} .
There is also a gating network, which we will denote by g(e\mid\mathbf{h}) – for any input \mathbf{h} to the MoE module, the gating network would predict a distribution over the E experts.
When k<E , we refer to this as a Sparse Mixture of Experts (SMoE), and g(e\mid\mathbf{h})>0 for only k experts, while maintaining that \sum_{e}^{E}g(e\mid\mathbf{h})=1 .
The final output of the SMoE is then given by y=\sum_{e=1}^{E}g(e\mid\mathbf{h})\cdot f_{e}(\mathbf{h}) , where g(e\mid\mathbf{h})=0 , f_{e}(\mathbf{h}) will not need to be evaluated, reducing computation cost during training and inference.
We replace the Feed-forward layer (FFD) in the Transformer block with a mixture of FFDs. Each Mixture of FFD can be described with a 3-tuple, (E,k,D) : E experts, k for the number of experts to be used in the top- k operation, and D for the dimension of the hidden layer for each FFD expert.
For the attention layer, we use the Mixture of Multi-head Attention (MoMHA) proposed by Zhang et al. ( 2022 ) and Shen et al. ( 2023 ) .
Each MoMHA module can be described by a 5-tuple, (E,k,H,D,W) ,
with E representing the number of experts, K representing the parameter k in a top- k operation, H representing the number of attention heads per expert, and D for the dimensions per head.
Like in MoA, MoMHA maintains only a single set of H key-value projections shared among the experts, while there are E query projections of H heads each. W represents the relative position embedding window size, parameterizing 2W+1 embeddings for W positions before and after the present position.
Figure 3 ( Left ) shows the schematic of a SUT block.

This technique has been used to reduce computation costs both during training and inference time for large models.
Mutual Information Maximisation
Like other models that rely on conditional activation, auxiliary losses are needed in order to aid learning a module that decides which experts are activated, and to ensure that all experts are used, balancing the load for processing.
For this, we use the Mutual Information Maximization introduced in Chen et al. ( 2022 ) for the auxiliary loss (to be maximised):

\mathcal{L}_{MIM}= \underbrace{\sum_{e=1}^{E}g(e)\log g(e)}_{-H(e)} \underbrace{-\frac{1}{|\mathcal{X}|}\sum_{\mathbf{h}\in\mathcal{X}}\sum_{e=1}^{E}g(e\mid\mathbf{h})\log g(e\mid\mathbf{h})}_{H(e\mid\mathbf{h})}, (1)

where,

g(e)=\frac{1}{|\mathcal{X}|}\sum_{\mathbf{h}\in\mathcal{X}}g(e|\mathbf{h})

.
Specifically, we use the unsupervised version proposed by Shen et al. ( 2023 ) that assumes a uniform distribution over all tokens and layers, resulting in the following auxiliary objective.
In the SUT setting, the gating network is used |\mathcal{X}|=L\cdot T times, where L is the number of layers, and T is the number of timesteps.

Intuitively, the entropy term increases the entropy of the marginal probability of the gating network predictions, which at its maximum means that the weight for each gating network across the entire minibatch is uniform.
The conditional entropy term decreases the conditional entropy, which causes the prediction of the gating network to be sharp, and also penalizes the uniform distribution solution for the gating network.
Experiments
First, we show that we can scale the UT with SUT on the WMT’14 English-German (Bojar et al., 2014 ) translation task.
We then ran experiments on Compositional Freebase Questions (CFQ; Keysers et al. 2019 ) to test for compositional generalization properties.
To further analyze the behaviour of the model under compositional generalization settings, we test our model on the Logical inference task from Bowman et al. ( 2015 ) .
All experiments were implemented within the Fairseq framework (Ott et al., 2019 ) .
English to German Translation
We perform experiments on the WMT’14 English-German translation dataset (Bojar et al., 2014 ) .
We use the pre-processing from Liu et al. ( 2020 ) .
We use a joined dictionary and share all word embeddings of the encoder and decoder.
For evaluation, we average the last 5 best models according to their negative log-likelihood scores.
We report the BLEU scores (Papineni et al., 2002 ) , and also report the MACs (Multiply-Accumulate Operations) to evaluate the runtime computational costs of the different models.
MACs of previous models were computed in Zhang et al. ( 2022 ) .

The results are reported in Table 1 .
We compare against strong baselines while accounting for the number of parameters in these models.
In addition, we train two UTs by setting E=1,k=1 , and parameterizing the FFD and Attention layers with parameters to match our \sim 65M, and \sim 110M setting for SUT.
The SUTs and UTs both demonstrate good parameter efficiency when compared to previous models.
In the \sim 110M parameter class, SUT and UT perform at around 29.4 and 29.6 BLEU respectively, while previous models require \sim 200M parameters.
While the SUT does not perform as well as the UT, but the computations required during runtime could be as low as one-fifth of UT.
Also, because we keep k constant for SUT, the MACs stays constant as SUT scales up.

We ran experiments removing different aspects of the model and its training process, including: MIM auxiliary loss, Mixture of MHA, the ACT loss, and the halting mechanism.
The results are in Table 2 .
The introduction of multiple heads to the MoA was crucial in seeing performance gains on this task, as well as having the MIM loss as a load-balancing auxiliary objective.
Interestingly, halting does contribute as much of a performance gain as it does in CFQ.

Additionally, we compute the top 5 tokens that occur in conjunction with each expert, regardless of layers, and find that certain associations exist.
We pick several experts in Table 3 that show a clear sign of co-occurring with tokens that seem to show a pattern.
This suggests that while there may be redundancy between the experts, groups of experts can specialize on certain tasks, resulting in some modularity.
Future work can investigate if such modularity can result in more robust generalization.
Compositional Freebase Questions
We run experiments on the Compositional Freebase Questions (CFQ; Keysers et al. 2019 ) dataset to determine the compositional generalization abilities of the SUT.
This is a translation task from natural language to a SPARQL query.
As an example, the sequence Who wrote M1 and wrote a film would be translated to the target sequence SELECT DISTINCT ?x0 WHERE { ?x0 a people.person . ?x0 film.writer.film ?x1 M1 . ?x1 a film.film } .
CFQ tests for compositional generalization using the notion of compound divergence , which measures how different the training set and test set are in terms of combinations of tokens, which they refer to as compounds.
To our knowledge, the current best-performing models either finetune a pretrained language model or, use knowledge about the task to design a suitable prompt for a large language model (Drozdov et al., 2022 ) .
While the prompting approach is extremely effective at the CFQ task, we view the task as a benchmark for compositional generalization in general and should be viewed in concert with other experiments, especially real-world data (like translation).
When using domain knowledge of the task in the prompt, the results may indicate better performance with a specific approach for CFQ (and perhaps other SQL translation tasks) but might be difficult to extrapolate to other settings.

In our experiments, we use preprocessing scripts from Zheng and Lapata ( 2021 ) .
The scripts perform preprocessing to the target sequence that simplifies the target sequence the same way performed in Furrer et al. ( 2020 ) .
Accordingly, we train a baseline Transformer on the transformed target.
We performed a search on the SUT hyperparameters, using the MCD1 validation set, and the best-performing set of parameters are
Attention (E=1,k=1,H=8,D=64,W=1) and FFD (E=1,k=1,D=1024) , which corresponds to the UT setting.
Refer to Appendix A for further details.
Since CFQ is a relatively small task, larger scale is not a factor and might suggest that expert specialization may not be as helpful.
The results are shown in Table 4 .
In cases with and without halting, the model already outperforms previous benchmarks, including the UT baseline from Bergen et al. ( 2021 ) .
For a fairer comparison, we use the same hyperparameters as our UT implementation, we modify our UT implementation to be more similar to the T5-based UT in Bergen et al. ( 2021 ) .
These changes include: the bucketed relative position bias used by T5, and going from post layer-norm to pre layer-norm.
While this results in much improved results compared to the original paper, our implementation of UT still outperforms it.

The Dangle (Zheng and Lapata, 2021 ) model, which beats our model, also requires re-running the encoder for every token decoded.
This is an expensive process, but given that both our method and Dangle perform well at this task, is additional evidence that iterative processes are beneficial for compositional generalization.
Logical Inference
We use the logical inference task from Bowman et al. ( 2015 ) as a test bench for UT.
Despite the apparent simplicity of the language, the task inherently requires the model to learn the hierarchical structure of the problem.
Each instance of the task comprises of two logical statements, and the goal is to predict if the statements are equivalent, contradictory, disjoint, or entail in either direction.
For example, given s_{1}= a and s_{2}= a ( or b ) , then s_{1}\sqsubset s_{2} .
The crux of the task is in training the model on sequences that have 0-6 logical operators and evaluating it on sequences that have 7-12 operators.
Given our sequence-to-sequence setting, we convert the task into a translation task.
The model takes sentence1 #SEP# sentence2 as its source sentence, with the target sentence being the single-token label for that pair.

We train a 12 layer model with Attention (E=12,k=4,H=2,D=32,W=1) and FFD (E=12,K=4,D=128) and halting. Refer to Appendix A for further details.
Training a 12-layer Vanilla Transformer achieves approximately the same results as in Shen et al. ( 2019 ) , so we report their results.
Our results in Table 5 confirm the findings of Tran et al. ( 2018 ) , showing that with recurrence in SUTs, we are able to generalize to longer sequences of the task.
While there are other models that induce a tree structure that performs exceedingly well on the task, we wanted to evaluate our model against other popular architectures.
The LSTM is a strong baseline, and we find that UT outperforms it in generalization.
We also evaluate UTs on the compositional generalization splits as proposed in Shen et al. ( 2019 ) , where the splits A, B, and C are in increasing difficulty.
The results show that UTs are able to generalize better for the A and B splits, outperforming the LSTM and VT.
Split C is still presents a challenge for the Transformer variants.

Additionally, we compute the average halting depth for the test data segmented by operator counts.
Because more operators require more nesting of expressions in the sequences, more recursion is required to properly parse the sequence.
As expected, in Figure 4 , the average halting depth increases as more operators are used.
The operator count for these clauses are correlated with length, which suggests that SUTs may be suited to generalize for length.
We include further experiments on length generalization in the Appendix Table 8 .
Post-training Computation Reduction
Does lowering \alpha_{\text{thresh}} after training cause the model to halt earlier, saving computation? How much would that cost us in terms of accuracy?

We estimate the skipped SUT block computations given different values of \alpha_{\text{thresh}}\in\{0.1,0.2,\ldots,0.9\} by looking at the halting patterns of the decoder given the ground truth source-target pairs.
We pass the source-target pair into the model and analyze the halting patterns of the model, giving us a rough estimate of how much computation would be saved as a percentage of computing all layers of the SUT.
Logical Inference
We observe the resulting performance on the hardest split of the test set with 12 operations.
Due to the already saturated halting pattern, the halting probability \alpha_{l} spikes rapidly from close to 0 to higher values,
resulting in a near constant \sim 50% reduction of the computation time regardless of the threshold.
CFQ
Using the MCD1 test split of the dataset, and our best-performing model on MCD1, we perform the \alpha_{\text{thresh}} adjustment.
The halting patterns reflect the repeated structure of SQL, using fewer steps for ‘.‘ and ‘WHERE‘, while the main bulk of the region within {...} requires more SUT steps before halting.
Surprisingly, when 0.8\leq\alpha_{\mathrm{thresh}}\leq 0.999 , the accuracy remains fairly constant.
An estimated 33% of the computation steps were skipped at \alpha_{\mathrm{thresh}}=0.8 .
At \alpha_{\text{thresh}}=0.1 , there is a slight increase in the number of computed steps, which is possible since halting earlier will result in different embeddings, and result in different halting decisions in other timesteps.
Overall, the results suggest that we can save about 20% of the SUT computation steps without any drop in accuracy, and about \sim 50% for a 0.2% decrease.
English-German Translation
For this larger dataset, we find that these translation models halt much later, suggesting that the translation task requires more computational steps than the 6-layer SUT we used.
However, further increasing the number of layers to 12 layers does not bring about much benefit, as evidenced by the halting in Figure 4 , which is an example of the halting mechanism using nearly all layers.
For comparison, Admin 60L-12L model, requires a 60-layer encoder to achieve its performance.
Even when \alpha_{\text{thresh}}=1 , the skipped computation steps remain at about 33%, compared to 80% in the CFQ task.
We find that we can reduce the computation by 9% while still retaining a BLEU score of 29.1.
Conclusion
We show that it is possible to scale up the UT via SUTs, and SUTs outperforms models of the same capacity in the WMT’14 English-to-German translation task.
The recursive nature of both UTs and SUTs allows for better inductive biases, which we have demonstrated in synthetic tasks like CFQ and logical inference.
VTs have been shown to be poor at these compositional generalization tasks without additional domain knowledge.
The stick-breaking dynamic halting mechanism also allows post-training adjustment of computation cost, which is a boon for deployment at scale.
Limitations
While the experiments in this paper show the desirable generalization properties of UTs, there are some aspects of compositional generalization that SUTs do not solve.
Importantly, while we demonstrate scaling UTs up via SMoEs, further experiments on larger settings are needed to ascertain viability in large scale systems.
Other issues may also crop up in the further scaling of SUTs, but we believe there is ample literature to draw on to finding solutions for these problems.
Figure 1: A VT has separate Transformer blocks for each layer, with different parameters.
For a UT with the same number of parameters, the UT block will be \sim 3 times the dimensions of each VT block. Running this block for 3 layers would then incur approximately 9 times the runtime memory.
Using SMoEs can recover approximately the same computational cost as the VT.; A VT has separate Transformer blocks for each layer, with different parameters.
For a UT with the same number of parameters, the UT block will be \sim 3 times the dimensions of each VT block. Running this block for 3 layers would then incur approximately 9 times the runtime memory.
Using SMoEs can recover approximately the same computational cost as the VT.
Figure 2: Example of the compositional generalization splits from Shen et al. ( 2019 ) .
The combination of not and and are never seen in successive combination during training, and a VT may learn a shortcut that prevents generalisation during test.; Example of the compositional generalization splits from Shen et al. ( 2019 ) .
The combination of not and and are never seen in successive combination during training, and a VT may learn a shortcut that prevents generalisation during test.
Figure 3: Left : Schematic of a SUT block. Right : While the input of each SUT block is the output of the previous layer, the attention mechanism attends to the halted state of the timestep. When the halting probability exceeds \alpha_{\text{thresh}} , the hidden state is simply copied. Finally, the halted state is used as the output of the SUT.; Left : Schematic of a SUT block. Right : While the input of each SUT block is the output of the previous layer, the attention mechanism attends to the halted state of the timestep. When the halting probability exceeds \alpha_{\text{thresh}} , the hidden state is simply copied. Finally, the halted state is used as the output of the SUT.
Figure 4: The average dynamic halting depth of the UT model as the number of operators increases in the test set.
The model learns to think more when the problem is harder.; The average dynamic halting depth of the UT model as the number of operators increases in the test set.
The model learns to think more when the problem is harder.
Figure 5: Above: Plot of 1-\sum_{l^{\prime}=1}^{l-1}\alpha^{(t)}_{l^{\prime}} , for an example Logical Inference input — x -axis: timesteps, y -axis: layers.
This visualizes the halting pattern of the model: dark blue represents halted, while yellow represents active. Below: Efficiency vs. Performance tradeoff curves when \alpha_{\text{thresh}} is adjusted.; Above: Plot of 1-\sum_{l^{\prime}=1}^{l-1}\alpha^{(t)}_{l^{\prime}} , for an example Logical Inference input — x -axis: timesteps, y -axis: layers.
This visualizes the halting pattern of the model: dark blue represents halted, while yellow represents active. Below: Efficiency vs. Performance tradeoff curves when \alpha_{\text{thresh}} is adjusted.
Figure 6: Halting plot and trade-off curves for CFQ. (See Figure 5 for description); Halting plot and trade-off curves for CFQ. (See Figure 5 for description)
Figure 7: Halting plot and trade-off curves for English-German Translation. (See Figure 5 for description); Halting plot and trade-off curves for English-German Translation. (See Figure 5 for description)
Figure 8: Expert-token affinity for the logical inference task. We plot the PMI scores between the tokens ( y -axis) and the experts ( x -axis).
Looking at the MoMHA experts, we see some patterns that suggest the experts have specialized in their tasks.
For example, the or operator is dominating expert 9.
Interestingly, the not operator is reliant on experts 4 and 5, unlike the or and and operators, which are binary and rely on experts 1 and 2.
The variables a through f seem to have similar expert usage patterns as well.
This suggests that the experts are specialized in their operation.; Expert-token affinity for the logical inference task. We plot the PMI scores between the tokens ( y -axis) and the experts ( x -axis).
Looking at the MoMHA experts, we see some patterns that suggest the experts have specialized in their tasks.
For example, the or operator is dominating expert 9.
Interestingly, the not operator is reliant on experts 4 and 5, unlike the or and and operators, which are binary and rely on experts 1 and 2.
The variables a through f seem to have similar expert usage patterns as well.
This suggests that the experts are specialized in their operation.