Isotropic Contextual Representations through Variational Regularization

1 INTRODUCTION . Self-supervised , attention-based language models are prone to overfitting or memorizing input data , especially when trained on smaller datasets . Nevertheless , models such as BERT ( Devlin et al. , 2019 ) , based on an encoder-only architecture and trained on very large datasets , exhibit state-of-theart performance on common natural language processing tasks . The resulting token representations of these models are distributed and highly contextual , but the latent space does not exhibit additional structural properties such as isotropy and resulting completeness ( Ethayarajh , 2019 ) . We develop an autoencoder architecture that avoids overfitting by increasing latent space isotropy . At the same time , enforcing a complete latent space leads to a decoder that is able to interpret ( i.e . decode ) any given point in latent space . While the primary goal are tasks that require paraphrase generation , the resulting encoders surprisingly also perform on par in transfer tasks with smaller training setups . Previous work has shown that isotropic latent representation , i.e . distributed across latent space instead of populating a narrow cone ( Ethayarajh , 2019 ) , can improve the performance of language models on transfer tasks . Gao et al . ( 2019 ) improve the performance of a Transformer-based model with increased isotropy by directly optimizing the cosine distance between the latent representations . Wang et al . ( 2020 ) suggest to control the singular value distribution of the output representation matrix and Li et al . ( 2020a ) transform the original representation distribution into a Gaussian distribution through normalizing flows to increase the isotropy of the underlying models . Similar to Gao et al . ( 2019 ) , we propose to apply a regularizing constraint during training to enforce an isotropic and , even more , complete latent space . A common regularization network is the socalled Variational Autoencoder ( VAE ) ( Kingma & Welling , 2014 ) . VAE networks consist of an encoder that maps given input data not to a point in latent space but a distribution . A VAE ’ s decoder is required to successfully reconstruct the original input from samples of the latent distribution . Thus , the VAE is optimized to reconstruct the given input sequence whilst enforcing the latent distributions to match a given prior distribution . We propose a Variational Auto-Transformer ( VAT ) , a VAE based on a Transformer architecture , where the variational loss is computed at token level . We show that by enforcing a Gaussian distribution as latent prior , the latent token-level representations become more isotropic in comparison to models such as BERT . We introduce the prior distribution ’ s standard deviation as model parameter to optimize isotropy and balance the language generation variety against the network ’ s reconstruction ability . Besides isotropy , the completeness of the latent space is demonstrated . The pretrained decoder can map back every point in latent space to a textual representation , be it actual encoded or synthetic , sampled latent points . This allows for text generation through ” variational sampling ” and other manipulations such as interpolation directly in latent space . While the token-level representations can be used to generate or manipulate text , we also show that sentence-level representations , e.g . obtained through averaging , are suitable for sentence classification tasks . 2 BACKGROUND . An autoencoder is a neural network architecture consisting of an encoder enc , that maps any input x ∈ Rd to a point z = enc ( x ) in latent space , and a decoder dec , such that x ≈ dec ( enc ( x ) ) ( Goodfellow et al. , 2016 ) . The chosen network architecture defines the families E and D of the encoder and decoder . During training of the autoencoder , the network parameters are optimized to find the optimal ( enc∗ , dec∗ ) pair that minimizes a given loss function L : ( enc∗ , dec∗ ) = argmin ( enc , dec ) ∈E×D ( L ) ( 1 ) For the standard autoencoder with continuous inputs x ∈ Rd , the reconstruction loss L = LREC is : LREC = ∥x− dec ( enc ( x ) ) ∥2 ( 2 ) Over-complete autoencoder architectures with many degrees of freedom are prone to a particular form of overfitting : The encoder maps each data point x to an isolated point z in the latent space such that the decoder memorizes the lossless reconstruction of each of these codes . This highly discrete latent space lacks completeness and is of limited use for advanced NLP applications . 2.1 VARIATIONAL AUTOENCODER . The variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ) can be thought of as generative autoencoder . A VAE encodes an input data point as a distribution over the latent space by adding a regularization term to the loss function . The regularization term LREG assesses the KulbackLeibler divergence between the latent distribution and a standard Gaussian based on their means µ and covariances Σ : LREG = −DKL [ N ( µ , Σ ) , N ( 0 , I ) ] ( 3 ) µ and Σ are either estimated using standard point estimators from the encoder outputs , or computed by learned functions such that µ = g ( x ) and Σ = h ( x ) with g ∈ G and h ∈ H , where G and H are families of network architectures . G and H can contribute to the desired properties of the latent representation by implicitly implementing dimensionality reduction or feature disentangling . Decoding from the latent representation requires sampling from z ∼ N ( µ , Σ ) . In order to enable backpropagation despite this sampling operation , Kingma and Welling ( Kingma & Welling , 2014 ) introduced the reparameterization trick : Instead of sampling from the latent distribution , a random sample ϵ ∼ N ( 0 , I ) from a standard Gaussian is drawn and then transformed by the computed mean and standard deviation z = LT ϵ+ µ , where LT is the Cholesky factor of Σ . The VAE ’ s total loss is the sum of the regularization and reconstruction term : ( enc∗ , dec∗ ) = argmin ( enc , dec ) ∈E×D ( LREC + LREG ) = argmin ( enc , dec ) ∈E×D ( ∥x− dec ( z ) ∥2 −DKL [ N ( µ , Σ ) , N ( 0 , I ) ] ) ( 4 ) 3 VARIATIONAL AUTO-TRANSFORMER . Initially proposed for machine translation , the Transformer ( Vaswani et al. , 2017 ) can be used as language model and framed as autoencoder . Extended with a token-level variational loss , the VAT maps input tokens to context-aware , distributed latent representations . Figure 1 illustrates the architecture of the VAT . The set of encoders E and decoders D is defined by the Transformer ’ s network architecture based on self-attention modules . The encoder maps the embedded input sequence X = { x1 , x2 , . . . , xT } , X ∈ RT×d to a latent representation Z = { z1 , z2 , . . . , zT } , Z ∈ RT×d , where T is the sequence length and d is the model dimension . To do so , the model predicts mut and σt from xt and samples zt ∼ N ( µt , σt ) . Z is then used as attention source in the decoder to predict the next word based on the already produced output . As X and Z depend on T , the objective is to minimize L = 1 T T∑ t=1 ( Eq ( Z|X ) ∥xt − dec ( zt ) ∥2 −DKL [ q ( zt|X ) , p ( zt ) ] ) ( 5 ) where p ( zt ) ∼ N ( 0 , I ) and q ( zt|X ) ∼ N ( µt , Σt ) . More precisely , we let Σt = diag ( expσt ) where the vector σt = ( σj ) j=1 ... d is predicted , so that computing the actual value of Σ involves the exponential as a non-linear activation function . The output of the encoder is passed to two linear layers in parallel . The linear layers have the same number of input and output nodes and learn a transformation to predict µt and σt for each token . The latent representation Z is used as input in the encoder-decoder attention layer of the Transformer ’ s decoder . 3.1 SCALING THE REGULARIZATION LOSS . During our experiments , we experienced overfitting on the training data for the reconstruction loss . The learning curves of a corresponding experiment are illustrated in Figure 2a . Weighting the regularization loss according to a logistic annealing function as suggested by Bowman et al . ( 2016 ) or by a scaling factor β similar to the scaling of the beta-VAE ( Higgins et al. , 2017 ) , but with β < 1 , only had little effect . Thus , we propose to scale the covariance matrix of the target distribution of LREG instead . This is motivated by the observation that standard Gaussians in d-dimensional latent spaces ( e.g . d ≥ 128 ) will overlap considerably if their mean values are - at the same time - regularized to zero . As a consequence , LREG can not be minimized by the model without increasing LREC in an undesirable way . Scaling the standard deviation to a too small value , though , might result in peaky distributions , having no regulatory effect and resulting in a trivial reconstruction objective . We adapt the computation of the regularization loss to incorporate a scaling factor s for the standard deviation as hyperparameter . The closed form of the VAE loss function1 with a prior standard deviation σp = s ∈ ( 0 , 1 ] and prior µp = 0 for a specific token ’ s layer outputs σt and µt becomes −DKL [ N ( µt , σt ) , N ( 0 , diag ( s ) ) ] = d∑ i=1 1 2 log ( σ2t , i ) − σ2t , i + µ 2 t , i 2s2 − log ( s ) + 1 2 ( 6 ) This results in codes zt = σt · ϵ′t + µ being distributed according to a Gaussian with zero mean and a predefined standard deviation , ϵ′ ∼ N ( 0 , diag ( s ) ) . The optimal balance between LREC and LREG as expressed by the value of s ∈ ( 0 , 1 ] is related to the representations ’ isotropy and will be determined experimentally . 4 EXPERIMENTS . The aim of the experiments is to balance the two loss terms . LREC should be low to maintain the reconstruction ability of the network , while a too small scaling value s implies that the VAT essentially behaves as regular AE . At the same time , we want to optimize the isotropy of the latent representations . In order to choose the best setting , we observe both , the resulting loss values and the properties of the latent space in terms of similarity between representations . Our VAT model architecture is smaller than the original Transformer architecture , as it is not intended for machine translation . The VAT consists of N = 4 encoder and decoder layers , respectively , with H = 8 attention heads each and a model dimensionality of d = 128 . A logistic annealing function ( Bowman et al. , 2016 ) is applied to weight LREG . The VAT is trained on the train split of the WMT19 de-en dataset ( Barrault et al. , 2019 ) using English sentences only . WMT19 contains data from the news commentary , Wiki titles , Europarl , ParaCrawl and Common Crawl corpora . The data is tokenized using subword tokenization with a target vocabulary of 215 tokens . The full details on model and training parameters are listed in the appendix .

Summary: To obtain an isotropic word embedding space, the authors proposed a new transformer-based autoencoder in which each word is represented as a normal distribution, and a variational loss is used to make it closer to an isotropic normal distribution. By carefully adjusting the weights of the variational loss (magnitude of variance), the learning converged, and an isotropic latent representation was obtained. They also tried to show some of the advantages of the proposed approach: 1. The proposed model can variationally sample sentences with slightly different meanings from the input sentences. 2. In some cases, the proposed model can generate a semantic "interpolation" between two given sentences. 3. Solving sentence-level tasks using word representations obtained by the proposed model outperforms MiniBERT.