summaries: Understanding drug-drug interactions (DDI) is an important task in drug development and prescription management. The authors proposed a new graph energy neural network (GENN) for DDI prediction. Comparing to the existing Decagon model (Zitnik et al. 2018), the proposed new model considered correlations between DDI types and used a new energy function to capture this information. By comparing to the previous baselines, the authors demonstrated their approach was able to achieve the SOAT performance in terms of prediction accuracy, be more robust to missing DDI data, and better capture the correlations between DDI types.

summaries: This paper proposes a framework to learn correlated drug-drug interaction based on structured prediction energy networks (SPEN). The core idea is to model the dependency structure of the labels (multi-label) by minimizing a designed energy function. The graph energy is designed as MLP over the mean of all nodes embeddings, where the nodes embeddings are obtained through a graph convolutional network. The edge information is included in the node embedding when aggregating neighborhood information. The proposed method also introduces an additional test inference network to jointly train with the cost-augmented training network under the semi-supervised setting. The authors tested on two DDI datasets and the result shows improvement compared to several baseline methods.

summaries: This paper presents a graph neural network for drug-to-drug interaction (DDI) prediction, which explicitly models link type correlation. Basically, the drug-to-drug interaction prediction problem is a specific type of link prediction task, with drugs as vertices and interaction as edges, and the authors propose a graph neural network with an energy-based formulation where the link types are encoded as the graph edges. The authors validate their method against feedforward GNNs on two DDI prediction datasets, and achieve significantly improved performances.

abstractText: Gaining more comprehensive knowledge about drug-drug interactions (DDIs) is one of the most important tasks in drug development and medical practice. Recently graph neural networks have achieved great success in this task by modeling drugs as nodes and drug-drug interactions as links and casting DDI predictions as link prediction problems. However, correlations between link labels (e.g., DDI types) were rarely considered in existing works. We propose the graph energy neural network (GENN) to explicitly model link type correlations. We formulate the DDI prediction task as a structure prediction problem, and introduce a new energy-based model where the energy function is defined by graph neural networks. Experiments on two real world DDI datasets demonstrated that GENN is superior to many baselines without consideration of link type correlations and achieved 13.77% and 5.01% PR-AUC improvement on the two datasets, respectively. We also present a case study in which GENN can better capture meaningful DDI correlations compared with baseline models.

authors:
Claudio Angione
Max Conway
Pietro Lió

title: Multiplex methods provide effective integration of multi-omic data in genome-scale models

venue: BMC bioinformatics,

authors:
David Belanger
Andrew McCallum

title: Structured prediction energy networks

venue: In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48,

title: Structured prediction energy networks

venue: In International Conference on Machine Learning,

authors:
David Belanger
Bishan Yang
Andrew McCallum

title: End-to-end learning for structured prediction energy networks

venue: In Proceedings of the 34th International Conference on Machine LearningVolume

authors:
David K Duvenaud
Dougal Maclaurin
Jorge Iparraguirre
Rafael Bombarell
Timothy Hirzel
Alan Aspuru-Guzik
Ryan P Adams

title: Convolutional networks on graphs for learning molecular fingerprints

venue: Advances in Neural Information Processing Systems

authors:
Matthias Fey
Jan E. Lenssen

title: Fast graph representation learning with PyTorch Geometric

venue: In ICLR Workshop on Representation Learning on Graphs and Manifolds,

authors:
J. Gilmer
S.S. Schoenholz
P.F. Riley
O. Vinyals
G.E. Dahl

title: Neural message passing for quantum chemistry

venue: In ICML,

authors:
Assaf Gottlieb
Gideon Y Stein
Yoram Oron
Eytan Ruppin
Roded Sharan

title: Indi: a computational framework for inferring drug interactions and their associated recommendations

venue: Molecular systems biology,

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: arXiv preprint arXiv:1502.03167,

authors:
Sabrina Jaeger
Simone Fulle
Samo Turk

title: Mol2vec: Unsupervised machine learning approach with chemical intuition

venue: Journal of chemical information and modeling,

authors:
TN. Kipf
M. Welling

title: Variational graph auto-encoders

venue: arXiv preprint arXiv:1611.07308,

authors:
John Lafferty
Andrew McCallum
Fernando CN Pereira

title: Conditional random fields: Probabilistic models for segmenting and labeling sequence data

authors:
John D. Lafferty
Andrew McCallum
Fernando C.N. Pereira

venue: In Proceedings of the Eighteenth International Conference on Machine Learning,

authors:
Yann LeCun
Sumit Chopra
Raia Hadsell

title: A tutorial on energy-based learning

venue: Predicting Structured Data,

authors:
Tengfei Ma
Cao Xiao
Jiayu Zhou
Fei Wang

title: Drug similarity integration through attentive multi-view graph auto-encoders

venue: arXiv preprint arXiv:1804.10850,

authors:
Xuezhe Ma
Eduard Hovy

title: End-to-end sequence labeling via bi-directional lstm-cnns-crf

venue: arXiv preprint arXiv:1603.01354,

authors:
Sagar Maheshwari Marinka Zitnik
Rok Sosič
Jure Leskovec

title: BioSNAP Datasets: Stanford biomedical network dataset collection

authors:
Maximilian Nickel
Volker Tresp
Hans-Peter Kriegel

title: A three-way model for collective learning on multi-relational data

authors:
Evangelos E Papalexakis
Christos Faloutsos
Nicholas D Sidiropoulos

title: Tensors for data mining and data fusion: Models, applications, and scalable algorithms

venue: ACM Transactions on Intelligent Systems and Technology (TIST),

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in pytorch

authors:
F. Pedregosa
G. Varoquaux
A. Gramfort
V. Michel
B. Thirion
O. Grisel
M. Blondel
P. Prettenhofer
R. Weiss
V. Dubourg
J. Vanderplas
A. Passos
D. Cournapeau
M. Brucher
M. Perrot
E. Duchesnay

title: Scikit-learn: Machine learning in Python

venue: Journal of Machine Learning Research,

authors:
Bryan Perozzi
Rami Al-Rfou
Steven Skiena

title: Deepwalk: Online learning of social representations

venue: In Proceedings of the 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining,

authors:
Meng Qu
Yoshua Bengio
Jian Tang

title: GMNN: Graph Markov neural networks

venue: In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning,

authors:
Jae Yong Ryu
Hyun Uk Kim
Sang Yup Lee

title: Deep learning improves prediction of drug–drug and drug–food interactions

venue: Proceedings of the National Academy of Sciences,

authors:
Motoki Sato
Hiroyuki Shindo
Ikuya Yamada
Yuji Matsumoto

title: Segment-level neural conditional random fields for named entity recognition

venue: In Proceedings of the Eighth International Joint Conference on Natural Language Processing

authors:
Ioannis Tsochantaridis
Thomas Hofmann
Thorsten Joachims
Yasemin Altun

title: Support vector machine learning for interdependent and structured output spaces

venue: In Proceedings of the twentyfirst international conference on Machine learning,

authors:
Lifu Tu
Kevin Gimpel

title: Learning approximate inference networks for structured prediction

authors:
B. Wang
A. Mezlini
F. Demir
M. Fiume
Z. Tu
M. Brudno
A. Goldenberg

title: Similarity network fusion for aggregating data types on a genomic scale

venue: Nature Methods,

authors:
F. Wang
P. Li
AC. Konig

title: Learning a bi-stochastic data similarity matrix

venue: In ICDM, pp

authors:
P. Zhang
F. Wang
J. Hu
R. Sorrentino

title: Label propagation prediction of drug-drug interactions based on clinical side effects

venue: Scientific reports,

authors:
W. Zhang
Y. Chen
F. Liu
F. Luo
G. Tian
X. Li

title: Predicting potential drug-drug interactions by integrating chemical, biological, phenotypic and network data

venue: BMC Bioinformatics.,

authors:
Shuai Zheng
Sadeep Jayasumana
Bernardino Romera-Paredes
Vibhav Vineet
Zhizhong Su
Dalong Du
Chang Huang
Philip H.S. Torr

title: Conditional random fields as recurrent neural networks

authors:
Shuai Zheng
Sadeep Jayasumana
Bernardino Romera-Paredes
Vibhav Vineet
Zhizhong Su
Dalong Du
Chang Huang
Philip HS Torr

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Xiaojin Zhu
Zoubin Ghahramani
John D Lafferty

title: Semi-supervised learning using gaussian fields and harmonic functions

venue: In Proceedings of the 20th International conference on Machine learning

authors:
Marinka Zitnik
Monica Agrawal
Jure Leskovec

title: Modeling polypharmacy side effects with graph convolutional

venue: networks. Bioinformatics,

authors:
Nansu Zong
Hyeoneui Kim
Victoria Ngo
Olivier Harismendy

title: Deep mining heterogeneous networks of biomedical linked data to predict novel drug–target

venue: associations. Bioinformatics,

heading: 1 INTRODUCTION

text:
The use of drug combinations is common and often necessary for treating patients with complex diseases. However, it also increases the risk of drug-drug interactions (DDI). DDIs are pharmacological interactions between drugs that can alter the action of either or both drugs and cause adverse effects. Overall DDIs result in a large number of fatalities per year and incur ∼ $177 billion DDI associated cost annually (Giacomini et al., 2007). To mitigate these risks and costs, accurate prediction of DDI becomes a clinically important task.
While DDI knowledge is expensive to gather, several deep learning approaches have been proposed to search the large biomedical data for predicting potential DDIs (Zitnik et al., 2018; Ma et al., 2018; Ryu et al., 2018). Among them, graph neural networks (GNN) that consider DDI prediction in the graph setting have obtained great performance. In DDI graphs, drugs are represented as nodes and embedded into a low-dimensional space. Further prediction can be made based on casting DDI prediction as a link prediction problem where DDI types are represented as link types (Kipf & Welling, 2016; Zitnik et al., 2018).
However, existing works rarely exploit the correlations between these link types (e.g., DDI types) despite of their importance. For example, in Fig. 1 Warfarin is a drug for treating blood clots. Its combined use with antibiotic or nonsteroidal anti-inflammatory drugs can cause multiple DDI types including nhibition of clotting, gastrointestinal bleeding, and hemorrhage. These DDI types are correlated since they are all bleeding related DDIs caused by the increase of Warfarin’s effect. Explicit modeling such correlations can help infer unseen DDI types.
To fill the gap, we propose GENN, a new deep architecture that predicts correlated DDIs based on graph neural networks and energy based models. Specifically, we leverage the dependency structures among DDI types and formulate multi-type DDI detection as a structure prediction problem (Lafferty et al., 2001a; LeCun et al., 2006; Belanger & McCallum, 2016b). Here we use an energy based approach to incorporate such dependency structures.
To summarize, GENN is enabled by the following technical contributions.
1. Modeling link type correlations. GENN bridges graph neural networks and structure prediction to directly capture link type correlations for more accurate link prediction in graphs through minimizing an energy function. 2. A new graph neural network based energy function. Inspired by a family of structure prediction model called structured prediction energy networks (SPENs) (Belanger & McCallum, 2016a), we design a new graph energy function based on graph neural networks to capture the global dependencies among link types in DDI graphs. This is more flexible and expressive than previous approaches which used CRF-type energy definitions Qu et al. (2019). 3. Efficient semi-supervised training. We designed one cost-augmented inference network to approximate the output in training and one test inference network to approximate the output in testing. Different from the previous inference method for SPENs, we propose a new semi-supervised joint optimization procedure to adapt to our problem on graphs. The new inference method optimizes the training-set structured hinge loss together with the test-set energy with respect to the parameters for both inference networks and the energy function.
We evaluated GENN on two real world DDI datasets with both quantitative and qualitative study. Results demonstrated GENN outperformed basic graph neural networks with 13.77% and 5.01% PR-AUC improvement on the two datasets, respectively.

heading: 2 RELATED WORK

heading: 2.1 DDI PREDICTION

text: To predict unseen DDIs based on known ones, drug similarity has been learned via nearest neighbor approaches (Zhang et al., 2017), and using random walk methods including label propagation (Zhang et al., 2015; Wang et al., 2010), and unsupervised methods (Wang et al., 2014; Angione et al., 2016). Recently, deep graph neural networks have been demonstrated to provide much improved performance in DDI prediction. Among them, (Ma et al., 2018) integrates different sources of drug-related information with heterogeneous formats into an coherent and informationpreserving representation using attentive multi-view graph auto-encoders (Kipf & Welling, 2016). Decagon (Zitnik et al., 2018) develops a new graph auto-encoder approach, which allows us to develop an end-to-end trainable model for link prediction on a multi-modal graph. As the experimental results in Zitnik et al. (2018) show graph neural networks achieved significantly improved performance than shallow models such as tensor decomposition (Nickel et al., 2011; Papalexakis et al., 2017), random walk based methods (Perozzi et al., 2014; Zong et al., 2017), and non-graph neural fingerprinting (Duvenaud et al., 2015; Jaeger et al., 2018), DeepDDI (Ryu et al., 2018). None of the existing DDI prediction works consider the correlation among multiple DDI types.

heading: 2.2 STRUCTURE PREDICTION AND ENERGY BASED MODELS

text:
Structure prediction is an important problem in various application domains where we want to predict structured outputs instead of independent labels, e.g., structured labels prediction for object detection (Zheng et al., 2015b), semantic structure prediction (Belanger et al., 2017) or part-ofspeeches tagging (Ma & Hovy, 2016). For structure prediction problems, feed-forward networks is insufficient since they cannot directly model the interaction or constraints of the outputs. Instead, energy based models (LeCun et al., 2006) provide one way to address the challenge by defining a much more flexible and expressive energy score, and the prediction is conducted by minimizing the energy function.
One of the well-known structure prediction approaches is the conditional random fields (CRF) (Lafferty et al., 2001b) which show success in application areas such as named entity recognition (Sato et al., 2017) and image segmentation (Zheng et al., 2015a). But as a structured linear model, CRF has limited representation ability. To enhance the flexibility of structure prediction models, Belanger & McCallum (2016b) proposed Structured Prediction Energy Networks (SPENs) which uses arbitrary neural networks to define the energy function and optimizes the energy over a continuous relaxation of labels. Thus it can capture high-arity interaction among output labels and approximate the optimization of energy function via gradient descent and repeated “loss-augmented inference”. Belanger et al. (2017) further developed an “end-to-end” method that unroll the approximate energy optimization to a different computation graph. However, after learning the energy function, they still have to use gradient descent for test-time inference. Later, Tu & Gimpel (2018) replaced the gradient descent with a neural network trained to do inference directly. However, it separates the training of cost-augmented inference network and the fine-tuning of another inference network for testing, which makes it inefficient. In addition, how to optimize the SPENs on graphs for semi-supervised learning remains challenging. In this paper our method improves the algorithm of Tu & Gimpel (2018) to solve this problem.
Very recently, GMNN (Qu et al., 2019) combines the statistical relational learning (SRL) and graph neural networks. It includes a CRF to model the joint distribution of object labels conditioned on object attributes, and then it is trained with a variational EM algorithm. The model has been used for node classification and link classification. However, it cannot be easily extended for prediction of missing links, since they have to use links as nodes in a “dual graph” for link classification, but the missing links generally cannot be modeled as nodes.

heading: 3 PRELIMINARY

text: In this section, we first summarize our task in Section 3.1, then a basic message passing graph neural network (MPNN) is described for embedding the DDI graph in Section 3.2.

heading: 3.1 TASK FORMULATION

text:
Definition 1 (DDI Graph) Given a DDI graph G = {V, E} , V is the node set which contains N drug nodes with node features X ∈ RN×D, and E is the edge set which contains all DDIs between drug pairs. For a specific drug pair xi ∈ RD,xj ∈ RD, the DDIs of this pair ei,j can have multiple types, i.e. ei,j ∈ {0, 1}L, where D is the feature dimension and L is the total number of DDI types. And these DDI vectors e form the edge features (or labels) Y ∈ {0, 1}|E|×L.
Problem 1 (DDI Prediction) We cast the DDI prediction into a multi-type link prediction problem. We assume there are some missing edges YU in the graph, i.e. Y = {YL,YU}. Given the node features X and some known DDI links YL, the goal is to predict the unknown DDIs YU .

heading: 3.2 MESSAGE PASSING NEURAL NETWORKS FOR DDI PREDICTION

text:
Recently graph neural networks (GNNs) have been successfully applied to DDI prediction (Zitnik et al., 2018; Ma et al., 2018). To predict the unknown DDIs in a graph, a GNN first learns the node embeddings of all the drug nodes and then make prediction based on these embedding. In this
section, we describe a basic message passing based graph neural network (Gilmer et al., 2017) for drug node embedding. Considering that our graph contains different types of edges, we select a schema including edge information in the message passing. Given the graph G with node features X, for each node with hidden state hti and its neighborhood N (i), we define the message passing layer to update the node hidden state as follows:
ht+1i = Wsh t i + ∑ j∈N (i) htj · f(ei,j); ∀i ∈ V,h0i = W0xi (1)
where ei,j ∈ RL is the edge (DDI type vector) between nodes xi,xj ∈ RD, Ws ∈ RM×M ,W0 ∈ RM×D are both learnable weight, and f(·) ∈ RM×M is a neural network, i.e. MLP. The message passing layer can then be stacked to get the final state hi for each node, i.e. the node embedding. Then we can use these final states to predict the probability of DDI types for each edge:
ŷi,j = Sigmoid(W[hi||hj ] + b) (2)
where || indicates concatenation, W ∈ RL×2M and b ∈ R2M are parameters and Sigmoid is the activation function commonly used for multi-label prediction.
The model is trained on the partially known DDI edges YL by optimizing cross entropy loss LCE(ŶL,YL) = ∑ ei,j∈YL LCE(ŷi,j , ei,j) between the predictions and the true edges. After training, it can be directly used for prediction of unknown DDIs YU .
4 THE GENN METHOD
We present our GENN framework in this section (Fig. 2). We first discuss the design of energy based framework for DDI prediction and the design of our graph energy functions based on MPNN. Then we propose a joint learning strategy with a semi-supervised learning setting in Section 4.3.
Obviously the aforementioned message passing neural network does not explicitly consider the DDI correlation in the prediction. However, in practice, the labels are often correlated as we stated in the introduction section. So we re-formulate our problem as a structure prediction problem and infer the labels by optimizing an energy function. In order to get extremely powerful representation for the energy model, we follow the ideas of SPENs to formulate the energy function as a neural network.

heading: 4.1 ENERGY FUNCTION DESIGN

text:
The SPENs define the energy E(X,Y ) as a neural network that takes both the input features X and labels Y as inputs and returns the energy. Although the SPENs can almost obtain arbitrarily high expressiveness, in practice it requires a trade-off between using increasingly expressive energy networks and being more vulnerable to overfitting (Belanger & McCallum, 2016b), so the energy function need to be properly designed.
Energy function via GNN Recall the message passing neural network we used before, the edge information is also included in the neural networks. If we aggregate all the nodes’ and edges’
information to get the whole graph’s representation, it is natural to derive an “energy” over the graph. Motivated by this intuition, we formulate a new and simple energy function from a graph neural network:
E(X,Y) = Readout(G0(X,Y)) = Relu(MLP( 1
N ∑ i hi)) (3)
Where G0 is an arbitrary graph neural network (GNN) which accepts node features X and edge features Y which can be initialized by multi-hot encoding storing labels, hi is the node embedding of final output of GNN which contains information of neighbor nodes and edges, MLP is a multilayer perceptron network. Obviously, Eq. (3) contains both the interaction between each X and Y and the interaction among all labels Y. When we have an energy function, the inference of the predicted label then simply becomes to minimize the energy function:
ŶU = arg min ŶU E(X,YL, ŶU ) (4)

heading: 4.2 TRAINING COST-AUGMENTED INFERENCE NETWORK

text:
The training of SPENs usually takes the structured hinge loss (Tsochantaridis et al., 2004; Belanger & McCallum, 2016b) and relaxes the edge labels to be continuous for easier optimization. Assume the parameters for the energy function Eq. (3) is Θ, our problem becomes:
min Θ max ŶL
[4(ŶL,YL))− EΘ(X, ŶL) + EΘ(X,YL)]+ (5)
where YL are the ground-truth DDI edges for training, and ŶL are the predictions on these edges. 4 is the structured error function which returns a non-negative value to indicate the difference between the ground truth and prediction; and [·]+ means max(0, ·). In this paper, we use L1 loss as4. This above training loss is expensive to directly optimize due to the cost-augmented inference step. Following the idea of Tu & Gimpel (2018), we use a cost-augmented inference network (i.e. a graph neural network) GΦ(X,YL) to approximate the output ŶL in the training phase.
min Θ max Φ
[4(GΦ(X,YL),YL)− EΘ(X, GΦ(X,YL)) + EΘ(X,YL)]+ (6)
The problem can be seen as minimax game and optimized by alternatively optimizing Θ and Φ.
(1) Fix Θ, we can optimize the parameters for the cost-augmented inference networks max
Φ [4(GΦ(X,YL),YL)− EΘ(X, GΦ(X,YL)) + EΘ(X,YL)]+ (7)
(2) Fix Φ, optimize Θ: min
Θ [4(GΦ(X,YL),YL)− EΘ(X, GΦ(X,YL)) + EΘ(X,YL)]+ (8)

heading: 4.3 SEMI-SUPERVISED JOINT TRAINING AND INFERENCE

text:
The trained cost-augmented inference network cannot be directly used for test inference because of the cost augmentation. Notice that objective function to infer the unknown labels in the testing phase is Eq. (4). Due to the additional term in the loss function Eq. (8) during training, the derived training inference network cannot lead to a minimized energy. Thus we need another test inference network to optimize Eq. (4) in the testing phase.
We propose the test inference network ŶU = GΨ(X,YL) to directly approximate the test output. Since the test edge features are not given(instead, we only know the indices of edges to test), so this test inference network also use X and YL as its inputs. In the DDI prediction problem, the training set and test set share the same set of nodes X in a single graph. In this graph, the energy optimization in the test phase should not be regarding only YU but the all graph including both YL and the predicted ŶU , as we showed in Eq. (4). Since this energy function is a graph neural network, that is to say, we use both YL and the prediction ŶU as edge attributes for the computation of Eq. (4). When the parameter Θ for the energy function is given, the objective for this test inference network is as below:
min Ψ EΘ(X,YL, GΨ(X,YL)) (9)
The previous methods (e.g. Tu & Gimpel (2018)) first fine-tune the trained inference network on training data again with respect to the original inference objective without cost augmentation arg minΦE(X, GΦ(X,YL)). But the real effect of this fine-tuning is actually doubted (Tu & Gimpel (2018)). Also considering that the definition of training-set energy and the test-set energy are in the same graph, and the test inference network should not deviate too much from the training cost-augmented inference network, we are proposing a new joint training strategy as follows.
We design the test inference network to share some bottom layers with the training cost-augmented inference networks. Then we do the optimization with the two following steps:
(1) Fix Θ, we optimize the parameters for both the cost-augmented inference networks and the test inference network.
max Φ,Ψ
[4(GLΦ(X,YL),YL)− EΘ(X, GLΦ(X,YL)) + EΘ(X,YL)]+ − λ1EΘ(X,YL, GUΨ(X,YL))
(10)
To better discriminate the two inference networks, we useGL to indicates the predictions on training edges, and GU are the predictions on and missing edges. λ1 is a hyperparameter which is often set as 1 in practice. Notice that Φ and Ψ are not independent but share some parameters.
To better understand the benefit of the joint training, let us rewrite the parameter Φ and Ψ into {Φ0,Φ1} and {Φ0,Ψ1}, where Φ0 is the shared parameter set between Φ and Ψ for the shared layers of the two inference networks. Then the above optimization problem becomes:
max Φ0,Φ1,Ψ1
[4(GLΦ0,Φ1(X,YL),YL)− EΘ(X, G L Φ0,Φ1(X,YL)) + EΘ(X,YL)]+ (11)
−λ1EΘ(X,YL, GUΦ0,Ψ1(X,YL)) From this equation we can see the optimization of the second term, i.e. the test-set energy, can also impact the training of the cost-augmented inference network. If we fix all other parameters, the Ψ1 still gives the test inference network enough flexibility to only optimize the test energy term; meanwhile the shared layers and the shared parameter Φ0 kept the two inference networks not deviate too much.
(2) Fix Φ and Ψ, we do not want the test inference network to impact the energy function (so that the objective for optimizing Θ is still from Eq. (6), so the min function is unchanged, i.e. still Eq. (8).

heading: 5 EXPERIMENT

text:
We evaluate GENN1 model to answer the following questions:
1. Q1: Does GENN provide more accurate DDI prediction than feed-forward GNNs? If so, we can demonstrate the usefulness of the new energy function. 2. Q2: Does GENN improves the supervised inference method (Section 4.2)? If so, we can demonstrate the effectiveness of our new optimization algorithm. 3. Q3: How does GENN respond to various fraction of edge missingness? 4. Q4: Does the model really capture meaningful label correlation?

heading: 5.1 EXPERIMENTAL SETUP

text:
Data We implemented our experiments on two public datasets (with some modification of the setting): DeepDDI and Decagon. For both datasets, we randomly select 80% of the drug-drug pairs as training data, 10% as validation data and the remaining 10% as test data. For a given drug-drug pair, we predict its DDI types. Note that one drug-drug pair can have zero, one or multiple DDI types, so it is a multi-label prediction problem.
DeepDDI Dataset (Ryu et al., 2018) consists of 1861 drugs (nodes) and 222,127 drug-drug pairs (edges) from DrugBank which results in 113 different DDI types as labels. 99.87% drug-drug pairs only have one type of DDI. The edge exists when two drugs has at least one DDI type. The input feature for each drug is generated based on structural similarity profile (SSP), then its dimension
1https://github.com/ICLR2020-code/GENN.git
is reduced to 50 using principal components analysis (PCA) as suggested by DeepDDI. Moreover, Chemical-based similarity measure is used to reduce the the effect of redundant drugs on prediction accuracy where redundant similar drugs in training and test dataset will overestimate the predictive performanc as noted in Gottlieb et al. (2012). In particular, we use Tanimoto score (Tanimoto, 1957) and find 90% drug pairs has less than 0.4721 similarity score which show that there is no need to tackle the effect of redundant drugs.
BIOSNAP-sub Dataset (Marinka Zitnik & Leskovec, 2018) consists of 645 drugs (nodes) and 46,221 drug-drug pairs (edges) from TWOSIDES dataset which results in 200 different DDI types as labels. We extract data from Zitnik et al. (2018) and only keep 200 medium commonly-occurring DDI types ranging from Top-600 to Top-800 that every ddi type has at least 90 drug combinations which contains proper number of edges for fast evaluation. 73.27% drug-drug pairs have more than one type of DDI. The input feature is 32-dimension vector by transforming one-hot encoding by Gaussian random projection using Pedregosa et al. (2011).
Baselines We compared GENN with the following baselines.
1. Label Propragation (LP) (Zhu et al., 2003) is a similarity-based semi-supervised method which makes use of unlabeled data to better generalize to new samples. 2. MLP is the model used in original DeepDDI dataset which accepts pairwise features (e.g. structural similarity profile (SSP)) between two drugs to predict DDI. 3. DeepWalk learns d-dimensional neural features for nodes based on a biased random walk procedure exploring network neighborhoods of nodes. For each drug pair, we concatenate the learned DeepWalk feature vectors and its original drug feature representation and use the MLP same as the above method to do classification. 4. GNN is a basic graph neural network based on MPNN framework as introduced in 3.2. 5. GLENN is designed to be an energy model with a CRF-style local energy function and the same
inference algorithm as GENN. Notice that each node and all its connected edges form a clique (because these edges are mutual neighbors and they are also connected to each other due to the correlation), so we use E = ∑ i f1(xi + ∑ j∈N(j) f2(ei,j)). The difference between Eq. (3) and Eq. (1) is that it does not include the neighbor node features in the equation. It is linear and has only one layer. We call it a graph local energy neural network (GLENN) and use the same optimization strategy as our final model.
6. GENN− is an ablation model without semi-supervised joint training and inference. 7. GENN is our final model that incorporates the power of GNNs and energy models.
The Implementation Details can be found in Appendix B.
Metrics. To measure the prediction accuracy, we used the following metrics
1. ROC-AUC: Area under the receiver operating characteristic curve: the area under the plot of the true positive rate against the false positive rate at various thresholds. 2. PR-AUC: Area under the precision-recall curve: the area under the plot of precision versus recall curve at various thresholds. 3. P@K: short for Precision at K which is the mean percentage of correct predicted labels among TOP-K over all samples. Here, we specially choose P@1 and P@5 to validate performence as usually done in multi-label learning setting.
Following Decagon (Zitnik et al., 2018), we calculated the ROC-AUC and PR-AUC for each single label, and then use the average scores as the final values.

heading: 5.2 RESULTS

text:
Performance Comparison For each dataset, we use 3 different random splits and run all the models on that split. The results are then averaged over 3 runs and we report the mean and std values in test dataset see Table 1. Note that the results on DeepDDI dataset in Table 1 is based on the 60% randomly sampled data, because it will take more than 2 days for models to converge with limited performance gain on total dataset.
From Table. 1, we can see DeepWalk and the methods based on MPNN framework outperform MLP and LP in both datasets with a large margin. It indicates the neighbor edges and nodes information
actually help the representation learning which improve the performance. Compared with GNN, GLENN achieves nearly the same performance and gets little better score with respect to ROCAUC. As for GENN, it outperforms others with respect to all metrics in both datasets which give the positive answer to Q1 that GENN does provide more accurate DDI prediction. Especially, compared to GNN and GLENN, it demonstrates the power of our new designed graph neural network based global energy function. It is more expressive and flexible to capture the label correlation. In addition, for variants of GENN, the supervised learning model GENN− setting performs worse compared with GENN which give the answer of Q2.
Analysis of Robustness
We design experiments to show that compared with baseline models, GENN achieved empirically more robust performance in the setting of high edge missingness. We let each method be trained on 5%, 10%, and 30% of edges on DeepDDI dataset, and predict on the rest of them. From Figure. 3, we can see that GENN is better than all the baselines in every percentage point (Q3).
Correlation Analysis
In correlation analysis experiment, we chose the models trained on 60% DeepDDI dataset same as the models reported in Table. 1 and made DDI type predictions on the rest of the edges (40% left test dataset). To illustrate that GENN is able to capture correlation between labels, i.e., different DDI types by incorperating energy function, we further calculate the pearson correlation coefficent between distribution of DDI types. For each of the 113 DDI types2, we define its distribution by counting the number of occurrences in the prediction on each drug node over all edges in test dataset.
The PairGrid figure 4 shows the pairwise relation between these randomly chosed four DDI types, numbered as 17, 36, 75, 102. For instance, DDI type 17 means “Drug a may increase the cardiotoxic activities of Drug b” and DDI type 36 means “Drug a may decrease the bronchodilatory activities of Drug b”. For the ground truth, DDI type 17 has medium correlation with DDI type 36 but MLP and DeepWalk overestimate the coefficient. Especially for DDI type 102, only GLENN and GENN capture the true distribution same with the ground truth. On the one hand, it is consistent with the high performance achieved by GENN. On the other hand, the consistency between Truth and GENN demonstrates that GENN really capture some label correlation which answers Q4.

heading: 6 CONCLUSION

text: In this paper, we proposed GENN to cast DDI detection as a structure prediction problem with a new GNN based energy function. Experiments on two real DDI datasets demonstrated that GENN is superior to the models without consideration of link type correlations and achieved up to 0.8635 (13.53% relative improvement) with respect to PR-AUC. Future works includes extending the correlated GNNs to heterogeneous networks and incorporating medical domain knowledge with structure information such as drug classification ontology in the learning.

heading: C CROSS ENTROPY REGULARIZATION

text:
From the experience of Tu & Gimpel (2018), adding an local cross entropy loss to the energy function could improve the performance a lot. It can be seen as a multi-task training with two objectives: energy optimization and reconstruction error minimization. We add the reconstruction cross-entropy loss of both the training and test inference networks in Eq. (10) to train better inference network approximations.
max Φ,Ψ
[4(GUΦ(X,YL),YL)− EΘ(X, GUΦ(X,YL)) + EΘ(X,YL)]+ (12)
−λ1EΘ(X, GallΨ (X,YL))− λ2LCEΦ − λ3LCEΨ where LCEΦ = LCE(GUΦ(X,YL),YL), LCEΨ = LCE(GUΨ(X,YL),YL) and λ2, λ3 are hyperparameters which can all be set as 1 the same as λ1 in practice. Note that these additional regularization terms are independent of Θ, so we do not need to add it when minimizing over Θ (Eq. (8) is still unchanged).
3https://github.com/phanein/deepwalk

heading: D CORRELATION ANALYSIS

text: As shown in table 3, there are three pairs of DDI type which shows large (0.6423), medium (0.3347) and small postive correlation (0.2854) indicated by Truth column. For the first pair of DDI type, decreasing metabolism has large correlation with increasing the bradycardic activities in ground truth but GNN only got medium correlation coefficient much lower than GENN’. As for the last two pairs of DDI type, GNN performed bad and gave a nearly zero coefficient for the second and overestimated the coefficient for the last. In the meantime, GENN corresponds with the ground truth which in some extent gives the evidence that GENN has the power to capture correlation between labels.

title: GENN: PREDICTING CORRELATED DRUG-DRUG IN-


summaries: This paper considers unsupervised (or self-supervised) discrete representation learning of speech using a combination of a recent vector quantized neural network discritization method and future time step prediction. Discrete representations are fine-tuned by using these as input to a BERT model; the resulting representations are then used instead of conventional speech features as the input to speech recognition models. New state-of-the-art results are achieved on two datasets.

summaries: The paper proposes a way to pre-train quantized representations for speech. The approach proposed is a two-stage process: 1. train a quantized version of wav2vec [my understanding is that wav2vec is the same thing as CPC for Audio except for using a binary cross-entropy loss instead of InfoNCE softmax-cross entropy loss]. the authors propose to use gumbel softmax / VQ codebook for the vector quantization.

summaries: Though rather dense in its exposition, this paper is an interesting contribution to the area of self-supervised learning  based on discrete representations. What would make it stronger imo is to address the issue of how much is gained from a discrete vs. continuous representation. The authors take it as a given that discrete is good because it allows us to leverage work in NLP. That makes sense -- but at what cost?

abstractText: We propose vq-wav2vec to learn discrete representations of audio segments through a wav2vec-style self-supervised context prediction task. The algorithm uses either a Gumbel-Softmax or online k-means clustering to quantize the dense representations. Discretization enables the direct application of algorithms from the NLP community which require discrete inputs. Experiments show that BERT pre-training achieves a new state of the art on TIMIT phoneme classification and WSJ speech recognition.1

authors:
Dario Amodei
Sundaram Ananthanarayanan
Rishita Anubhai
Jingliang Bai
Eric Battenberg
Carl Case
Jared Casper
Bryan Catanzaro
Qiang Cheng
Guoliang Chen

title: Deep speech 2: Endto-end speech recognition in english and mandarin

authors:
Mathilde Caron
Piotr Bojanowski
Julien Mairal
Armand Joulin

title: Unsupervised pre-training of image features on non-curated data

venue: In Proceedings of the International Conference on Computer Vision (ICCV),

authors:
Jan Chorowski
Ron J. Weiss
Samy Bengio
Aäron van den Oord

title: Unsupervised speech representation learning using wavenet

venue: autoencoders. arXiv,

authors:
Yu-An Chung
James Glass

title: Speech2vec: A sequence-to-sequence framework for learning word embeddings from

venue: speech. arXiv,

authors:
Yu-An Chung
Wei-Ning Hsu
Hao Tang
James Glass

title: An unsupervised autoregressive model for speech representation learning

authors:
Ronan Collobert
Christian Puhrsch
Gabriel Synnaeve

title: Wav2letter: an end-to-end convnetbased speech recognition

venue: system. arXiv,

authors:
Ronan Collobert
Awni Hannun
Gabriel Synnaeve

title: A fully differentiable beam search

venue: decoder. arXiv,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova. Bert

title: Pre-training of deep bidirectional transformers for language

venue: understanding. arXiv,

authors:
Ewan Dunbar
Robin Algayres
Julien Karadayi
Mathieu Bernard
Juan Benjumea
Xuan-Nga Cao
Lucie Miskic
Charlotte Dugrain
Lucas Ondel
Alan W Black

title: The zero resource speech challenge 2019

venue: Tts without t. arXiv,

authors:
Ryan Eloff
André Nortje
Benjamin van Niekerk
Avashna Govender
Leanne Nortje
Arnu Pretorius
Elan Van Biljon
Ewald van der Westhuizen
Lisa van Staden
Herman Kamper

title: Unsupervised acoustic unit discovery for speech synthesis using discrete latent-variable neural networks

authors:
John S. Garofolo
David Graff
Doug Paul
David S. Pallett

title: CSR-I (WSJ0) Complete LDC93S6A

venue: Web Download. Linguistic Data Consortium,

authors:
John S. Garofolo
Lori F. Lamel
William M. Fisher
Jonathon G. Fiscus
David S. Pallett
Nancy L. Dahlgren

title: The DARPA TIMIT Acoustic-Phonetic Continuous Speech Corpus CDROM

venue: Linguistic Data Consortium,

authors:
Pegah Ghahremani
Vimal Manohar
Hossein Hadian
Daniel Povey
Sanjeev Khudanpur

title: Investigation of transfer learning for asr using lf-mmi trained neural networks

venue: In Proc. of ASRU,

title: Statistical theory of extreme values and some practical applications: a series of lectures, volume 33

venue: US Government Printing Office,

authors:
Hossein Hadian
Hossein Sameti
Daniel Povey
Sanjeev Khudanpur

title: End-to-end speech recognition using lattice-free mmi

venue: In Proc. of Interspeech,

authors:
Kenneth Heafield
Ivan Pouzyrevsky
Jonathan H. Clark
Philipp Koehn

title: Scalable modified Kneser-Ney language model estimation

venue: In Proc. of ACL,

authors:
Kazuki Irie
Rohit Prabhavalkar
Anjuli Kannan
Antoine Bruguier
David Rybach
Patrick Nguyen

title: On the choice of modeling unit for sequence-to-sequence speech recognition

venue: Sep 2019. doi: 10.21437/interspeech.2019-2277. URL http://dx.doi.org/

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with gumbel-softmax

authors:
Herve Jegou
Matthijs Douze
Cordelia Schmid

title: Product quantization for nearest neighbor search

venue: IEEE Trans. Pattern Anal. Mach. Intell.,

authors:
Mandar Joshi
Danqi Chen
Yinhan Liu
Daniel S. Weld
Luke Zettlemoyer
Omer Levy

title: Spanbert: Improving pre-training by representing and predicting

venue: spans. arXiv,

authors:
Tatiana Likhomanenko
Gabriel Synnaeve
Ronan Collobert

title: Who needs words? lexicon-free speech recognition

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

authors:
Ilya Loshchilov
Frank Hutter

title: SGDR: stochastic gradient descent with restarts

authors:
Tomas Mikolov
Ilya Sutskever
Kai Chen
Greg S Corrado
Jeff Dean

title: Distributed representations of words and phrases and their compositionality

venue: In Proc. of NIPS,

authors:
Abdelrahman Mohamed
Dmytro Okhonko
Luke Zettlemoyer

title: Transformers with convolutional context for ASR

authors:
Myle Ott
Sergey Edunov
David Grangier
Michael Auli

title: Scaling neural machine translation

venue: In Proc. of WMT,

authors:
Myle Ott
Sergey Edunov
Alexei Baevski
Angela Fan
Sam Gross
Nathan Ng
David Grangier
Michael Auli

title: fairseq: A fast, extensible toolkit for sequence modeling

venue: In Proc. of NAACL System Demonstrations,

authors:
Vassil Panayotov
Guoguo Chen
Daniel Povey
Sanjeev Khudanpur

title: Librispeech: an asr corpus based on public domain audio books

venue: In Proc. of ICASSP,

authors:
Daniel S. Park
William Chan
Yu Zhang
Chung-Cheng Chiu
Barret Zoph
Ekin D. Cubuk
Quoc V. Le

title: Specaugment: A simple data augmentation method for automatic speech recognition, 2019

authors:
Mirco Ravanelli
Philemon Brakel
Maurizio Omologo
Yoshua Bengio

title: Light gated recurrent units for speech recognition

venue: IEEE Transactions on Emerging Topics in Computational Intelligence,

authors:
Steffen Schneider
Alexei Baevski
Ronan Collobert
Michael Auli

title: wav2vec: Unsupervised pre-training for speech recognition

authors:
Rico Sennrich
Barry Haddow
Alexandra Birch

title: Neural machine translation of rare words with subword units

authors:
Andros Tjandra
Berrak Sisman
Mingyang Zhang
Sakriani Sakti
Haizhou Li
Satoshi Nakamura

title: Vqvae unsupervised unit discovery and multi-scale code2spec inverter for zerospeech challenge 2019

authors:
Aaron van den Oord
Oriol Vinyals

title: Neural discrete representation learning

venue: In Advances in Neural Information Processing Systems,

authors:
Aäron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive

venue: coding. arXiv,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N. Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Maarten Versteegh
Xavier Anguera
Aren Jansen
Emmanuel Dupoux

title: The zero resource speech challenge 2015: Proposed approaches and results

venue: Procedia Computer Science,

authors:
Neil Zeghidour
Nicolas Usunier
Iasonas Kokkinos
Thomas Schaiz
Gabriel Synnaeve
Emmanuel Dupoux

title: Learning filterbanks from raw speech for phone recognition

venue: In Proc. of (ICASSP),

heading: 1 INTRODUCTION

text:
Learning discrete representations of speech has gathered much recent interest (Versteegh et al., 2016; Dunbar et al., 2019). A popular approach to discover discrete units is via autoencoding (Tjandra et al., 2019; Eloff et al., 2019; Chorowski et al., 2019) sometimes coupled with an autoregressive model (Chung et al., 2019). Another line of research is to learn continuous speech representations in a self-supervised way via predicting context information (Chung & Glass, 2018; van den Oord et al., 2018; Schneider et al., 2019).
In this paper, we combine these two lines of research by learning discrete representations of speech via a context prediction task instead of reconstructing the input. This enables us to directly apply well performing NLP algorithms to speech data (Figure 1a).
Our new discretization algorithm, vq-wav2vec, learns discrete representations of fixed length segments of audio signal by utilizing the wav2vec loss and architecture (Schneider et al, 2019; §2). To
∗Equal contribution. †Work done during a Facebook AI residency. 1The code will be made available at http://github.com/pytorch/fairseq.
choose the discrete variables, we consider a Gumbel-Softmax approach (Jang et al., 2016) as well as online k-means clustering, similar to VQ-VAE (Oord et al., 2017; Eloff et al., 2019; §3).
We then train a Deep Bidirectional Transformer (BERT; Devlin et al., 2018; Liu et al., 2019) on the discretized unlabeled speech data and input these representations to a standard acoustic model (Figure 1b; §4). Our experiments show that BERT representations perform better than log-mel filterbank inputs as well as dense wav2vec representations on both TIMIT and WSJ benchmarks. Discretization of audio enables the direct application of a whole host of algorithms from the NLP literature to speech data. For example, we show that a standard sequence to sequence model from the NLP literature can be used to perform speech recognition over discrete audio tokens (§5, §6).

heading: 2 BACKGROUND

heading: 2.1 WAV2VEC

text:
wav2vec (Schneider et al., 2019) learns representations of audio data by solving a self-supervised context-prediction task with the same loss function as word2vec (Mikolov et al., 2013; van den Oord et al., 2018). The model is based on two convolutional neural networks where the the encoder produces a representation zi for each time step i at a rate of 100 Hz and the aggregator combines multiple encoder time steps into a new representation ci for each time step i. Given an aggregated representation ci, the model is trained to distinguish a sample zi+k that is k steps in the future from distractor samples z̃ drawn from a distribution pn, by minimizing the contrastive loss for steps k = 1, . . . ,K:
Lwav2veck = − T−k∑ i=1 ( log σ(z>i+khk(ci)) + λE z̃∼pn [log σ(−z̃>hk(ci))] ) (1)
where T is the sequence length, σ(x) = 1/(1 + exp(−x)), and where σ(z>i+khk(ci)) is the probability of zi+k being the true sample. We consider a step-specific affine transformation hk(ci) = Wkci + bk that is applied to ci (van den Oord et al., 2018). We optimize the loss L = ∑K k=1 Lk, summing (1) over different step sizes. After training, the representations produced by the context network ci are input to the acoustic model instead of log-mel filterbank features.

heading: 2.2 BERT

text: BERT (Devlin et al., 2018) is a pre-training approach for NLP tasks, which uses a transformer encoder model to build a representation of text. Transformers uses self-attention to encode the input sequence as well as an optional source sequence (Vaswani et al., 2017). The original BERT model combined two tasks for training: first, masked language modeling randomly removes some of the input tokens and the model has to predict those missing tokens. Second, next sentence prediction splices two different text passages together into a single example and the model needs to predict whether the passages are from the same document.

heading: 3 VQ-WAV2VEC

text:
Our approach, vq-wav2vec, learns vector quantized (VQ) representations of audio data using a future time-step prediction task. We follow the same architectual choices as wav2vec (§2.1) with two convolutional networks f : X 7→ Z and g : Ẑ 7→ C for feature extraction and aggregation, as well as a new quantization module q : Z 7→ Ẑ to build discrete representations (Figure 1a). We first map 30ms segments of raw speech to a dense feature representation z at a stride of 10ms using the encoder network f . Next, the quantizer (q) turns these dense representations into discrete indices which are mapped to a reconstruction ẑ of the original representation z. We feed ẑ into the aggregator g and optimize the same context prediction task as wav2vec outlined in §2.1.
The quantization module replaces the original representation z by ẑ = ei from a fixed size codebook e ∈ RV×d which contains V representations of size d. We consider the Gumbel-Softmax which is a differentiable approximation of the argmax for computing one-hot representations (§3.1; Figure 2a)
as well as online k-means clustering, similar to the vector quantized variational autoencoder (VQVAE; Oord et al., 2017; §3.2; Figure 2b). Finally, we perform multiple vector quantizations over different parts of z to mitigate mode collapse (§3.3).

heading: 3.1 GUMBEL-SOFTMAX

text:
The Gumbel-Softmax (Gumbel, 1954; Jang et al., 2016; Maddison et al., 2014) enables selecting discrete codebook variables in a fully differentiable way and we use the straight-through estimator of Jang et al. (2016). Given the dense representation z, we apply a linear layer, followed by a ReLU and another linear which outputs l ∈ RV logits for the Gumbel-Softmax. At inference, we simply pick the largest index in l. At training, the output probabilities for choosing the j-th variable are
pj = exp(lj + vj)/τ∑V
k=1 exp(lk + vk)/τ , (2)
where v = − log(− log(u)) and u are uniform samples from U(0, 1). During the forward pass, i = argmaxjpj and in the backward pass, the true gradient of the Gumbel-Softmax outputs is used.

heading: 3.2 K-MEANS

text:
The vector quantization approach of van den Oord et al. (2017) is an alternative to making the index selection procedure fully differentiable. Different to their setup, we optimize a future time step prediction loss instead of the reconstruction loss of an autoencoder.
We choose the codebook variable representation by finding the closest variable to the input features z in terms of the Euclidean distance, yielding i = argminj‖z − ej‖22. During the forward pass, we select ẑ = ei by choosing the corresponding variable from the codebook. We obtain gradients for the encoder network by back-propagating dLwav2vec/dẑ (van den Oord et al., 2017). The final loss has two additional terms:
L = K∑
k=1
Lwav2veck + ( ‖sg(z)− ẑ‖2 + γ‖z− sg(ẑ)‖2 ) , (3)
where sg(x) ≡ x, ddx sg(x) ≡ 0 is the stop gradient operator and γ is a hyperparameter. The first term is the future prediction task and gradients do not change the codebook because of the straightthrough gradient estimation of mapping z to ẑ. The second term ‖sg(z)− ẑ‖2 moves the codebook vectors closer to the encoder output, and the third term ‖z − sg(ẑ)‖2 makes sure that the encoder outputs are close to a centroid (codeword).

heading: 3.3 VECTOR QUANTIZATION WITH MULTIPLE VARIABLE GROUPS

text:
So far, we considered replacing the encoder feature vector z by a single entry ei in the codebook. This is prone to mode collapse where only some of the codewords are actually used. Previously, this problem has been mitigated by workarounds such as re-initializing codewords or applying additional regularizers to the loss function (Caron et al., 2019). In the following, we describe another strategy where we independently quantize partitions of z, similar to product quantization (Jegou et al., 2011). This results in larger dictionaries and increased downstream performance (Appendix A).
The dense feature vector z ∈ Rd is first organized into multiple groups G into the matrix form z′ ∈ RG×(d/G). We then represent each row by an integer index, and hence can represent the full feature vector by the indices i ∈ [V ]G, where V again denotes the possible number of variables for this particular group and each element ij corresponds to a fixed codebook vector. For each of the G groups, we apply either one of the two VQ approaches (§3.1 and §3.2).
The codebook itself can be initialized in two possible ways: Codebook variables can be shared across groups, i.e., a particular index in group j would reference the same vector as the same index in group j′. This yields a codebook e ∈ RV×(d/G). In contrast, not sharing the codebook variables yields a codebook of size e ∈ RV×G×(d/G). In practise, we observe that sharing the codebook variables generally yields competitive results to a non-shared representation.

heading: 4 BERT PRE-TRAINING ON QUANTIZED SPEECH

text:
Once we trained a vq-wav2vec model we can discretize audio data and make it applicable to algorithms that require discrete inputs. One possibility is to use the discretized training data and apply BERT pre-training where the task is to predict masked input tokens based on an encoding of the surrounding context (Devlin et al., 2018). Once the BERT model is trained, we can use it to build representations and feed them into an acoustic model to improve speech recognition. We follow recent advances in BERT training which only use the masked input token prediction (Liu et al., 2019).
Since each of the discretized tokens represents around 10 ms of audio it is likely too easy to predict a single masked input token. We therefore change BERT training by masking spans of consecutive discretized speech tokens, similar to Joshi et al. (2019). To mask the input sequence, we randomly sample p = 0.05 of all tokens to be a starting index, without replacement, and mask M = 10 consecutive tokens from every sampled index; spans may overlap. This makes the masked token prediction harder and we show later that it improves accuracy over masking individual tokens (§6.5).

heading: 5 EXPERIMENTAL SETUP

heading: 5.1 DATASETS

text: We generally pre-train vq-wav2vec and BERT on the full 960h of Librispeech (Panayotov et al., 2015) and after vq-wav2vec training it is discretized to 345M tokens. Where indicated we perform ablations on a clean 100h subset which is discretized to 39.9M tokens. We evaluate models on two benchmarks: TIMIT (Garofolo et al., 1993b) is a 5h dataset with phoneme labels and Wall Street Journal (WSJ; Garofolo et al. 1993a) is a 81h dataset for speech recognition. For TIMIT, we apply the standard evaluation protocol and consider 39 different phonemes. For WSJ, we train acoustic models directly on 31 graphemes, including the English alphabet, the apostrophe, the silence token and tokens for repeating characters.

heading: 5.2 VQ-WAV2VEC

text:
We adapt the fairseq implementation of wav2vec (Schneider et al., 2019; Ott et al., 2019) and use vqwav2vec/wav2vec models with 34× 106 parameters. The encoder has 8 layers with 512 channels each, kernel sizes (10,8,4,4,4,1,1,1) and strides (5,4,2,2,2,1,1,1), yielding a total stride of 160. Each layer contains a convolution, followed by dropout, group normalization with a single group (Wu & He, 2018) and a ReLU non-linearity. The aggregator is composed of 12 layers, with 512 channels, stride 1, and kernel sizes starting at 2 and increasing by 1 for every subsequent layer. The block
structure is the same as for the encoder network, except we introduce skip connections between each subsequent block.
We train with the wav2vec context prediction loss (Equation 1) for 400k updates, predicting K = 8 steps into the future and sample 10 negatives from the same audio example. Training is warmed up for 500 steps where the learning rate is increased from 1× 10−7 to 5× 10−3, and then annealed to 1× 10−6 using a cosine schedule (Loshchilov & Hutter, 2016). The batch size is 10, and we crop a random section of 150k frames for each example (approximately 9.3 seconds for 16kHz sampling rate). All models are trained on 8 GPUs.
For ablations and experiments on the 100h Librispeech subset, we use a smaller model with kernels (10,8,4,4,4) and strides (5,4,2,2,2) in the encoder and seven convolutional layers with stride one and kernel size three in the aggregator. This model is trained for 40k updates.
Gumbel-Softmax Models. We use G = 2 groups and V = 320 latents per group and the linear layer projects the features produced by the encoder into G · V = 640 logits. The Gumbel-Softmax produces a one-hot vector for each group G. The temperature τ is linearly annealed from 2 to 0.5 over the first 70% of updates and then kept constant at 0.5. This enables the model to learn which latents work best for each input before committing to a single latent. After training this model on 960h of Librispeech and quantizing the training dataset, we are left with 13.5k unique codewords combinations (out of V G = 102k possible codewords).
k-means Models. We use G = 2 groups and V = 320 variables per group. vq-wav2vec on full Librispeech yields 23k unique codewords. Following van den Oord et al. (2017), we found γ = 0.25 to be a robust choice for balancing the VQ auxiliary loss.

heading: 5.3 BERT

text:
BERT base models have 12 layers, model dimension 768, inner dimension (FFN) 3072 and 12 attention heads (Devlin et al., 2018). The learning rate is warmed up over the first 10,000 updates to a peak value of 1× 10−5, and then linearly decayed over a total of 250k updates. We train on 128 GPUs with a batch size of 3072 tokens per GPU giving a total batch size of 393k tokens (Ott et al., 2018). Each token represents 10ms of audio data.
BERT small. For ablations we use a smaller setup with model dimension 512, FFN size 2048, 8 attention heads and dropout 0.05. Models are trained for 250k updates with a batch size of 2 examples per GPU.

heading: 5.4 ACOUSTIC MODEL

text: We use wav2letter as accoustic model (Collobert et al., 2016; 2019) and train for 1,000 epochs on 8 GPUs for both TIMIT and WSJ using the auto segmentation criterion. For decoding the emissions from the acoustic model on WSJ we use a lexicon as well as a separate language model trained on the WSJ language modeling data only. We consider a 4-gram KenLM language model (Heafield et al., 2013) and a character based convolutional language model (Likhomanenko et al., 2019) and tune the models with the same protocol as Schneider et al. (2019).

heading: 6 RESULTS

heading: 6.1 WSJ SPEECH RECOGNITION

text:
We first evaluate on the WSJ speech recognition benchmark. We train a vq-wav2vec model on the unlabeled version of Librispeech, then discretize the same data with the resulting model to estimate a BERT model. Finally, we train a wav2letter acoustic model on WSJ by inputting either the BERT or vq-wav2vec representations instead of log-mel filterbanks.2
We compare to various results from the literature, including wav2vec (Schneider et al., 2019) and we consider three setups: performance without any language model (No LM), with an n-gram LM
2For vq-wav2vec we input the dense representations corresponding to the learned discrete units.
(4-gram LM) and with a character convolutional LM (Char ConvLM). We report the accuracy of wav2letter with log-mel filterbanks as input (Baseline) and wav2vec. For vq-wav2vec we first experiment with the Gumbel-Softmax, with and without a BERT base model (§5.3).
Table 1 shows that vq-wav2vec together with BERT training can achieve a new state of the art of 2.34 WER on nov92. Gains are largest when no language model is used which is the fastest setting. vq-wav2vec with Gumbel-Softmax uses only 13.5k distinct codewords to represent the audio signal and this limited set of codewords is not sufficient to outperform the baseline. However, it does enable training BERT models which require a relatively small vocabulary.
Next, we compare Gumbel-Softmax to k-means for vector quantization. For this experiment we use the faster to train BERT small configuration (§5.3). We also train a vq-wav2vec k-means model with a very large number of codewords (39.9M) to test whether a more expressive model can close the gap to wav2vec. Table 2 shows that Gumbel-Softmax and k-means clustering perform relatively comparably: in the no language model setup without BERT, Gumbel-Softmax is more accurate than k-means but these differences disappear with BERT. For 4-gram LM setup, k-means is better but those differences disappear again after BERT training. Finally, the large codeword model can substantially reduce the gap to the original wav2vec model.

heading: 6.2 TIMIT PHONEME RECOGNITION

text: Next, we experiment on the much smaller TIMIT phoneme recognition task where we also pre-train vq-wav2vec on the full Librispeech corpus. Table 3 shows that vq-wav2vec and BERT achieve a new state of the art of 11.64 PER which corresponds to a 21% reduction in error over the previous best result of wav2vec.

heading: 6.3 SEQUENCE TO SEQUENCE MODELING

text: So far we used vq-wav2vec to train BERT on discretized speech. However, once the audio is discretized we can also train a standard sequence to sequence model to perform speech recognition. In preliminary experiments, we trained an off-the-shelf Big Transformer (Vaswani et al., 2017; Ott et al., 2019) on the vq-wav2vec Gumbel-Softmax discretized Librispeech corpus and evaluated on the Librispeech dev/test sets; we use a 4k BPE output vocabulary (Sennrich et al., 2016). Table 4 shows that results are promising, even though they are not as good as the state of the art (Park et al., 2019) which depends on data augmentation that we do not use.

heading: 6.4 ACCURACY VS. BITRATE

text:
Next, we investigate how well vq-wav2vec can compress the audio data. Specifically, we train models with different numbers of groupsG and variables V to vary the size of the possible codebook size V G and measure accuracy on TIMIT phoneme recognition without BERT training.
We measure compression with the bitrate r · G log2 V at sampling rate r = 100Hz and report the trade-off between bitrate and accuracy on our phoneme recognition task. We experiment with vq-wav2vec k-means and train models with 1,2,4,8,16 and 32 groups, using 40,80,160,...,1280 variables, spanning a bitrate range from 0.53 kbit/s (G = 1, V = 40) to 33.03 kbit/s (G = 32, V = 1280). We place the quantization module after the aggregator module and train all models in the small vq-wav2vec setup (§5.2) on the 100h clean Librispeech subset.
As baselines, we consider various lossy compression algorithms applied to the TIMIT audio data and train wav2letter models on the resulting audio: Codec23 as a low bitrate codec, Opus (Terriberry & Vos, 2012) as a medium bitrate codec and MP3 and Ogg Vorbis (Montgomery, 2004) as high bitrate codecs. We use the whole spectrum of both variable and constant bitrate settings of the codecs; we encode and decode with ffmpeg (ffmpeg developers, 2016). Figure 3 shows the trade-off between the bitrate and TIMIT accuracy. Acoustic models on vq-wav2vec achieve the best results across most bitrate settings.

heading: 6.5 ABLATIONS

text: Table 5a shows that masking entire spans of tokens performs significantly better than individual tokens (M = 1). Furthermore, BERT training on discretized audio data is fairly robust to masking large parts of the input (Table 5b).

heading: 7 CONCLUSION

text:
vq-wav2vec is a self-supervised algorithm that quantizes unlabeled audio data which makes it amenable to algorithms requiring discrete data. This approach improves the state of the art on the WSJ and TIMIT benchmarks by leveraging BERT pre-training. In future work, we plan to apply
3https://github.com/drowe67/codec2
other algorithms requiring discrete inputs to audio data and to explore self-supervised pre-training algorithms which mask part of the continuous audio input. Another future work avenue is to finetune the pre-trained model to output transcriptions instead of feeding the pre-trained features to a custom ASR model.

heading: APPENDIX A NUMBER OF VARIABLES VS. GROUPS

text: We investigate the relationship between number of variables V and groups G. Table 6 shows that multiple groups are beneficial compared to a single group with a large number of variables. Table 7 shows that with a single group and many variables, only a small number of codewords survive.

title: DISCRETE SPEECH REPRESENTATIONS