Probabilistic Transformer: A Probabilistic Dependency Model for Contextual Word Representation

Given a sentence (a sequence of words), denote $n$ as the sequence length. For the $i$-th word, we define $Z_{i}$ as a discrete latent label that represents the syntactic (and possibly semantic) property of the word in the sentence (i.e., it is a contextual representation) with a label set of size $d$. Such a discrete representation deviates from the common practice of representing a word with a continuous vector, but it is sufficient at least for syntactic processing Kitaev et al. (2022) and it greatly simplifies our probabilistic model. For the $i$-th word, we also define $H_{i}\in\{1,2,\cdots,n\}$ representing the syntactic dependency head of the word. So the set of variables $\{H_{i}\}_{i=1}^{n}$ specifies a dependency structure. We may also allow $H_{i}$ to point to a dummy root node, which will be discussed in Section 2.3.5. We follow the head-selection paradigm of dependency parsing and do not enforce the tree constraint, which again simplifies our model design.

Next, we define two types of potential functions. For the $i$-th word $w_{i}$, we define a unary potential function (corresponding to the unary factors in Figure 1) evaluating the compatibility of the word and its label $Z_{i}$:

where $\mathbf{S}\in\mathbb{R}^{|\mathcal{V}|\times d}$ is a score matrix, $|\mathcal{V}|$ is the size of the vocabulary. For simplicity, we do not exploit any morphological or contextual features for computing the scores. For every pair of words $w_{i}$ and $w_{j}$ ($i\neq j$), we define a ternary potential function (corresponding to the ternary factors in Figure 1) over $Z_{i}$, $Z_{j}$ and $H_{i}$, which evaluates the compatibility between the labels of the two words if $w_{j}$ is the dependency head of $w_{i}$:

where $\mathbf{T}\in\mathbb{R}^{d\times d}$ is a score matrix.

Inspired by the multi-head structure in transformers, we allow multiple dependency structures for the same sentence, which may represent different flavors of dependencies. Each dependency structure resides in a different channel with its own dependency head variables and ternary potential functions. For the $c$-th channel, we denote the set of dependency head variables by $\{H_{i}^{(c)}\}_{i=1}^{n}$ and the score matrix of the ternary potential function by $\mathbf{T}^{(c)}$. Let $h$ denote the total number of channels. We may stack all the score matrices $\mathbf{T}^{(c)}$ for $c=1,\cdots,h$ to form a score tensor $\mathbf{T}\in\mathbb{R}^{d\times d\times h}$. Note that all the channels share the same set of latent label variables $\{Z_{i}\}_{i=1}^{n}$.

Following Wang and Tu (2020), we use Mean Field Variational Inference (MFVI) to perform approximate inference. Different from the previous work, however, we need to run inference over latent labels in addition to dependency heads.

MFVI iteratively passes messages between random variables and computes an approximate posterior marginal distribution for each random variable (denoted by $Q(\cdot)$). Let $\mathcal{F}_{ic}^{(t)}$ denote the message received by variable $H_{i}^{(c)}$ at time step $t$ from ternary factors, and $\mathcal{G}_{i}^{(t)}$ denote the message received by variable $Z_{i}$ at time step $t$ from ternary factors. We have

where

are the approximate marginal distributions at time step $t$, with $Q_{i}^{(t)}(\cdot)$ over $Z_{i}$ and $Q_{ic}^{(t)}(\cdot)$ over $H_{i}^{(c)}$. We initialize these distributions by

After a fixed number of $T>0$ iterations, we obtain the final posterior marginal distribution $Q_{i}^{(T)}(Z_{i})$ for $i=1,\cdots,n$. Resulted from interactions with all the words of the sentence, the distribution $Q_{i}^{(T)}(Z_{i})$ incorporates information of not only the $i$-th word, but also its context. Therefore, we can treat the probability vector of this distribution as a contextual vector representation for the $i$-th word. In practice, we find that using unnormalized scores in log space as contextual word representations produces better results, i.e., we skip exponentiation and normalization when computing $Q_{i}^{(T)}(Z_{i})$ using Equation 5 during the final iteration.

Since all the computation during MFVI is fully differentiable, we can regard the corresponding computation graph as a recurrent or graph neural network parameterized with score matrix $\mathbf{S}$ and tensor $\mathbf{T}$. We can use the contextual word representations for downstream tasks by connecting the network to any downstream task-specific network, and we can update the model parameters using any task-specific learning objective through gradient descent. This is exactly the same as how transformers are used.

We introduce a few extensions and variants to the basic model that are empirically beneficial. Additional variants are discussed in Appendix B.

Consider a probabilistic transformer using a distance-insensitive ternary potential function without a dummy root node. We tensorize the update formulas in the inference process of probabilistic transformers. Suppose $Q_{z}^{(t)}\in\mathbb{R}^{n\times d}$ is a tensor that represents the posterior distributions of all the $Z$ variables, and $Q_{h,c}^{(t)}\in\mathbb{R}^{n\times n}$ is a tensor that represents the posterior distributions of all the $H$ variables in channel $c$ (with a zero diagonal to rule out self-heading). We can rewrite Equation 3 and 4 as

where

and $\sigma$ is the softmax function. We still set $\lambda_{Z}$ to its default value $1$ but regard $\lambda_{H}$ as a hyperparameter. With asynchronous update, Equation 18 becomes:

We assume that $\mathbf{T}^{(c)}$ is symmetric for $c=1,\cdots,h$. This is the only assumption that we make in this section beyond the original definition from the previous section. Symmetric score matrices indicate that the ternary factors are insensitive to the head-child order, which is related to undirected dependency parsing Sleator and Temperley (1993). If $\mathbf{T}^{(c)}$ is symmetric, then $Q_{h,c}^{(t)}$ is also symmetric based on Formula 15 and 19. Thus, we can simplify Equation 16 to

Suppose we decompose the ternary score tensor into two tensors $\mathbf{U},\mathbf{V}\in\mathbb{R}^{d\times h\times r}$ according to Equation 14, which can be rewritten as:

where $\mathbf{U}^{(c)},\mathbf{V}^{(c)}\in\mathbb{R}^{d\times r}$ are the $c$-th channel of tensor $\mathbf{U}$ and $\mathbf{V}$ respectively. Substitute 21 into 15 and 20, we have

We define

For time step $t-1$, we could rewrite Formula 22 and 23 as

Apply Equation 27, 19, 26 to 17, we have

where

We call the computation of $\operatorname{channel}_{c}$ a single-channel update for channel $c$.

Now we have a tensorized formulation of the computation in probabilistic transformers and we are ready for its comparison with transformers at three different levels.

Scaled dot-product attention in transformers is formulated as:

As we can see, our single-channel update in Equation 29 is almost identical to scaled dot-product attention in transformers. The only difference is that the diagonal of the tensor $Q_{c}K_{c}^{T}$ is zero in our model because the head of a word cannot be itself.

Multi-head attention in transformers is formulated as:

where

It is equivalent to

where $W^{O}\equiv\operatorname{Concat}(W_{1}^{O},\ldots,W_{h}^{O})$ and $W_{i}^{Q},W_{i}^{K},W_{i}^{V},W_{i}^{O}\in\mathbb{R}^{d\times r}$. Our multi-channel update formula (the second term within the softmax function in Equation 28) is similar to the multi-head attention in transformers, as shown in Figure 2. The main difference is that probabilistic transformers use the same parameters for $W^{K}$ and $W^{V}$ (both are $\mathbf{V}$, shown in green color in Figure 2(b)) and for $W^{Q}$ and $W^{O}$ (both are $\mathbf{U}$, shown in orange color in Figure 2(b)).

Recall that $\mathbf{U}$ and $\mathbf{V}$ are obtained from matrix decomposition (Equation 14). Therefore, the correspondence between $\mathbf{U}$, $\mathbf{V}$ and $W^{Q},W^{K},W^{O},W^{V}$ in transformers suggests that the latter can also be seen as derived from tensor decomposition. Previous work on transformers has the same findings Elhage et al. (2021).

Figure 3 compares the full computation graphs of the two models, which have a similar overall structure that repeats a module recurrently until outputting contextual word representations. Within the module, we have also established the correspondence between multi-channel update and multi-head attention. On the other hand, there are a few interesting differences.

First, our model does not have a feed-forward structure as in a transformer. However, we do propose a variant of our model that contains global variables representing topics (Appendix B.3), which may have similar functionality to the feed-forward structure.

Second, our model does not have residual connections or layer norms. Instead, it adds the initial distributions (unary scores) to the updated message at each iteration. This may replace the functionality of residual connections and may even make more sense when the downstream task strongly depends on the original word information.

Third, we have an additional softmax in each iteration. Note that we do softmax before the first iteration (Equation 7) and also at the end of each iteration (Equation 28), but bypass it in the last iteration when producing the output word representations, so our model could be equivalently formulated as doing softmax before each iteration, which we show in Figure 3(c). Doing softmax in this way is similar to the layer norm in pre-LN transformers Xiong et al. (2020) (Figure 3(b)).

Finally, our model shares parameters in all iterations. This is similar to some variants of transformers that share parameters between layers, such as Universal Transformer Dehghani et al. (2019) and ALBERT Lan et al. (2019).

One consequence of these differences is that probabilistic transformers have much fewer parameters than transformers with the same number of layers, heads and embedding dimensions, because of shared parameters between iterations, absence of a feed-forward structure, and tied parameter matrices in multi-channel updates.

Here we briefly introduce our tasks and datasets. A detailed description is presented in Appendix D.

Masked Language Modeling (MLM). We perform MLM tasks on two corpora: the Penn TreeBank (PTB) Marcus et al. (1993) and Brown Laboratory for Linguistic Information Processing (BLLIP) Charniak et al. (2000). Following Shen et al. (2022), we randomly replace words with a mask token <mask> at a rate of 30%. The performance of MLM is evaluated by measuring perplexity (lower is better) on masked words.

We project the final word representation of each mask token to the vocabulary. For transformers, we tie the projection parameters to the initial word embeddings. We find that this trick improves the performance of transformers.

Sequence Labeling. For sequence labeling tasks, we perform part-of-speech (POS) tagging on two datasets: the Penn TreeBank (PTB) Marcus et al. (1993) and the Universal Dependencies (UD) De Marneffe et al. (2021). We also perform named entity recognition (NER) on CoNLL-2003 Tjong Kim Sang and De Meulder (2003).

We directly project the final word representation of each word to the target tag set. For POS tagging, we evaluate the results by the accuracy of word-level predictions. For NER, we evaluate the results by measuring the F1 score of named entities.

Text Classification. We use the Stanford Sentiment Treebank (SST) Socher et al. (2013) as the dataset. It has two variants: binary classification (SST-2) and fine-grained classification (SST-5).

For transformers, we add a <CLS> token at the front of the sentence and then project its representation to the tag set. For our model, we use the variant with a root node introduced in Section 2.3.5 and project the representation of the root node to the tag set.

Syntactic Test. To evaluate the compositional generalization abilities of our model, we perform a syntactic test on the COGS dataset Kim and Linzen (2020). We follow the settings in Ontanón et al. (2021), who cast the task as a sequence labeling task.

As in sequence labeling, we project word representations to tag sets. If all words in a sentence are correctly predicted, the sentence prediction will be counted as correct. We evaluate the results by the sentence-level accuracy of the predictions.

We tune transformers and our model separately for each task except the syntactic test. For the syntactic test, we find that both transformers and our model easily reach 100% accuracy on the validation set. This observation is consistent with Ontanón et al. (2021). Therefore, instead of tuning, we use the best-performed setting of transformers in Ontanón et al. (2021) for our experiments. The hyperparameters of our model are determined by their counterparts of transformers based on the correspondence discussed in Section 3.

For our model, we integrate all the variants mentioned in Section 2.3 except the root node variant, which we only use for text classification tasks. We tune the tensor decomposition strategy on different tasks. For MLM tasks, we add a small L2 regularization term to the ternary scores in our model, which we experimentally find beneficial. We optimize both models using the Adam optimizer Kingma and Ba (2015) with $\beta_{1}=0.9$, $\beta_{2}=0.999$.

We report the average and standard deviation results of 5 random runs in Table 1. It shows that our model has a competitive performance compared with transformers. In most tasks, probabilistic transformers perform competitively to transformers. It is worth noting that in these experiments, probabilistic transformers have much fewer parameters than transformers. For most tasks, the number of parameters of our best model is about one-fifth to one-half of that of the best transformer.

We also conduct case studies of the dependency structures inferred by our model after training on downstream tasks. Similar to the case of self-attentions in transformers, the inferred dependency structures are only partially consistent with human intuition. See Appendix F for details.