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: This paper presents a variational autoencoder model for learning representations of tokens in a sequence of text. The authors build upon prior work which argues that the point-estimate representations learned by Transformer-based language models are degenerate, i.e. they only populate a narrow subspace. The authors address this problem by learning Gaussian-distributed representations (instead of point-estimates) which are regularized using a KL divergence loss. Empirical results on classification tasks are competitive with point-estimate baselines.