Stick-breaking Dynamic Halting
There have been several methods for imbuing models with the ability to make a prediction without having to use all layers of the model Graves ( 2016 ); Tan and Sim ( 2016 ); Dehghani et al. ( 2018 ); Elbayad et al. ( 2019 ); Schuster et al. ( 2022 ) .
Motivations for this include: (1) different inputs require different amounts of iteration to make a prediction, (2) reducing computation cost.

UT implements a similar mechanism, but the UT version of halting is difficult to interpret.
Here we choose a principled version of the dynamic halting mechanism based on the stick-breaking process, viewing it as a probability distribution.
First, \hat{\alpha}^{(t)}_{l} are the halting probabilities predicted by \mathrm{halt}(\mathbf{h}^{(t)}_{l}) , a function which is implemented by an MLP that takes in the previous layer’s embedding.
Then, the probability of any layer halting is computed by

\alpha^{(t)}_{l}=\hat{\alpha}^{(t)}_{l}\prod_{l^{\prime}=1}^{l-1}(1-\hat{\alpha}^{(t)}_{l^{\prime}})

A similar formulation is described in Graves ( 2016 ) and Tan and Sim ( 2016 ) .
Algorithm 1 shows how the mechanism is implemented at any given timestep. \mathbf{h}_{l-1} is the output of the previous layer for the current timestep.

Conditioned on the fact that we are computing \mathbf{h}^{(t)}_{l} , time-step t must not have halted before or at l-1 .
So we can use \mathbf{h}^{(t)}_{l} , the unhalted state, as input to the computation of the attention query of the block.
However, since time-step t can attend to all other timesteps, and it these other steps may have halted, we use the halted states \mathbf{S}_{l-1} for the previous layers.

However, because the halting is a ‘soft’ decision, we can relax the requirement for evaluating all possible halted states and use the expected halted state as a substitute.
Previous halting mechanisms use a ‘gating’ mechanism of convex sums between previously gated outputs and the current step’s output \mathbf{h}_{l}=\alpha_{l}\cdot\hat{\mathbf{h}}_{l}+(1-\alpha_{l})\cdot\mathbf{h}_{l-1} (Dehghani et al., 2018 ) .
This can lead to vanishingly small gradients going up the layers as (1-\alpha_{l}) multiplies.
We can instead compute the expected halted embedding at any l ,

\mathbf{s}^{(t)}_{l}=\underbrace{\left(1-\sum_{l^{\prime}=1}^{l-1}\alpha^{(t)}_{l^{\prime}}\right)\cdot\mathbf{h}^{(t)}_{l}}_{previous layer if not halted}+\underbrace{\sum_{l^{\prime}=1}^{l-1}\alpha^{(t)}_{l^{\prime}}~{}\mathbf{h}^{(t)}_{l^{\prime}}}_{halted at $<l$} (3)

If \alpha^{(t)}_{l}=1 for some l , \mathbf{s}^{(t)}_{l}=\mathbf{h}^{(t)}_{l} , recovering the behavior of the discrete halting decision.
We use \mathbf{s}^{(t)}_{l} as input to the attention key and value transformations.

This probabilistic interpretation also allows us to impose a loss on the expected number of layers used at each step, biasing the model towards fewer iterations, thereby saving computational cost.

\mathcal{L}_{ACT}=\frac{1}{T}\sum_{t=1}^{T}\sum_{l=1}^{L}\alpha^{(t)}_{l}\cdot l

We use a threshold \alpha_{\text{thresh}}=0.999 , such that the cumulative sum of the halting probabilities has exceeded this, no computation will be performed for that time step, and the previous layer’s embeddings will be copied.
Due to the routing operation required in the implementation fo SMoEs, we can simply route halted states to a “No Op” expert, leading to real savings in computation cost when halting hits the threshold early.
We find that adjusting this threshold after training can maintain performance while saving computation steps.