Improving Semantic Control in Discrete Latent Spaces with Transformer Quantized Variational Autoencoders (2024)

Abstract: Achieving precise semantic control over the latent spaces of Variational AutoEncoders (VAEs) holds significant value for downstream tasks in NLP as the underlying generative mechanisms could be better localised, explained and improved upon. Recent research, however, has struggled to achieve consistent results, primarily due to the inevitable loss of semantic information in the variational bottleneck and limited control over the decoding mechanism. To overcome these challenges, we investigate discrete latent spaces in Vector Quantized Variational AutoEncoder (VQVAE) to improve semantic control and generation in Transformer-based VAEs. In particular, We propose T5VQVAE, a novel model that leverages the controllability of VQVAE to guide the self-attention mechanism in T5, exploiting its full generalization capabilities. Experimental results indicate that T5VQVAE outperforms existing state-of-the-art VAE models, including Optimus, in terms of control and preservation of semantic information across different tasks such as auto-encoding of sentences and mathematical expressions, text transfer, and inference. Moreover, T5VQVAE exhibits improved reasoning capabilities, suggesting potential applications for downstream natural language and symbolic inference tasks.

## introduction the emergence of deep generative neural networks supported by variational autoencoders (vaes) @xcite enables the localisation of syntactic and semantic properties within complex sentence latent spaces. by localising and manipulating these generative factors within the latent spaces, one can better control the properties of the textual output, enhancing performance on downstream tasks @xcite , and providing mechanisms for representing and disentangling syntactic and semantic features animals require warmth in cold environment x l2 index ... 10 x 43 argmin lookup vq index index ... ... figure 1 : by controlling the token-level discrete latent space in vaes, we aim to explicitly guide the crossattention mechanism in t5 to improve the generation process. we focus on three challenging tasks to assess precise semantic control and inference. recent work @xcite @xcite investigated controllable text generation via latent sentence geometry based on the canonical optimus architecture (the first large pretrained language vae, @xcite ). however, the optimus architecture brings its associated challenges since (i) the optimus setup does not allow for a fine-grained (i.e., token-level) semantic control as sentence-level representation features are ignored by most attention heads especially in lower layers, where lexical-level semantics is captured @xcite ; (ii) the sentence bottleneck in the vae architecture leads to inevitable information loss during inference @xcite . this work concentrates on addressing these architectural limitations by aiming to minimise the information loss in the latent space and effectively control the decoder and its attention mechanism. the vector quantized variational autoencoder (vqvae) @xcite , as a discrete latent variable model, can be considered an ideal mechanism to alleviate these issues since it preserves and closely integrates both a coarse-grained continuous latent sentence space and a fine-grained latent token space that can preventinformation loss. more importantly, its latent token space can directly work on the cross-attention module @xcite to guide the generation in seq2seq models, such as t5 @xcite . therefore, we hypothesise that such a mechanism can enable better generalisation and semantic control in transformerbased vaes. following these insights, we propose a novel approach named t5vqvae, a model that leverages the controllability of vqvae to guide the token-level self-attention mechanism during the generation process. we evaluate t5vqvae on three challenging and diverse downstream tasks including (1) language modelling, (2) text transfer (guided text generation via the movement of latent vectors), and (3) natural language and symbolic inference tasks. an illustration of the complete model architecture and experimental setup can be found in figure 1 . the overall contribution of the paper can be summarised as follows: the experiments reveal that the proposed model outperforms previous state-of-the-art vae models, including optimus @xcite , on three target tasks, as well as delivering improved semantic control when compared to the previous state-of-the-art. our experimental code is available online foot_0 to encourage future work in the field. ## methodology in this section, we first present our model, t5vqvae, whose primary goal is to learn a latent space by reconstructing input sentences. next, we illustrate its objective function, which consists of three parts designed to improve semantic control: reconstruction term, latent space optimization term, and encoder constraint term. finally, we highlight the architectural advantages of t5vqvae compared to transformer-based vaes. model architecture. van den oord et al. (2017) first proposed the vqvae architecture for learning a discretised latent space of images, showing that it can alleviate the issue of posterior collapse, in which the latent representations produced by the encoder are ignored by the decoder (kingma and welling, 2013). in this work, we propose to integrate t5 encoder/decoder into the vqvae architecture for representation learning with natural language. t5 was selected due to its consistent performance across a large range of nlp tasks and its accessibility. to cast t5 into a vqvae model, we first establish a latent token embedding space, denoted as the codebook, represented by @xmath0. here, k refers to the number of tokens in the codebook, and i represents the dimensionality of each token embedding. when given a token x, the encoder e maps it into a vector representation, denoted as e(x). then, the nearest latent representation z k from the codebook z is selected based on the l2 distance. the input of the crossattention module can then be formalised as follows: here, z k is the key and value and d(x), which represents the input token embedding of the decoder, is the query. x represents the reconstructed token, while w q , w k , and w v are trainable weights of query, key, and value. training t5vqvae the training of t5vqvae can be then considered as the optimisation of three independent parts, including d(z k ), z k , and e(x). starting from d, the model can be trained by maximising the reconstruction probability p (x|d(z k )) via the teach-forcing scheme. next, the z k is optimised by minimising the l2 distance between e(x) and z k , which can be described as (sg[e(x)]-z k ) 2 where sg is the stop gradient operation. finally, e(x) can be trained via the l2 distance. by ensuring that e(x) can learn the latent embedding under the constraint of r k×i rather than learning an embedding directly, we can guide the model to achieve better performance. a commitment weight β < 1 is used to constraint the e close to z k , which can be described as: β(e(x)-sg[z k ]) 2 . β is set to 0.25 following the same setup as @xcite to preserve a behaviour consistent with their findings. the final objective function of t5vqvae can be formalised as follows: training the latent space. there are two possible strategies to update the latent space: i. k-means and ii. gumbel softmax. regarding k-means, for each token embedding w i in a sentence, it selects the nearest latent token embedding, z k , to its token embedding e w i . this process is equivalent to classifying e w i using k-means and then choosing the corresponding central point z k as the input for d(z k ). this can be expressed as follows: to improve the stability of latent space training (term 2), we adapted the exponential moving average (ema) training scheme to update z @xcite . figure 2 displays the training and testing loss curves of t5vqvae with ema or not. more details of ema are provided in appendix a. instead of using k-means, which performs a soft selection of the index k, we can utilize the gumbel softmax trick @xcite for a hard sampling of the index k. this trick involves sampling a noise value g k from the gumbel distribution and then using the softmax function to normalize the output, resulting in a probability distribution. by selecting the index with the highest probability, we obtain a discrete choice. this entire process can be described as follows: in this context, t k represents the probability of the k-th token, which can be obtained through a linear transformation before being fed into the gumbel softmax. the parameter τ serves as a temperature hyper-parameter that controls the closeness of the new distribution to a discrete distribution. as τ approaches zero, the distribution becomes one-hot, while a non-zero value of τ leads to a more uniform distribution. in our experiments, we experienced convergence issues when using the gumbel softmax scheme, and therefore decided to adopt the k-means mechanism which generally leads to better results. ## controllability evaluation next, we put forward two metrics for quantitatively evaluating the controllability of the proposed model (t5vqvae), which we refer to as semantic disentanglement and interpolation smoothness. the former evaluates the controllability from the perspective of disentanglement of semantic factors (e.g., arguments and associated semantic roles). the latter evaluates the smoothness and coherence of the latent space geometry during interpolation. ## conclusion and future works in this work, we build a model for improving the semantic and inference control for vae-enabled language model (autoencoding) architectures. we propose a new model (i.e., t5vqvae) which is based on the close integration of a vector-quantized vae and a consistently accessible and high-performing language model (t5). the proposed model was extensively evaluated with regard to its syntactic, semantic and inference controls using three downstream tasks (autoencoding, text transfer, and inference task). our experimental results indicate that the t5vqvae can outperform the canonical stateof-the-art models in those tasks and can deliver a quasi-symbolic behaviour in the inference task (via the direct manipulation of the latent space). as future work, we plan to further explore applications on symbolic natural language inference via the direct manipulation of the latent space, and to investigate the controllability of recent large language models through the vqvae architecture. moreover, additional research directions could be informed by the current work: word-level disentanglement. our architecture provides a foundation to explore token/word-level disentanglement for more general sentence and inference representation tasks. while sentence-level disentanglement is widely explored in the nlp domain, such as sentiment-content @xcite , semantic-syntax @xcite , and negationuncertainty @xcite , or syntacticlevel disentanglement @xcite , this mechanism is still under-explored in other nlp tasks @xcite . interpretability. discrete properties derived from vector quantization can enable the further probing and interpretability of neural networks by discretizing continuous neural latent spaces, where symbolic concepts are emerging in both images @xcite and natural language @xcite domains.