Bayesian Neural Network Language Modeling for Speech Recognition

In conventional LSTM-RNN language models, the word probability is computed as:

where $w_{0:t-1}$ denotes the word sequence $(w_0,\dots ,w_{t-1})$ and $h_{t-1}\in R^D$ is the hidden state that encodes the previous word sequence $(w_0,\dots ,w_{t-2})$ into a continuous vector representation, and D is the number of hidden nodes. In NNLMs, the most recent history $w_{t-1}$ is represented by a one-hot vector ${\tilde{w}}_{t-1}\in R^N$ as the input, where N is the vocabulary size. Generally, the LSTM-RNN LM contains three parts: the word embedding layer, the recurrent layer and the output layer. In order to address data sparsity, the word embedding layer projects the one-hot word ${\tilde{w}}_t$ as the input vector into a neural network. After being forward fed through non-linear transformations contained in multiple hidden layers, the one-hot word ${\tilde{w}}_t$ is then projected into a continuous space $x_t\in R^M$, where M is the embedding size and usually $M\ll N$:

where ${\Theta }_U\in R^{M\times N}$ is a projection matrix that can be learned during training process before being further fed into the hidden layers. After word embedding, the hidden state $h_t$ is obtained within the memory cell in the LSTM architecture [18] where the information of previous hidden state $h_{t-1}$ is combined with the word $w_{t-1}$. As shown in Fig.1 (a), at time step $t$ the respective outputs from the gates are obtained, including the forget gate $f_t$, the input gate $i_t$, the cell input ${\tilde{c}}_t$ and the output gate $o_t$. These gating functions are used to control the information flow within the cells to store the historical contexts over longer time steps and address the vanishing gradient issue with the additive connection in Equation (8) found in conventional RNN LMs. The respective outputs from these gates are calculated by

where $\sigma$ denotes the Sigmoid activation function and ${\Theta }_f,{\Theta }_i,{\Theta }_c,{\Theta }_o\in R^{D\times (M+D+1)}$ denote the weight parameter matrices associated with the above gates respectively. The final hidden state is normally computed recursively from left to right in uni-directional LSTM-RNN LMs [7, 8]. Given the gating outputs, the hidden state $h_t$ and cell memory $c_t$ are:

where $\circ$ is the Hadamard product. Finally, the output layer uses the hidden state vector $h_t$ to compute the word probabilities via a Softmax function:

where ${\Theta }_v^{\left(w\right)}\in R^{1\times D}$ is the weight vector for word $w$ in the output layer, and $V$ represents the vocabulary.

The Transformer model architecture considered in this paper features a deep stacking of multiple Transformer decoder blocks. As shown in Fig.2 (a), each Transformer decoder block consists of a multi-head self-attention [21, 22, 23] and feed forward neural network modules in each block. Residual connections [24] and layer normalization [25] are also inserted between the two modules. Assume that $x_t^{l-1}\in R^M$ represents the output of the $(l-1)$-th Transformer layer at time step $t$, where $M$ is the dimensionality of word embedding. In the $l$-th layer, the multi-head self-attention module transforms $x_t^{l-1}$ to $z_t^l$ as follows:

where ${\Theta }_q^l,{\Theta }_k^l,{\Theta }_v^l\in R^{M\times (M+1)}$ are projection matrices of $l$-th self-attention module that maps the input $x_t^{l-1}$ into query $q_t^l$, key $k_t^l$ and value $v_t^l$ respectively. $\text{bf{Attn}}(\cdot )$ represents the scaled multi-head dot product self-attention mechanism [19]. Only self-attention mechanism with a single head is used in Equation (12). $y_t^l$ is the sequence of cached key-value vector pairs up to time step $t$, which can be restricted to only contain the history context information and thus prevent the model from using any future context, and $\frac{1}{\sqrt{M}}$ denotes the scaling factor. $\text{bf LayerNorm}(\cdot )$ is the layer normalization operation [25]. ${\Theta }_h^l\in R^{M\times (M+1)}$ is the projection matrix applied to the outputs of the $\text{{bf Attn}}(\cdot )$ operation for residual connection [24]. The normalized output $z_t^l$ is then fed into the feed forward network module:

where ${\Theta }_1^l\in R^{D\times (M+1)},{\Theta }_2^l\in R^{M\times (D+1)}$ denote the weight matrices of the feed forward networks and $D$ is the size of hidden nodes. The Gaussian error linear unit $\text{bf GELU}(\cdot )$ [82, 83] is adopted as the activation function in this work. The output of the last Transformer layer, $x_t^l$, is fed into the output layer, where the word probabilities are computed as in Equation (10), assuming $h_t=x_t^l$.

Conventional NNLMs using point estimate parameters fail to account for the model uncertainty associated with the word prediction. When given limited training data, these NNLMs are prone to over-fitting and poor generalization. To this end, Bayesian neural networks (BNNs) [53, 54, 55, 56, 57] can be used by modeling the uncertainty over model parameters $\Theta$ using as a probability distribution $p_r\left(\Theta \right)$. The model parameters $\Theta$ include all the activation function parameters, for example, those of various LSTM unit gates ${\Theta }_i$, ${\Theta }_f$, ${\Theta }_c$, ${\Theta }_o$ of Equation (4)-(7) in LSTM-RNN LMs, and the feed-forward layer parameters ${\Theta }_1$ of Equation (15) in Transformer LMs. The log probability of a given $n$ word sentence $W=(w_0,w_1,\dots ,w_n)$ is obtained as:

In classical Bayesian neural networks, including Bayesian NNLMs considered in this paper, the estimation of the latent variable distribution $p\left(\Theta |D\right)$ must be learned in order to perform inference on any unseen test data sentence $W$ using the predictive distribution $P\left(W|D\right)$, for example, given in Equation (III-A) for Bayesian NNLMs. The latent variable distribution $p\left(\Theta |D\right)$ needs to be learned through maximizing the model evidence expressed as the training data marginal likelihood ([55], Chapter 5.7). Since directly computing the evidence integral is intractable for neural networks, a range of approximated inference schemes were proposed in previous works including Laplace approximation [55, 54], Markov chain Monte Carlo [53, 84], and the more recently proposed variational inference [56, 58, 73, 74] that is also used in this paper. The following variational lower bound of the marginal log-likelihood is maximized instead [58]:

The second part $L_2$ of Equation (19) is the KL divergence between $q\left(\Theta \right)$ and $p_r\left(\Theta \right)$. Under the Gaussian assumption, an analytical closed-form can be derived as:

The combined use of both the parameter sampling of Equation (III-A)-(22), and the closed-form KL regularization term of Equation (23) under the Gaussian assumption over the parameter posterior and prior distributions, allows the variational lower bound of Equation (19) to be differentiable with respect to the variational distribution hyper-parameters $\mu$ and ${\Sigma }_{diag}$, and can be fully integrated into the conventional back-propagation algorithm during Bayesian LSTM-RNN or Transformer LM training.

Gaussian Processes (GPs) [86] are non-parametric distributions over continuous functions in probabilistic modeling for many machine learning applications including regression and classification tasks and beyond. A function $f(\cdot )$ modeled using Gaussian Process is represented as:

A series of prior research works were conducted in the machine learning community to build the connection between neural networks and Gaussian processes. Neal [87] proved that Bayesian neural networks [54] with a single hidden layer of infinite width are equivalent to Gaussian Processes. Hazan and Jaakkola [88] and later Lee [89] proposed the use of GP kernels to approximate infinitely wide deep neural networks. Such connection was further studied in deep Gaussian process (DGP) [90] models where deep belief neural network layers were replaced by Gaussian Processes.

The form of the traditional Bayesian neural networks introduced in Sec.III-A only considers the uncertainty associated with model parameters, but not the network structural configurations. For example, the choice over the activation functions that are widely used in Equation (4)-(7), (9) and (15), can be learned using a learnable weighted interpolation in Equation (25) over commonly used basis activation functions, i.e Sigmoid, tanh, ReLU and GELU, as considered in this paper. This interpolated form of hidden activation function, for example, the activation function in Equation (4), is given by

In addition to treating model parameters inside the activation functions as random variables as Bayesian NNLMs, further considering the uncertainty over the basis activation coefficients $\lambda$ leads to a more general form of Gaussian Process neural networks that simultaneously account for the uncertainty over both the NNLM model parameters and the choice of hidden activation functions. In GP NNLMs, assuming that the activation functions parameters $\Theta$ (linear weight parameters that are applied to the inputs of a neuron, before any non-linearity) and basis activation coefficients $\lambda$ are statistically independent of each other, the sentence level probability of a given word sequence $W$ is then computed as follows:

where $p\left(\Theta |D\right)$ and $p\left(\lambda |D\right)$ denote the posterior probability distributions over the basis activation function parameters $\Theta$ and the basis coefficients $\lambda$ respectively, both to be learned from the training data $D$.

Similar to the variational inference approach considered in Sec.III-A for Bayesian NNLMs, the following variational lower bound of marginal log-likelihood is maximized:

where $q\left(\Theta \right)$ and $q\left(\lambda \right)$ are variational approximations of posterior distributions $p\left(\Theta |D\right)$ and $p\left(\lambda |D\right)$ that separately model the uncertainty over model parameters and the choice of hidden neural activation functions. Their respective KL-divergence against their prior distributions in $L_2$ and $L_3$ serve as regularization terms during GP NNLM inference.

In common with Bayesian NNLMs, both the prior distributions, $p_r\left(\Theta \right)$ and $p_r\left(\lambda \right)$, and the variational distributions, $q\left(\Theta \right)$ and $q\left(\lambda \right)$, for model parameters and basis activation coefficients are assumed to be Gaussian distributions [71, 72, 73, 74] to produce closed-form, analytical solutions to the two KL terms.

The expectation of log-likelihood in the first term $L_1$ of the variational lower bound in Equation (28) is calculated using sampling method that requires activation parameter samples and basis activation coefficient samples to be drawn from their respective variational distributions, $q\left(\Theta \right)$ and $q\left(\lambda \right)$, with the re-parameterization method of Equation (22).

In contrast to the BNNs and GPNNs presented in Sec.III-A and III-B, variational neural networks (VNNs) [58, 59, 91, 70] can be used to model uncertainty associated with hidden representations. Instead of modeling the uncertainty over the weight parameters $\Theta$ inside the activation functions in BNNs or assuming additional uncertainty over the activation basis coefficients $\lambda$ in GPNNs, VNNs introduce a latent variable $z_t$ to encode the uncertainty associated with, for example, the hidden node output vector $h_t$ of LSTM-RNNs in Equation (9), or the outputs of feed-forward networks $s_t^l$ of Transformer models in Equation (15). The general form of sentence level probability of a given word sequence $W=(w_0,w_1,\dots ,w_n)$ under variational NNLMs is computed as:

During variational NNLMs training, the following training data marginal log-likelihood lower bound is maximized,

The hyper-parameters ${\mu }_t,{\Sigma }_{diag,t}$, where ${{\sigma }_{t,i}}^2$$=$${\Sigma }_{diag,t,ii}$, are calculated using an inference network ${\Phi }^{infer}\left(w_0,\dots ,w_{t-1}\right)$, and ${\mu }_t^r,{\Sigma }_{diag,t}^r$, where ${{\sigma }_{t,i}^r}^2$$=$${\Sigma }_{diag,t,ii}^r$, through an prior network ${\Phi }^{prior}\left(w_0,\dots ,w_{t-1}\right)$ respectively. Both the inference network and prior network are regression networks. The inputs of the two networks are history sequence $\left(w_0,\dots ,w_{t-1}\right)$. Each of the network produce two outputs including the mean and variance vectors of $q\left(z_t\right)$ and $p_r\left(z_t\right)$ as follows.

The input word sequence $\left(w_0,\dots ,w_{t-1}\right)$ are approximated by the hidden state $h_t$ in LSTM-RNNs, and the output of the Transformer feed-forward modules $s_t^l$ as shown in Fig.1 (d) and Fig.2 (d) respectively. In this work, both the inference networks and prior networks we use contain one single hidden layer for the Variational NNLMs. Both the inference and prior network modules are differentiable with respect to their internal parameters. Similar to the Bayesian and Gaussian Process NNLMs of Sec.III-A and III-B, sampling the hidden output vector $z_t$ from the variational distribution $q\left(z_t\right)$ allows the marginal log-likelihood item $L_1$ of Equation (30) to be efficiently approximated. The maximization of the variational lower bound of Equation (30) can then be integrated into the back-propagation algorithm during variational LSTM-RNN or Transformer LM training.

When training all three Bayesian learning approaches based NNLMs, suitable choices of parameter prior distributions need to be set. In this paper, the prior distribution Gaussian means for various Bayesian learned LSTM-RNN and Transformer LMs are based on the parameter estimates of comparable deterministic, point estimated NNLMs that have converged in model training. Throughout this paper the prior Gaussian variances used for Bayesian learning based LSTM-RNN and Transformer LMs are empirically set as 1 and 1.0e-3 respectively. In addition, all the other standard, non-Bayesian estimated parameters in the Bayesian learning based LSTM-RNN and Transformer models are initialized using the parameters obtained from the comparable pre-trained standard models reaching half convergence. The combination of these two settings in practice was found to yield a good balance between convergence speed and performance¹¹1Several methods to initialize the other standard non-Bayesian estimated model parameters include: 1) random weights; 2) the weights from both half-converged standard models; and 3) fully converged pre-trained standard models for Bayesian learning based NNLMs. Experimental results suggested that the model parameters initialized from the half-converged or fully converged pre-trained standard models produced similar performance, while both marginally better than using random weights..