summaries: The paper proposes a novel, feedforward, end-to-end trainable, deep, neural network for abstract diagrammatic reasoning with significant improvements over the state of the art. The proposed model architecture is reasonable and is designed to exploit the information present at multiple granularities – at the level of objects in the diagram, their relations across diagrams, and diagram subsets. As a multimodule neural pipeline, it seems a reasonable design. Further, it shows significant performance gains over the state of the art.

summaries: In this paper the authors solve for the task of Raven Progressive Matrices (RPM) reasoning. They do so by considering multiplexed graph networks. They present an architecture for the same. The basic premise is a combination of object level representation that is obtained by a method similar to region proposal and combining them with graph network. The approach uses gated graph networks that also uses an aggregation function. These are combined and result in node embeddings. Detailed analysis of the network is provided. This provides improved results over earlier WREN method. However, the performance is slightly lesser than another paper simultaneously submitted that achieves similar results. That approach uses transformer network for spatial attention while here the spatial attention is just based on object level representation.

title: Logic representation: Aristotelian syllogism by diagram. In Applied Computing and Information Technology

venue: Intl Conf on Computational Science/intelligence and Applied Informatics/ Intl Conf on Big Data, Cloud Computing, Data Science and Engineering,

authors:
David Barrett
Felix Hill
Adam Santoro
Ari Morcos
Timothy Lillicrap

title: Measuring abstract reasoning in neural networks

venue: In International Conference on Machine Learning,

authors:
David Eigen
Marc’Aurelio Ranzato
Ilya Sutskever

title: Learning factored representations in a deep mixture of experts

venue: arXiv preprint arXiv:1312.4314,

authors:
Will Hamilton
Zhitao Ying
Jure Leskovec

title: Inductive representation learning on large graphs

venue: In Advances in Neural Information Processing Systems,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Ronghang Hu
Jacob Andreas
Marcus Rohrbach
Trevor Darrell
Kate Saenko

title: Learning to reason: End-to-end module networks for visual question answering

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
Drew A Hudson
Christopher D Manning

title: Compositional attention networks for machine reasoning

venue: arXiv preprint arXiv:1803.03067,

authors:
Sergey Ioffe
Christian Szegedy

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

venue: arXiv preprint arXiv:1502.03167,

authors:
Max Jaderberg
Karen Simonyan
Andrew Zisserman

title: Spatial transformer networks. In Advances in neural information processing

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with gumbel-softmax

venue: arXiv preprint arXiv:1611.01144,

authors:
Justin Johnson
Bharath Hariharan
Laurens van der Maaten
Li Fei-Fei
C Lawrence Zitnick
Ross Girshick

title: Clevr: A diagnostic dataset for compositional language and elementary visual reasoning

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Ta-Chu Kao
Mason A Porter

title: Layer communities in multiplex networks

venue: Journal of Statistical Physics,

authors:
Thomas N Kipf
Max Welling

title: Semi-supervised classification with graph convolutional networks

venue: arXiv preprint arXiv:1609.02907,

authors:
Yujia Li
Daniel Tarlow
Marc Brockschmidt
Richard Zemel

title: Gated graph sequence neural networks

venue: arXiv preprint arXiv:1511.05493,

authors:
Liyuan Liu
Haoming Jiang
Pengcheng He
Weizhu Chen
Xiaodong Liu
Jianfeng Gao
Jiawei Han

title: On the variance of the adaptive learning rate and beyond

authors:
Chris J Maddison
Andriy Mnih
Yee Whye Teh

title: The concrete distribution: A continuous relaxation of discrete random variables

venue: arXiv preprint arXiv:1611.00712,

authors:
Jiayuan Mao
Chuang Gan
Pushmeet Kohli
Joshua B Tenenbaum
Jiajun Wu

title: The neuro-symbolic concept learner: Interpreting scenes, words, and sentences from natural supervision

title: The raven’s progressive matrices: change and stability over culture and time

venue: Cognitive psychology,

authors:
Shaoqing Ren
Kaiming He
Ross Girshick
Jian Sun

title: Faster r-cnn: Towards real-time object detection with region proposal networks. In Advances in neural information processing

authors:
Adam Santoro
David Raposo
David G Barrett
Mateusz Malinowski
Razvan Pascanu
Peter Battaglia
Timothy Lillicrap

title: A simple neural network module for relational reasoning

venue: In Advances in neural information processing systems,

authors:
Yuri Sato
Sayako Masuda
Yoshiaki Someya
Takeo Tsujii
Shigeru Watanabe

title: An fmri analysis of the efficacy of euler diagrams in logical reasoning

venue: In Visual Languages and Human-Centric Computing (VL/HCC),

authors:
Noam Shazeer
Azalia Mirhoseini
Krzysztof Maziarz
Andy Davis
Quoc Le
Geoffrey Hinton
Jeff Dean

title: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer

venue: arXiv preprint arXiv:1701.06538,

authors:
Xander Steenbrugge
Sam Leroux
Tim Verbelen
Bart Dhoedt

title: Improving generalization for abstract reasoning tasks using disentangled feature representations

venue: arXiv preprint arXiv:1811.04784,

authors:
Damien Teney
Lingqiao Liu
Anton van den Hengel

title: Graph-structured representations for visual question answering

authors:
Aaron van den Oord
Oriol Vinyals

title: Neural discrete representation learning

title: Attention is all you need

authors:
Duo Wang
Mateja Jamnik
Pietro Li

title: Investigating diagrammatic reasoning with deep neural networks

venue: In International Conference on Theory and Application of Diagrams,

title: Non-local neural networks

authors:
Kexin Yi
Jiajun Wu
Chuang Gan
Antonio Torralba
Pushmeet Kohli
Josh Tenenbaum

title: Neural-symbolic vqa: Disentangling reasoning from vision and language understanding

authors:
Rowan Zellers
Mark Yatskar
Sam Thomson
Yejin Choi

title: Neural motifs: Scene graph parsing with global context

authors:
Chi Zhang
Feng Gao
Baoxiong Jia
Yixin Zhu
Song-Chun Zhu

title: Raven: A dataset for relational and analogical visual reasoning

title: Architecture of a single Residual Convolution Block. Spatial Attention Object-level representation: The second approach is to use spatial attention to attend to locations of objects, and extract representations for each object attended. This is similar to object detection models such as faster R-CNN Ren et al. (2015), which use a Region Proposal Network to propose bounding boxes of objects

title: Gumbel-Sigmoid distribution Maddison et al

title: The original model

text:
ABSTRACT DIAGRAMMATIC REASONING WITH MULTIPLEX GRAPH NETWORKS
Duo Wang ∗& Mateja Jamnik & Pietro Lio Department of Computer Science and Technology University of Cambridge Cambridge, United Kingdom {Duo.Wang,Mateja.Jamnik,Pietro.Lio}@cl.cam.ac.uk
ABSTRACT
Abstract reasoning, particularly in the visual domain, is a complex human ability, but it remains a challenging problem for artificial neural learning systems. In this work we propose MXGNet, a multilayer graph neural network for multi-panel diagrammatic reasoning tasks. MXGNet combines three powerful concepts, namely, object-level representation, graph neural networks and multiplex graphs, for solving visual reasoning tasks. MXGNet first extracts object-level representations for each element in all panels of the diagrams, and then forms a multi-layer multiplex graph capturing multiple relations between objects across different diagram panels. MXGNet summarises the multiple graphs extracted from the diagrams of the task, and uses this summarisation to pick the most probable answer from the given candidates. We have tested MXGNet on two types of diagrammatic reasoning tasks, namely Diagram Syllogisms and Raven Progressive Matrices (RPM). For an Euler Diagram Syllogism task MXGNet achieves state-of-the-art accuracy of 99.8%. For PGM and RAVEN, two comprehensive datasets for RPM reasoning, MXGNet outperforms the state-of-the-art models by a considerable margin.

heading: 1 INTRODUCTION

text:
Abstract reasoning has long been thought of as a key part of human intelligence, and a necessary component towards Artificial General Intelligence. When presented in complex scenes, humans can quickly identify elements across different scenes and infer relations between them. For example, when you are using a pile of different types of LEGO bricks to assemble a spaceship, you are actively inferring relations between each LEGO brick, such as in what ways they can fit together. This type of abstract reasoning, particularly in the visual domain, is a crucial key to human ability to build complex things.
Many tests have been proposed to measure human ability for abstract reasoning. The most popular test in the visual domain is the Raven Progressive Matrices (RPM) test (Raven (2000)). In the RPM test, the participants are asked to view a sequence of contextual diagrams, usually given as a 3× 3 matrices of diagrams with the bottom-right diagram left blank. Participants should infer abstract relationships in rows or columns of the diagram, and pick from a set of candidate answers the correct one to fill in the blank. Figures 1 (a) shows an example of RPM tasks containing XOR relations across diagrams in rows. More examples can be found in Appendix C. Another widely used test for measuring reasoning in psychology is Diagram Syllogism task (Sato et al. (2015)), where participants need to infer conclusions based on 2 given premises. Figure 1c shows an example of Euler Diagram Syllogism task.
Barrett et al. (2018) recently published a large and comprehensive RPM-style dataset named Procedurally Generated Matrices ‘PGM’, and proposed Wild Relation Network (WReN), a state-of-the-art neural net for
∗Corresponding Author
RPM-style tasks. While WReN outperforms other state-of-the-art vision models such as Residual Network He et al. (2016), the performance is still far from deep neural nets’ performance on other vision or natural language processing tasks. Recently, there has been a focus on object-level representations (Yi et al. (2018); Hu et al. (2017); Hudson & Manning (2018); Mao et al. (2019); Teney et al. (2017); Zellers et al. (2018)) for visual reasoning tasks, which enable the use of inductive-biased architectures such as symbolic programs and scene graphs to directly capture relations between objects. For RPM-style tasks, symbolic programs are less suitable as these programs are generated from given questions in the Visual-Question Answering setting. In RPM-style tasks there are no explicit questions. Encoding RPM tasks into graphs is a more natural choice. However, previous works on scene graphs (Teney et al. (2017); Zellers et al. (2018)) model a single image as graphs, which is not suitable for RPM tasks as there are many different layers of relations across different subsets of diagrams in a single task.
In this paper we introduce MXGNet, a multi-layer multiplex graph neural net architecture for abstract diagram reasoning. Here ’Multi-layer’ means the graphs are built across different diagram panels, where each diagram is a layer. ‘Multiplex’ means that edges of the graphs encode multiple relations between different element attributes, such as colour, shape and position. Multiplex networks are discussed in detail by Kao & Porter (2018). We first tested the application of multiplex graph on a Diagram Syllogism dataset (Wang et al. (2018a)), and confirmed that multiplex graph improves performance on the original model. For RPM task, MXGNet encodes subsets of diagram panels into multi-layer multiplex graphs, and combines summarisation of several graphs to predict the correct candidate answer. With a hierarchical summarisation scheme, each graph is summarised into feature embeddings representing relationships in the subset. These relation embeddings are then combined to predict the correct answer.
For PGM dataset (Barrett et al. (2018)), MXGNet outperforms WReN, the previous state-of-the-art model, by a considerable margin. For ‘neutral’ split of the dataset, MXGNet achieves 89.6% test accuracy, 12.7% higher than WReN’s 76.9%. For other splits MXGNet consistently performs better with smaller margins. For the RAVEN dataset (Zhang et al. (2019)), MXGNet, without any auxiliary training with additional labels, achieves 83.91% test accuracy, outperforming 59.56% accuracy by the best model with auxiliary training for the RAVEN dataset. We also show that MXGNet is robust to variations in forms of object-level representations. Both variants of MXGNet achieve higher test accuracies than existing best models for the two datasets.

heading: 2 RELATED WORK

text:
Raven Progressive Matrices: Hoshen & Werman (2017) proposed a neural network model on Raven-style reasoning tasks that are a subset of complete RPM problems. Their model is based on Convolutional Network, and is demonstrated to be ineffective in complete RPM tasks (Barrett et al. (2018)). Mandziuk & Zychowski also experimented with an auto-encoder based neural net on simple single-shape RPM tasks. Barrett et al. (2018) built PGM, a complete RPM dataset, and proposed WReN, a neural network architecture based on Relation Network (Santoro et al. (2017)).Steenbrugge et al. (2018) replace CNN part of WReN with a pre-trained Variational Auto Encoder and slightly improved performance. Zhang et al. (2019) built RAVEN, a RPM-style dataset with structured labels of elements in the diagrams in the form of parsing trees, and proposed Dynamic Residual Trees, a simple tree neural network for learning with these additional structures. Anonymous (2020) applies Multi-head attention (Vaswani et al. (2017)), originally developed for Language model, on RPM tasks.
Visual Reasoning: RPM test falls in the broader category of visual reasoning. One widely explored task of visual reasoning is Visual Question Answering(VQA). Johnson et al. (2017) built CLEVR dataset, a VQA dataset that focuses on visual reasoning instead of information retrieval in traditional VQA datasets. Current leading approaches (Yi et al. (2018); Mao et al. (2019)) on CLEVR dataset generate synthetic programs using questions in the VQA setting, and use these programs to process object-level representations extracted with objection detection models (Ren et al. (2015)). This approach is not applicable to RPM-style problems as there is no explicit question present for program synthesis.
Graph Neural Networks: Recently there has been a surge of interest in applying Graph Neural Networks (GNN) for datasets that are inherently structured as graphs, such as social networks. Many variants of GNNs (Li et al. (2015); Hamilton et al. (2017); Kipf & Welling (2016); Veličković et al. (2017)) have been proposed, which are all based on the same principle of learning feature representations of nodes by recursively aggregating information from neighbour nodes and edges. Recent methods (Teney et al. (2017); Zellers et al. (2018)) extract graph structures from visual scenes for visual question answering. These methods build scene graphs in which nodes represent parts of the scene, and edges capture relations between these parts. Such methods are only applied to scenes of a single image. For multi-image tasks such as video classification, Wang et al. (2018b) proposed non-local neural networks, which extract dense graphs where pixels in feature maps are connected to all other feature map pixels in the space-time dimensions.

heading: 3 REASONING TASKS

heading: 3.1 DIAGRAM SYLLOGISM

text: Syllogism is a reasoning task where conclusion is drawn from two given assumed propositions (premises). One well-known example is ’Socrates is a man, all man will die, therefore Socrates will die’. Syllogism can be conveniently represented using many types of diagrams (Al-Fedaghi (2017)) such as Euler diagrams and Venn diagrams. Figure 1 (c) shows an example of Euler diagram syllogism. Wang et al. (2018a) developed Euler-Net, a neural net architecture that tackles Euler diagram syllogism tasks. However Euler-Net is just a simple Siamese Conv-Net, which does not guarantee scalability to more entities in diagrams. We show that the addition of multiplex graph both improves performance and scalability to more entities.

heading: 3.2 RAVEN PROGRESSIVE MATRICES

text: In this section we briefly describe Raven Progressive Matrices (RPM) in the context of the PGM dataset (Barrett et al. (2018)) and the RAVEN dataset (Zhang et al. (2019)). RPM tasks usually have 8 context diagrams and 8 answer candidates. The context diagrams are laid out in a 3×3 matrix C where c1,1, ..c3,2 are context diagrams and c3,3 is a blank diagram to be filled with 1 of the 8 answer candidates A = {a1, . . . , a8}. One or more relations are present in rows or/and columns of the matrix. For example, in Figure 1 (a), there is XOR relation of positions of objects in rows of diagrams. With the correct answer filled in, the third row and column must satisfy all relations present in the first 2 rows and columns (in the RAVEN dataset, relations are only present in rows). In addition to labels of correct candidate choice, both datasets also provide labels of meta-targets for auxiliary training. The meta-target of a task is a multi-hot vector encoding tuples of (r, o, a) where r is the type of a relation present, o is the object type and a is the attribute. For example, the meta-target for Figure 1 (a) encodes (XOR,Shape, Position). The RAVEN dataset also provides additional structured labels of relations in the diagram. However, we found that structured labels do not improve results, and therefore did not use them in our implementation.

heading: 4 METHOD

text:
MXGNet is comprised of three main components: an object-level representation module, a graph processing module and a reasoning module. Figure 1a shows an overview of the MXGNet architecture. The object-level representation module Fρ, as the name suggests, extracts representations of objects in the diagrams as nodes in a graph. For each diagram di ⊂ C ∪A, a set of nodes vi,j ; i = 1 . . . L, j = 1 . . . N is extracted where L is the number of layers and N is the number of nodes per layer. We experimented with both fixed and dynamically learnt N values. We also experimented with an additional ‘background’ encoder that encodes background lines (See Appendix C for an example containing background lines) into a single vector, which can be considered as a single node. The multiplex graph module Gφ, for a subset of diagrams, learns the multiplex edges capturing multiple parallel relations between nodes in a multi-layer graph where each layer corresponds to one diagram in the subset, as illustrated in Figure 1 (c). In MXGNet, we consider a subset of cardinality 3 for 3 × 3 diagram matrices. While prior knowledge of RPM rules allows us to naturally treat rows and columns in RPM as subsets, this prior does not generalise to other types of visual reasoning problems. Considering all possible diagram combinations as subsets is computationally expensive. To tackle this, we developed a relatively quick pre-training method to greatly reduce the search space of subsets, as described below.
Search Space Reduction: We can consider each diagram as node vdi in a graph, where relations between adjacent diagrams are embedded as edges edij . Note here we are considering the graph of ’diagrams’, which is different from the graph of ’objects’ in the graph processing modules. Each subset of 3 diagrams in this case can be considered as subset of 2 edges. We here make weak assumptions that edges exist between adjacent diagrams (including vertical, horizontal and diagonal direction) and edges in the same subset must be adjacent (defined as two edges linking the same node), which are often used in other visual reasoning
problems. We denote the subset of edges as {edij , edjk}. We use 3 neural nets to embed nodes, edges and subsets. We use CNNs to embed diagram nodes into feature vectors, and MLPs to embed edges based on node embeddings and subsets based on edge embeddings. While it is possible to include graph architectures for better accuracy, we found that simple combinations of CNNs and MLPs train faster while still achieving the search space reduction results. This architecture first embeds nodes, then embeds edges based on node embedding, and finally embed subsets based on edge embedding. The subset embeddings are summed and passed through a reasoning network to predict answer probability, similar to WReN (Barrett et al. (2018)). For the exact configuration of the architecture used please refer to Appendix A. For each subset{edij , edjk} , we define a gating variable Gijk, controlling how much does each subset contributes to the final result. In practice we use tanh function, which allows a subset to contribute both positively and negatively to the final summed embeddings. In training we put L1 regularization constraint on the gating variables to suppress Gijk of non-contributing subsets close to zero. This architecture can quickly discover rows and columns as contributing subsets while leaving gating variables of other subsets not activated. We describe the experiment results in section 5.1. While this method is developed for discovering reasoning rules for RPM task, it can be readily applied to any other multi-frame reasoning task for search space reduction. In the rest of the paper, we hard-gate subsets by rounding the gating variables, thereby reducing subset space to only treat rows and columns as valid subsets.
We treat the first 2 rows and columns as contextual subsets ci,j where i and j are row and column indices. For the last row and column, where the answers should be filled in, we fill in each of the 8 answer candidates, and make 8 row subsets ai, i ⊂ [1, 8] and 8 column subsets ai, i ⊂ [1, 8].
The graph module then summarises the graph of objects in a subset into embeddings representing relations present in the subset. The reasoning module Rθ takes embeddings from context rows/columns and last rows/columns with different candidate answers filled in, and produce normalised probability of each answer being true. It also predicts meta-target for auxiliary training using context rows/columns. Next, we describe each module in detail.

heading: 4.1 OBJECT-LEVEL REPRESENTATION

text:
In the PGM dataset there are two types of objects, namely ‘shapes’ and background ‘lines’. While it is a natural choice to use object-level representation on shapes as they are varying in many attributes such as position and size, it is less efficient on background lines as they only vary in colour intensity. In this section we first describe object-level representation applied to ‘shapes’ objects, and then discuss object-level representation on ’lines’ and an alternative background encoder which performs better.
In MXGNet we experiment with two types of object-level representations for ‘shapes’, namely CNN grid features and representation obtained with spatial attention. For CNN grid features, we use each spatial location in the final CNN feature map as the object feature vector. Thus for each feature maps of width W and height H , N =W ×H object representations are extracted. This type of representation is used widely, such as in Relation Network (Santoro et al. (2017)) and VQ-VAE (van den Oord et al. (2017)). For representation obtained with attention, we use spatial attention to attend to locations of objects, and extract representations for each object attended. This is similar to objection detection models such as faster R-CNN (Ren et al. (2015)), which use a Region Proposal Network to propose bounding boxes of objects in the input image. For each attended location a presence variable zpres is predicted by attention module indicating whether an object exists in the location. Thus the total number of objects N can vary depending on the sum of zpres variables. As object-level representation is not the main innovation of this paper, we leave exact details for Appendix A.1.
For background ‘lines’ objects, which are not varying in position and size, spatial attention is not needed. We experimented with a recurrent encoder with Long-Short Term Memory (Hochreiter & Schmidhuber (1997)) on the output feature map of CNN, outputting M number of feature vectors. However, in the experiment
we found that this performs less well than just feature map embeddings produced by feed-forward conv-net encoder.

heading: 4.2 MULTIPLEX GRAPH NETWORK

text:
Multiplex Edge Embedding:The object-level representation module outputs a set of representations vi,j ; i ⊂ [1, L], j ⊂ [1, N ] for ‘shapes’ objects, where L is the number of layers (cardinality of subset of diagrams) and N is the number of nodes per layer. MXGNet uses an multiplex edge-embedding network Eγ to generate edge embeddings encoding multiple parallel relation embeddings:
et(i,j),(l,k) = E t γ(P k(vi,j , vl,k)); i 6= l, t = 1 . . . T (1)
Here P t is a projection layer projecting concatenated node embeddings to T different embeddings. Et is a small neural net processing tth projections to produce the tth sub-layer of edge embeddings. Here, we restricted the edges to be inter-layer only, as we found using intra-layer edges does not improve performance but increases computational costs. Figure 2 illustrates these multiplex edge embeddings between nodes of different layers. We hypothesise that different layers of the edge embeddings encode similarities/differences in different feature spaces. Such embeddings of similarities/differences are useful in comparing nodes for subsequent reasoning tasks. For example,for Progessive relation of object sizes, part of embeddings encoding size differences can be utilized to check if nodes in later layers are larger in size. This is similar to Mixture of Experts layers (Eigen et al. (2013); Shazeer et al. (2017)) introduced in Neural Machine Translation tasks. However, in this work we developed a new cross-multiplexing gating function at the node message aggregation stage, which is described below.
Graph Summarisation: After edge embeddings are generated, the graph module then summarises the graph into a feature embedding representing relations present in the subset of diagrams. We aggregate information in the graph to nodes of the last layer corresponding to the third diagram in a row or column, because in RPM tasks the relations are in the form Diagram3 = Function(Diagram1, Diagram2). All edges connecting nodes in a particular layer vi,j ; i 6= L, to a node vL,k in the last layer L are aggregated by a function Fag composed of four different types of set operations, namely max, min, sum and mean:
fvi,k = Fag(e(i,1),(L,k) . . . e(i,1),(L,k));Fag = concat(max(),min(), sum(),mean()) (2)
We use multiple aggregation functions together because different sub-tasks in reasoning may require different types of summarization. For example, counting number of objects is better suited for sum while checking if there is a object with the same size is better suited for max. The aggregated node information from each layer is then combined with a cross-multiplexing gating function. It is named ’cross-multiplexing’ because each embeddings in the set are ’multiplexing’ other embeddings in the set with gating variables that regulate which stream of information pass through. This gating function accepts a set of summarised node embeddings {fv1,k . . . fvN,k} as input, and output gating variables for each layer of node embeddings in the set:
g1,k . . .gN,k = G(fv1,k . . . fvN,k);gi,k = {g1i,k . . . gTi,k} (3) In practice G is implemented as an MLP with multi-head outputs for different embeddings, and Sigmoid activation which constrains gating variable g within the range of 0 to 1. The node embeddings of different layers are then multiplied with the gating variables, concatenated and passed through a small MLP to produce the final node embeddings: fvk = MLP (concat({fvi,k × g(i, k)|i = 1 . . . N})). Node embeddings and background embeddings are then concatenated and processed by a residual neural block to produce final relation feature embeddings r of the diagram subset.

heading: 4.3 REASONING NETWORK

text:
The reasoning network takes relation feature embeddings r from all graphs, and infers the correct answer based on these relation embeddings. We denote the relation embeddings for context rows as rcri ; i = 1, 2 and context columns as rcci ; i = 1, 2. The last row and column filled with each answer candidate ai are denoted rari ; i = 1, . . . , 8 and r ac i ; i = 1, . . . , 8. For the RAVEN dataset, only row relation embeddings r
cr and rar are used, as discussed in Section 3.2. The reasoning network Rθ is a multi-layer residual neural net with a softmax output activation that processes concatenated relation embeddings and outputs class probabilities for each answer candidate. The exact configuration of the reasoning network can be found in Appendix A.3.
For meta-target prediction, all relation information is contained in the context rows and columns of the RPM task. Therefore, we apply a meta-predicting network Rmeta with Sigmoid output activation to all context rows and columns to obtain probabilities of each meta-target categories:
pmeta = Rmeta(r cr 1 + r cr 2 + r cc 1 + r cc 2 ) (4)

heading: 4.4 TRAINING

text: The full pipeline of MXGNet is end-to-end trainable with any gradient descent optimiser. In practice, we used RAdam optimiser (Liu et al. (2019)) for its fast convergence and robustness to learning rate differences. The loss function for the PGM dataset is the same as used in WReN (Barrett et al. (2018)): L = Lans + βLmeta−target where β balances the training between answer prediction and meta-target prediction. For the RAVEN dataset, while the loss function can include auxiliary meta-target and structured labels as L = Lans + αLstruct + βLmeta−target, we found that both auxiliary targets do not improve performance, and thus set α and β to 0.

heading: 5 EXPERIMENTS

heading: 5.1 SEARCH SPACE REDUCTION

text: The Search Space Reduction model is applied on both PGM and RAVEN dataset to reduce the subset space. After 10 epochs, only gating variables of rows and columns subset for PGM and of rows for RAVEN have value larger than 0.5. The Gating variables for three rows are 0.884, 0.812 and 0.832. The gating variables for three columns are 0.901, 0.845 and 0.854. All other gating variables are below the threshold value of 0.5. Interestingly all activated (absolute value > 0.5) gating variables are positive. This is possibly because it is easier for the neural net to learn an aggregation function than a comparator function. Exact experiment statistics can be found in Appendix D.

heading: 5.2 DIAGRAM SYLLOGISM PERFORMANCE

text: We first test how well can the multiplex graph network capture relations for the simple Diagram Syllogism task. We simply add the multiplex graph to the original Conv-Net used in (Wang et al. (2018a)). MXGNet achieved 99.8% accuracy on both 2-contour and 3-contour tasks, higher than the original paper’s 99.5% and 99.4% accuracies. The same performance on 2-contour and 3-contour tasks also show that MXGNet scales better for more entities in the diagram. For more details please refer to Appendix E.

heading: 5.3 RPM TASK PERFORMANCES

text:
In this section we compare all variants of MXGNet against the state-of-the-art models for the PGM and the RAVEN datasets. For the PGM dataset, we tested against results of WReN (Barrett et al. (2018)) in the auxiliary training setting with β value of 10. In addition, we also compared MXGNet with VAE-WReN (Steenbrugge et al. (2018))’s result without auxiliary training. For the RAVEN dataset, we compared with WReN and ResNet model’s performance as reported in the original paper (Zhang et al. (2019)). We evaluated MXGNet with different object-level representations (Section 4.1) on the test data in the ‘neutral’ split of the PGM dataset.
Table 1 (a) shows test accuracies of model variants compared with WReN and VAE-WReN for the case without auxiliary training (β = 0) and with auxiliary training (β = 10) for the PGM dataset. Both model variants of MXGNet outperform other models by a considerable margin, showing that the multi-layer graph is indeed a more suitable way to capture relations in the reasoning task. Model variants using grid features from the CNN feature maps slightly outperform model using spatial-attention-based object representations for both with and without auxiliary training settings. This is possibly because the increased number of parameters for the spatial attention variant leads to over-fitting, as the training losses of both model variants are very close. In our following experiments for PGM we will use model variants using CNN features to report performances.
Table 1 (b) shows test accuracies of model variants compared with WReN the best performing ResNet models for RAVEN dataset. WReN surprisingly only achieves 14.69% as tested by Zhang et al. (2019). We include results of the ResNet model with or without Dynamic Residual Trees (DRT) which utilise additional structure labels of relations. We found that for the RAVEN dataset, auxiliary training of MXGNet with meta-target or structure labels does not improve performance. Therefore, we report test accuracies of models trained only with the target-prediction objective. Both variants of MXGNet significantly outperform the ResNet models. Models with spatial attention object-level representations under-perform simpler CNN features slightly, most probably due to overfitting, as the observed training losses of spatial attention models are in fact lower than CNN feature models.

heading: 5.4 GENERALISATION EVALUATION FOR PGM

text:
In the PGM dataset, other than the neutral data regime in which test dataset’s sampling space is the same as the training dataset, there are also other data regimes which restrict the sampling space of training or test data to evaluate the generalisation capability of a neural network. In the main paper, due to space limitations, we selected 2 representative regimes, the ‘interpolation’ regime and the ‘extrapolation’ regime to report results. For results of other data splits of PGM, please refer to Appendix G. For ‘interpolation’ regime, in the training dataset, when attribute a = color and a = size, the values of a are restricted to even-indexed values in the spectrum of a values. This tests how well can a model ‘interpolate’ for missing values. For ‘Extrapolation’ regime, in the training dataset, the value of a is restricted to be the lower half of the value spectrum. This tests how well can a model ‘extrapolate’ outside of the value range in the training dataset. Table 2 shows validation and test accuracies for all three data regimes with and without auxiliary training. In addition, differences between validation and test accuracies are also presented to show how well can models generalise. MXGNet models consistently perform better than WReN for all regimes tested. Interesting for ’Interpolation’ regime, while validation accuracy of MXGNet is lower than WReN, the test accuracy is higher. In addition, for regime
‘Interpolation’ and ‘Extrapolation’, MXGNet also shows a smaller difference between validation and test accuracy. These results show that MXGNet has better capability of generalising outside of the training space.

heading: 6 DISCUSSION AND CONCLUSION

text:
We presented MXGNet, a new graph-based approach to diagrammatic reasoning problems in the style of Raven Progressive Matrices (RPM). MXGNet combines three powerful ideas, namely, object-level representation, graph neural networks and multiplex graphs, to capture relations present in the reasoning task. Through experiments we showed that MXGNet performs better than previous models on two RPM datasets. We also showed that MXGNet has better generalisation performance.
One important direction for future work is to make MXGNet interpretable, and thereby extract logic rules from MXGNet. Currently, the learnt representations in MXGNet are still entangled, providing little in the way of understanding its mechanism of reasoning. Rule extraction can provide people with better understanding of the reasoning problem, and may allow neural networks to work seamlessly with more programmable traditional logic engines.
While the multi-layer multiplex graph neural network is designed for RPM style reasoning task, it can be readily extended to other diagrammatic reasoning tasks where relations are present between multiple elements
across different diagrams. One example of a real-world application scenario is robots assembling parts of an object into a whole, such as building a LEGO model from a room of LEGO blocks. MXGNet provides a suitable way of capturing relations between parts, such as ways of piecing and locking two parts together.

heading: A ARCHITECTURE

text:
In this section we present exact configurations of all model variants of MXGNet. Due to the complexity of architectures, we will describe each modules in sequence. The object-level representation has two variations which are (o1) CNN features and (o2) Spatial Attention features. Also the models for PGM and RAVEN dataset differ in details. Unless otherwise stated, in all layers we apply Batch Normalization Ioffe & Szegedy (2015) and use Rectified Linear Unit as activation function.
A.1 OBJECT-LEVEL REPRESENTATION ARCHITECTURE
CNN features: The first approach applies a CNN on the input image and use each spatial location in the final CNN feature map as the object feature vector. This type of representation is used widely, such as in Relation Network Santoro et al. (2017) and VQ-VAE van den Oord et al. (2017). Formally, the output of a CNN is a feature map tensor of dimension H ×W ×D where H , W and D are respectively height, width and depth of the feature map. At each H and W location, an object vector is extracted. This type of object representation is simple and fast, but does not guarantee that the receptive field at each feature map location fully bounds objects in the image.
We use a residual module He et al. (2016) with two residual blocks to extract CNN features, as shown in figure 4.This is because Residual connections show better performance in experiments. The structure of a single Residual Convolution Block is shown in figure 3.Unless otherwise stated, convolutional layer in residual blocks has kernel size of 3× 3. The output feature map processed by another residual block is treated as background encoding because we found that convolutional background encoding gives better results than feature vectors.
Spatial Attention Object-level representation: The second approach is to use spatial attention to attend to locations of objects, and extract representations for each object attended. This is similar to object detection models such as faster R-CNN Ren et al. (2015), which use a Region Proposal Network to propose bounding boxes of objects in the input image. In practice, we use Spatial Transformer Jaderberg et al. (2015) as our spatial attention module. Figure 5 shows the architecture used for extracting object-level representation using spatial attention. A CNN composed of 1 conv layr and 2 residual blocks is first applied to the input image, and the last layer feature map is extracted. This part is the same as CNN grid feature module. A spatial attention network composed of 2 conv layer then processes information at each spatial location on the feature map, and outputs k numbers of z = (zpres, zwhere), corresponding to k possible objects at each location. Here, zpres is a binary value indicating if an object exists in this location, and zwhere is an affine
transformation matrix specifying a sampling region on the feature maps. zpres, the binary variable, is sampled from Gumbel-Sigmoid distribution Maddison et al. (2016); Jang et al. (2016), which approximates the Bernoulli distribution. We set Gumbel temperature to 0.7 throughout the experiments. For the PGM dataset we restricted k to be 1 and zwhere to be a translation and scaling matrix as ‘shapes’ objects do not overlap and do not have affine transformation attributes other than scaling and translation. For all zi; i ⊂ [1, H ×W ], if zpresi is 1, an object encoder network samples a patch from location specified by z where i using a grid sampler with a fixed window size of 4× 4 pixels. More details of the grid sampler can be found in Jaderberg et al. (2015). The sampled patches are then processed by a conv-layer to generate object embeddings.
A.2 GRAPH NETWORKS
Multiplex Edge Embeddings:Figure 2 in the main paper shows an overview of the multiplex graph architecture. While motivation and overview of architecture is explained in section 4.2 of the main paper, in this section we provide exact configurations for each part of the model. Each sub-layer of the multiplex edge is
embedded by a small MLP. For PGM dataset, we use 6 parallel layers for each multiplex edge embeddings , with each layer having 32 hidden units and 8 output units. For RAVEN dataset we use 4 layers with 16 hidden units and 8 output units because RAVEN dataset contains fewer relations types than PGM dataset. Gating function is implemented as one Sigmoid fully connected layer with hidden size equal to the length of concatenated aggregated embeddings. Gating variables are element-wise multiplied with concatenated embeddings for gating effects. Gated embeddings are then processed with a final fully connected layer with hidden size 64.
Graph Summarization: This module summarizes all node summary embeddings and background embeddings to produce a diagram subset embedding representing relations present in the set of diagrams. We experimented with various approaches and found that keeping embeddings as feature maps and processing them with residual blocks yields the best results. Background feature map embeddings are generated with one additional residual block of 48 on top of lower layer feature-extracting resnet. For object representations obtained from CNN-grid features, we can simply reshape node embeddings into a feature map, and process it with additional conv-nets to generate a feature map embeddings of the same dimension to background feature map embeddings. For object representations with spatial attention, we can use another Spatial Transformer to write node summary embeddings to its corresponding locations on a canvas feature map. Finally we concatenate node summary embeddings and background embeddings and process it with 2 residual blocks of size 64 to produce the relation embeddings.
A.3 REASONING NETWORK
Figure 6 shows the reasoning network configuration for RPM tasks. We experimented with the approach introduced in Barrett et al. (2018), which compute scores for each answer candidates and finally normalize the scores. We found this approach leads to severe overfitting on the RAVEN dataset, and therefore used a simpler approach to just concatenate all relation embeddings and process them with a neural net. In practice we used two residual blocks of size 128 and 256, and a final fully connected layer with 8 units corresponding to 8 answer candidates. The output is normalized with softmax layer. For Meta-target prediction, all context relation embeddings (context rows and columns for PGM while only rows for RAVEN dataset) are summed and fed into a fully connected prediction layer with Sigmoid activation. For PGM there are 12 different meta-targets while for RAVEN there are 9.

heading: B TRAINING DETAILS

text: The architecture is implemented in Pytorch framework. During training, we used RAdam optimizer Liu et al. (2019) with learning rate 0.0001, β1 = 0.9,β2 = 0.999. We used batch size of 64, and distributed the training across 2 Nvidia Geforce Titan X GPUs. We early-stop training when validation accuracy stops increasing.

heading: C MORE DETAILS OF RPM DATASETS

text:
In PGM dataset there are two types of elements present in the diagram, namely shapes and lines. These elements have different attributes such as colour and size. In the PGM dataset, five types of relations can be present in the task: {Progression,AND,OR,XOR,ConsistentUnion}. The RAVEN dataset, compared to PGM, does not have logic relationsAND,OR,XOR, but has additional relationsArithmetic, Constant. In addition RAVEN dataset only allow relations to be present in rows.
Figure 7a and 7b show two examples from the PGM dataset(Image courtesy Barrett et al. (2018)). The first example contains a ’Progression’ relation of the number of objects across diagrams in columns. The second examples contains a ’XOR’ relation of position of objects across diagrams in rows.
In addition to shape objects, diagrams in the PGM dataset can also contain background line objects that appear at fixed locations. Figure 8a and 8b show two examples of PGM tasks containing line objects.

heading: D MORE DETAILS ON SEARCH SPACE REDUCTION

text:
In this section we provide detailed architecture used for Search Space reduction, and present additional experimental results.
The node embeddings are generated by applying a Conv-Net of 4 convolutional layer (32 filters in each layer) of kernel size 3, and a fully connected layer mapping flattened final-layer feature maps to a feature vector of size 256. Edge embeddings are generated by a 3-layer MLP of 512 − 512 − 256 hidden units. Subset embeddings are generated by a fully connected layer of 512 units. The subset embeddings are gated with the gating variables and summed into a feature vector, which is then feed into the reasoning net, a 3-layer MLP with 256− 256− 13. The output layer contains 13 units. The first unit gives probability of currently combined answer choice being true. The rest 12 units give meta-target prediction probabilities. This is the same as Barrett et al. (2018). The training loss function is:
L = Lans + βLmeta−target + λ ∥∥∥∥∥∥ ∑
(i,j,k)⊂S
Gi,j,k ∥∥∥∥∥∥ L1
(5)
In our experiment we have tested various values of λ, and found 0.01 to be the best. This model is trained with RAdam optimizer with learning rate of 0.0001 and batch size of 64. After 10 epochs of training, only gating variables of subsets that are rows and columns are above the 0.5 threshold. The Gating variables for three rows are 0.884, 0.812 and 0.832. The gating variables for three columns are 0.901, 0.845 and 0.854. All other gating variables are below 0.5. Among these, the one with highest absolute value is 0.411. Table 3 shows the top-16 ranked subsets, with each subset indexed by 2 connecting edges in the subset. Figure 9 illustrates this way of indexing the subset. For example, the first column with red inter-connecting arrows is indexed as 0-3-6. This indicates that there two edges, one connecting diagram 0 and 3, and the other connecting diagram 3-6. Similarly the subset connected by blue arrows is indexed as 1-2-5. Note that 1-2-5 and 2-1-5 is different because the 1-2-5 contains edge 1-2 and 2-5 while 2-1-5 contains edges 1-2 and 1-5.

heading: E MORE DETAILS ON EULER DIAGRAM SYLLOGISM

text: The original model in Wang et al. (2018a) uses a Siamese Conv-Net model to process two input premise diagrams and output all consistent conclusions. Convolutional layers with shared weights are first applied to two input diagrams. The top layer feature maps are then flattened and fed into a reasoning network to make predictions. We simply use CNN grid features of the top layer feature maps as object-level representations, and use the multi-layer multiplex graph to capture object relations between the two input premise diagrams. We use a multiplex edge embeddings of 4 layers, with each layer of dimension 32. The cross-multiplexing here becomes self-multiplexing as there are only 2 diagrams (Only 1 embedding of node summary for edges from first diagram to second diagram). Final node embeddings are processed by a convolutional layer to produce the final embedding, which is also fed into the reasoning network along with the conv-net embeddings.

heading: F ABLATION STUDY

text: We performed ablation study experiments to test how much does the multiplex edges affects performance. We have tested two model variants, one without any graph modules, and the other model graphs using vanilla edge embeddings produced by MLPs, on PGM dataset. We found that without graph modules, the model only achieved 83.2% test accuracy. While this is lower than MXGNet’s 89.6%, it is still higher than WReN’s 76.9%. This is possibly because the search space reduction, by trimming away non-contributing subsets, allow the model to learn more efficiently. The graph model with vanilla edge embeddings achieves 88.3% accuracy, only slightly lower than MXGNet with multiplex edge embeddings. This shows that while general graph neural network is a suitable model for capturing relations between objects, the multiplex edge embedding does so more efficiently by allowing parallel relation multiplexing.

heading: G ADDITIONAL GENERALIZATION PERFORMANCE ON PGM DATASET

text:
Table 4 shows performance of MXGNet on other splits of PGM dataset. MXGNet consistently outperforms WReN for test accuracy, except for H.O. Triple Pairs and H.O. shape-color in the case β = 0 Additionally
here we provide the analysis according to Sec 4.2 and Sec 4.6 in Barrett et al. (2018). unfortunately sec 4.3 of this paper, namely the analysis of distractors, cannot be performed as the publicly available dataset does not include any ground truth labels about distractors, nor any labels of present objects that can be used to synthesize distractor labels. For Meta-target prediction, MXG-Net achieves 84.1% accuracy. When Metatarget is correctly predicted, the model’s target prediction accuracy increases to 92.4%. When Meta-target is incorrectly predicted, the model only has 75.6% accuracy. For three logical relations the model performs best for OR relation (95.3%), and worst for XOR relation(92.6%). Accuracy for line-type tasks (86.5%) is only
slightly better than for shape tasks (80.1%), showing that object representation with graph modeling does improve on relations between shapes. The type of relation with worst performance isConsistentUnion, with only 75.1% accuracy. This is expected as ConsistentUnion is in fact a memory task instead of relational reasoning task.

title: MULTIPLEX GRAPH NETWORKS


summaries: The authors propose the adaptive thermostat Monte Carlo sampler for feedforward neural networks. The proposed approach dynamically adjust the amount of momentum and noisy applied to each model parameter during updates. ResNet++ (ResNet without batchnorm/dropout but adding SELU, fixup and weight normalization) is introduced. Further, the authors claim that the need for hyperparameter setup is reduced provided that early stopping, stochastic regularization and carefully tuned learning rate schedules are not required.

summaries: This paper proposes a novel MCMC algorithm (ATMC) that estimates and samples from the posterior distribution of neural network weights. The motivation for this approach is that applying Bayesian inference to deep learning should lead to less overfitting and better uncertainty-calibrated models. Unlike previous work, the proposed method scales to large models (ResNet) and data sets (ImageNet).

abstractText: Bayesian inference promises to ground and improve the performance of deep neural networks. It promises to be robust to overfitting, to simplify the training procedure and the space of hyperparameters, and to provide a calibrated measure of uncertainty that can enhance decision making, agent exploration and prediction fairness. Markov Chain Monte Carlo (MCMC) methods enable Bayesian inference by generating samples from the posterior distribution over model parameters. Despite the theoretical advantages of Bayesian inference and the similarity between MCMC and optimization methods, the performance of sampling methods has so far lagged behind optimization methods for large scale deep learning tasks. We aim to fill this gap and introduce ATMC, an adaptive noise MCMC algorithm that estimates and is able to sample from the posterior of a neural network. ATMC dynamically adjusts the amount of momentum and noise applied to each parameter update in order to compensate for the use of stochastic gradients. We use a ResNet architecture without batch normalization to test ATMC on the Cifar10 benchmark and the large scale ImageNet benchmark and show that, despite the absence of batch normalization, ATMC outperforms a strong optimization baseline in terms of both classification accuracy and test log-likelihood. We show that ATMC is intrinsically robust to overfitting on the training data and that ATMC provides a better calibrated measure of uncertainty compared to the optimization baseline.

authors:
Sungjin Ahn
Anoop Korattikara
Max Welling

title: Bayesian posterior sampling via stochastic gradient fisher scoring

venue: arXiv preprint arXiv:1206.6380,

authors:
Anoop Korattikara Balan
Vivek Rathod
Kevin P Murphy
Max Welling

title: Bayesian dark knowledge

venue: In Advances in Neural Information Processing Systems,

authors:
Charles Blundell
Julien Cornebise
Koray Kavukcuoglu
Daan Wierstra

title: Weight uncertainty in neural network

venue: In International Conference on Machine Learning,

authors:
Changyou Chen
Nan Ding
Lawrence Carin

title: On the convergence of stochastic gradient mcmc algorithms with high-order integrators

authors:
Tianqi Chen
Emily Fox
Carlos Guestrin

title: Stochastic gradient hamiltonian monte carlo

venue: In International conference on machine learning,

authors:
Nan Ding
Youhan Fang
Ryan Babbush
Changyou Chen
Robert D Skeel
Hartmut Neven

title: Bayesian sampling using stochastic gradient thermostats

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Sergey Ioffe
Christian Szegedy

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

authors:
Alex Kendall
Yarin Gal

title: What uncertainties do we need in bayesian deep learning for computer vision

authors:
Durk P Kingma
Tim Salimans
Max Welling

title: Variational dropout and the local reparameterization trick

authors:
Günter Klambauer
Thomas Unterthiner
Andreas Mayr
Sepp Hochreiter

title: Selfnormalizing neural networks

authors:
Chunyuan Li
Changyou Chen
David Carlson
Lawrence Carin

title: Preconditioned stochastic gradient langevin dynamics for deep neural networks

venue: In Thirtieth AAAI Conference on Artificial Intelligence,

authors:
Xiaoyu Lu
Valerio Perrone
Leonard Hasenclever
Yee Whye Teh
Sebastian J Vollmer

title: Relativistic monte carlo

venue: arXiv preprint arXiv:1609.04388,

authors:
Yi-An Ma
Tianqi Chen
Emily Fox

title: A complete recipe for stochastic gradient mcmc

authors:
Wesley Maddox
Timur Garipov
Pavel Izmailov
Dmitry Vetrov
Andrew Gordon Wilson

title: A simple baseline for bayesian uncertainty in deep learning

authors:
Stephan Mandt
Matthew D Hoffman
David M Blei

title: Stochastic gradient descent as approximate bayesian inference

venue: The Journal of Machine Learning Research,

title: Mcmc using hamiltonian dynamics. Handbook of markov chain monte carlo

title: Dynamical theories of Brownian motion, volume 3

venue: Princeton university press,

authors:
Ian Osband
Benjamin Van Roy

title: Why is posterior sampling better than optimism for reinforcement learning

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Bala Rajaratnam
Doug Sparks

title: Mcmc-based inference in the era of big data: A fundamental analysis of the convergence complexity of high-dimensional chains

venue: arXiv preprint arXiv:1508.00947,

authors:
Tim Salimans
Durk P Kingma

title: Weight normalization: A simple reparameterization to accelerate training of deep neural networks

authors:
Nitish Srivastava
Geoffrey Hinton
Alex Krizhevsky
Ilya Sutskever
Ruslan Salakhutdinov

title: Dropout: a simple way to prevent neural networks from overfitting

authors:
Mattias Teye
Hossein Azizpour
Kevin Smith

title: Bayesian uncertainty estimation for batch normalized deep networks

venue: arXiv preprint arXiv:1802.06455,

authors:
Yu-Xiang Wang
Stephen Fienberg
Alex Smola

title: Privacy for free: Posterior sampling and stochastic gradient monte carlo

authors:
Max Welling
Yee W Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: In Proceedings of the 28th international conference on machine learning

authors:
Xiaoxia Wu
Rachel Ward
Léon Bottou

title: Wngrad: learn the learning rate in gradient descent

venue: arXiv preprint arXiv:1803.02865,

authors:
Hongyi Zhang
Yann N Dauphin
Tengyu Ma

title: Fixup initialization: Residual learning without normalization

venue: arXiv preprint arXiv:1901.09321,

authors:
Ruqi Zhang
Chunyuan Li
Jianyi Zhang
Changyou Chen
Andrew Gordon Wilson

title: Cyclical stochastic gradient mcmc for bayesian deep learning

venue: arXiv preprint arXiv:1902.03932,

heading: 1 Introduction

text:
In contrast to optimization approaches in machine learning that derive a single estimate for the weights of a neural network, Bayesian inference aims at deriving a posterior distribution over the weights of the network. This makes it possible to sample model instances from the distribution over the weights and offers unique advantages. Multiple model instances can be aggregated to obtain robust uncertainty estimates over the network’s predictions; uncertainty estimates are crucial in domains such as medical diagnosis and autonomous driving where following a model’s incorrect predictions can result in catastrophe (Kendall & Gal, 2017). Sampling a distribution, as opposed to optimizing a loss, is less prone to overfitting and more training doesn’t decrease test performance. Bayesian inference can also be applied to differential privacy, where each individual sample has increased privacy guarantees (Wang et al., 2015), and to reinforcement learning, where one can leverage model uncertainty to balance between exploration and exploitation (Osband & Van Roy, 2017).
Traditional Markov Chain Monte Carlo (MCMC) methods like HMC (Neal et al., 2011) are a standard class of methods for generating samples from the posterior distribution over model parameters. These methods are seldom applied in deep learning because they have traditionally failed to scale well with large datasets and many parameters (Rajaratnam & Sparks, 2015). Stochastic Gradient MCMC (SG-MCMC) methods have fared somewhat better in scaling to large datasets due to their close relationship to stochastic optimization methods. For example the SGLD sampler (Welling & Teh, 2011) amounts to performing stochastic gradient descent while adding Gaussian noise to each parameter update. Despite these improvements, samplers like SGLD are only guaranteed to converge to the correct
Algorithm 1 The ATMC sampler. The algorithm accepts the initialized model parameters θ0, step size h, pre-conditioner m, and momentum noise Dc.
1: procedure atmc training(θ0, h,m,D) 2: p0 ← 0 3: ξ0 ← 0 4: while t < T do 5: Gt ← minibatch gradient(θt) 6: ηt ← random normal() 7: αt ← max(D − ξt, 0) 8: βt ← αt + ξt 9: pt+h ← eβth [ pt − exp[βth]−1βt Gt + √ exp[2βth]−1 βt αt ηt
] 10: θt+h ← θt + hpt+hm 11: ξt+h ← ξt + h [ p2t+h m − 1
]

heading: 12: t← t+ h

text:
distribution when the step size is annealed to zero; additional control variates have been developed to mitigate this to some extent (Ahn et al., 2012; Ding et al., 2014).
The objective of this work is to make Bayesian inference practical for deep learning by making SG-MCMC methods scale to large models and datasets. The contributions described in this work fall in three categories. We first propose the Adaptive Thermostat Monte Carlo (ATMC) sampler that offers improved convergence and stability. ATMC dynamically adjusts the amount of momentum and noise applied to each model parameter. Secondly, we improve an existing second order numerical integration method that is needed for the ATMC sampler. Third, since ATMC, like other SG-MCMC samplers, is not directly compatible with stochastic regularization methods such as batch normalization (BatchNorm) and Dropout (see Sect. 4), we construct the ResNet++ network by taking the original ResNet architecture (He et al., 2016), removing BatchNorm and introducing SELUs (Klambauer et al., 2017), Fixup initialization (Zhang et al., 2019a) and weight normalization (Salimans & Kingma, 2016). We design ResNet++ so that its parameters are easy to sample from and the gradients are well-behaved even in the absence of BatchNorm.
We show that the ATMC sampler is able to outperform optimization methods in terms of accuracy, log-likelihood and uncertainty calibration in the following settings. First, when using the ResNet++ architecture for both the ATMC sampler and the optimization baseline, the ATMC sampler significantly outperforms the optimization baseline on both Cifar-10 and ImageNet. Secondly, when using the standard ResNet for the optimization baseline and the ResNet++ for the ATMC sampler, multiple samples of the ATMC that approximate the predictive posterior of the model are still able to outperform the optimization baseline on ImageNet. Using the ResNet++ architecture, the ATMC sampler reduces the need for hyper-parameter tuning since it does not require early stopping, does not use stochastic regularization, is not prone to over-fitting on the training data and avoids a carefully tuned learning rate decay schedule.

heading: 2 ATMC Sampler

text: In this section we define the Stochastic Differential Equation (SDE) that gives rise to the ATMC sampler described in Algorithm 1. A detailed background and framework for constructing SDEs that converge to a target distribution can be found in (Ma et al., 2015).

heading: 2.1 General form of the SDE

text:
Our starting point for constructing the ATMC sampler is the framework of Stochastic Differential Equations. We are interested in SDEs that converge to a distribution p(z) over the vector z ∈ Rd for which we can evaluate ∇ log p(z). Because only the gradient of log p(z)
is required, it is sufficient to define an energy function H(z) = − log p(z) + C up to a constant C. As a consequence, we can sample from the posterior distribution p(θ|x) by only evaluating the energy function gradient ∇H(θ) = −∇ log p(x, θ). The general form of SDEs converging to p(z) for which only the gradient of p(z) is required is as follows (Ma et al., 2015):
dz = − [D(z) +Q(z)]∇H(z)dt+ Γ(z)dt+ √
2D(z)dWt, Γi(z) = d∑ j ∂ [Dij(z) +Qij(z)] ∂zj ,
(1)
where D(z) is a positive-definite matrix that determines the amount of noise, Q(z) is a skew-symmetric matrix that mixes energy between variables, Wt is a Wiener process, and Γ(z) is a correction factor that compensates for dynamics that depend on the current state z. The ATMC sampler that we propose is an instance of (1) for specific definitions of H(z), D(z), and Q(z).

heading: 2.2 Energy Function

text:
We start by defining the energy function H(z). The energy function for the model posterior p(θ|x) is defined by the loss function L(θ) = − log p(x, θ). Because the dataset x is generally large, we would like to only evaluate a mini-batch loss L̃(θ). However, naively using a stochastic gradient in (1) will result in significant bias (Chen et al., 2014). Motivated by the Central Limit Theorem, the stochastic gradient is assumed to follow a Gaussian distribution ∇L̃(θ) ∼ N (∇L(θ), B) where the covariance B is additionally assumed to be diagonal and constant w.r.t. θ. The energy function for the ATMC sampler is defined as:
H(θ, p, ξ) = L(θ) +K(p) + 1 2
( ξ − diag(B)
2m
)2 , (2)
where p is the momentum, K(p) defines the momentum distribution, and ξ is a control variate referred to as the temperature. Both p and ξ have the same dimensionality as θ. The hyper-parameter m controls the strength of the coupling between ξ and p. The distribution of the control variate p(ξ) depends on the amount of noise B in the stochastic gradient estimate L̃(θ).

heading: 2.3 Noise robust dynamics

text:
Next we define the dynamics Q(z) and D(z) such that the SDE that results from (1) can be simulated without the need to evaluate B:
D(θ, p, ξ) = 0 0 00 α(ξ)m+ 12B 0 0 0 0  , Q(θ, p, ξ) = (0 −I 0I 0 m∇K(p) 0 −m∇K(p) 0 ) , (3)
where α(ξ) is a non-negative function that determines how the temperature ξ affects the amount of noise added to the momentum update.
We first illustrate the resulting SDE by using a simpler Gaussian momentum distribution K(p) = ‖p‖2 /(2m). Note that the variance of the momentum Var(p) = m is reused in (2) and (3) to control the strength of the coupling between ξ and p. This will result in a temperature control with a momentum friction term proportional to ξ, unlike previously reported thermostat MCMC methods (Ding et al., 2014; Lu et al., 2016) where the friction
term is proportional to ξ/m. We substitute the dynamics Q(z) and D(z) defined in (3) and energy function H(z) defined in (2) into (1):
( dθ dp dξ ) =  p/m−∇L̃(θ)− β(ξ)p p2/m− 1  dt+ 0 0 00 √2α(ξ)m 0 0 0 0  dWt, β(ξ) = α(ξ) + ξ, (4) where we use ∇L̃(θ)dt = ∇L(θ)dt+ √ BdWt to replace the gradient of the loss with the minibatch estimate. The momentum p is dampened by a friction term β(ξ) that depends on the choice of α(ξ). The stochastic gradient noise B does not show up in (4) due to the particular choice of energy function H(z) and dynamics Q(z), D(z). Note however this analysis relies on the assumption that the covariance of the stochastic gradient noise B is constant in θ and a single temperature variable per parameter can only correct for a diagonal covariance B. We do not expect that this assumption will hold in practice and the approximation will therefore lead to bias in the samples. However, annealing the step size h will reduce the error due to mini-batching together with other sources of discretization error (Welling & Teh, 2011).

heading: 2.4 Adaptive Noise Thermostat

text:
Finally, we must choose a function α(ξ) which controls the amount of noise and momentum damping β(ξ). Previous work uses the Nosé-Hoover thermostat that is defined by α(ξ) = Dc where Dc is a constant determining the amount of noise added to the momentum update (Ding et al., 2014). Although the Nosé-Hoover thermostat is able to correct the stochastic gradient noise B, the correction comes at the cost of slower convergence because additional friction β(ξ) is applied as B increases. Another drawback of the Nosé-Hoover thermostat is that it causes negative friction when ξ < −Dc. In the negative friction phase β(ξ) < 0, previous gradient terms are amplified rather than dampened. Although this behavior is mathematically sound we find that it can cause exploding momentum variables.
Our choice of α(ξ) is based on the idea that negative friction should not occur and convergence speed should not be reduced by the stochastic gradient noise. Based on this intuition, we define the ATMC sampler by α(ξ) = max(D− ξ, 0). The ATMC sampler is best characterized by the various temperature stages. For 0 < ξ < Dc the total amount of noise added to the momentum is Dc and the friction coefficient β(ξ) = Dc. At this stage, the stochastic gradient noise is compensated for by adding less noise to the momentum update. If B Dc the dominant stage will be ξ > D resulting in β(ξ) < Dc and zero noise being added to the momentum. Finally, when ξ < 0 the friction coefficient β(ξ) = Dc and the noise added to the momentum is proportional to Dc − ξ. Thus, the momentum always experiences a minimum amount of friction β(ξ) ≥ Dc determined by the hyper-parameter Dc and the noise added to the momentum update is automatically adjusted based on the amount of noise present in the stochastic gradients.

heading: 2.5 Momentum energy function

text:
Following (Lu et al., 2016), we generalize the momentum energy function K(p) to the symmetric hyperbolic distribution which is defined as follows (Lu et al., 2016):
K(p) = ∑ i mc2 [√ p2i m2c2 + 1− 1 ] , (5)
where m and c are hyper-parameters. The Gaussian kinetic energy K(p) = ‖p‖2 /(2m) is a special case obtained by taking the limit c → ∞. The magnitude of parameter updates ‖∆θ‖ is determined by the gradient of the momentum:
‖∆θ‖ = ‖∇K(p)‖ = ∥∥∥∥ pM(p) ∥∥∥∥ , M(p) = m √ p2 m2c2 + 1. (6)
Hence, the hyperbolic distribution results in relativistic momentum dynamics where the parameter updates are upper bounded by c and the pre-conditioner M(p) depends on p. The average update magnitude E [‖∇K(p)‖] ≈ 1/ √ m for c m. Consequently, the parameters m and c are interpretable hyper-parameters controlling the average and maximum parameter update per step together with the step size h.
The SDE we derive in (4) and integrate in Sec. 3 uses a Gaussian momentum energy function for clarity. Deriving ATMC with a different momentum distribution like the hyperbolic distribution amounts to substituting (2), (3), and the alternative momentum distribution into (1). For the hyperbolic distribution, the dynamic friction coefficient β(ξ) will also depend on p. For the numerical integration of (4) with a hyperbolic momentum distribution we assume β(ξ) to be constant in p.

heading: 3 Improved numerical integrator for MCMC samplers

text:
In this section we construct the numerical integrator required to numerically approximate the ATMC sampler defined in (4). An efficient numerical integrator can be constructed by splitting the SDE into two terms:
( dθ dp dξ ) =  p/m0 p2/m− 1  dt ︸ ︷︷ ︸
A
+  0−∇L̃(θ)− β(p, ξ)p 0  dt+ 0 0 00 √2α(ξ)m 0 0 0 0  dWt︸ ︷︷ ︸ B . (7)
Hence, we obtain a linear ODE in part (A) that updates the parameters θ and the thermostats ξ and a linear SDE in part (B) that updates the momentum p. The operators that simulate these dynamics exactly for a time step h are denoted φhA and φ h B , respectively. Using the Strang splitting scheme yields a second order method (Chen et al., 2015):
φh = φ h/2 B ◦ φ h A ◦ φ h/2 B . (8)
The first operator φhA is given by φhA(zt) = ( θt + h pt m pt ξt + h [ p2t m − 1 ])T . (9)
The second operator φB is an instance of the Ornstein–Uhlenbeck process which can also be computed analytically as follows (Nelson, 1967):
φhB(zt) = ( θt e β(ξt)h [ pt − γ1(ξt)∇L̃(θt) + √ γ2(ξt)α(ξt) ηt ] ξt )T ,
γa(ξt) = exp[a β(ξt)h]− 1
β(ξt) ,
(10)
where ηt is isotropic Gaussian noise. Previous work (Chen et al., 2015) on higher order integrators for samplers splits the SDE into three parts where the third term is obtained from separating the friction term from the other terms in the momentum update φB . By integrating (10) exactly the gradient step and the noise and gradient term are directly affected by the friction. An exact momentum update provides additional robustness to large gradients because the temperature will increase in order to compensate for momentum updates that would lead to excessively large steps. Another advantage of a two-way split integrator is that the first and last steps in (8) can be fused together such that only a momentum update is performed per iteration. Algorithm 1 shows the pseudocode for the ATMC sampler with the split integrator defined in (9) and (10).

heading: 4 The ResNet++ Architecture

text:
The generalization performance of large neural nets trained using optimization depend on stochastic regularization methods like Dropout (Srivastava et al., 2014) and BatchNorm
>99.999%99.99%99.9%99%90%0% confidence
>99.999%
99.99%
99.9%
99%
90%
0%
a cc
u ra
cy
SGD SGD + BatchNorm SGNHT posterior predictive ATMC posterior predictive ideal
Figure 2: Calibration plot for Cifar10 99.999%99.99%99.9%99%90%0% confidence
99.9%
99%
90%
0%
a cc
u ra
cy
SGD SGD + BatchNorm SGNHT posterior predictive ATMC posterior predictive ideal
Figure 3: Calibration plot for ImageNet
(Ioffe & Szegedy, 2015). These methods implicitly add noise into the model parameters (Kingma et al., 2015; Teye et al., 2018) and significantly boost training performance and generalization for image classifiers. These methods can be interpreted as a coarse approximation of Bayesian Inference (Kingma et al., 2015; Teye et al., 2018). But a stochastic gradient sampler like ATMC already adds the necessary amount of noise and combined with BatchNorm or Dropout it leads to underfitting. We thus define a BatchNorm free version of ResNet called ResNet++ that includes SELUs (Klambauer et al., 2017), Fixup initialization (Zhang et al., 2019a) and weight normalization (Salimans & Kingma, 2016) (see Fig. 1). We use ATMC to fill the significant gap in performance due to the absence of BatchNorm in ResNet++.

heading: 4.1 SELU

text: We find the SELU activation to work well in BatchNorm free networks. SELU forces the statistics of the activations towards zero mean and unit variance (Klambauer et al., 2017). The SELU activation function additionally has a non-zero gradient everywhere which could improve the mixing of the sampler by providing a more informative gradient.

heading: 4.2 Fixup initialization

text: ResNets are known to scale well with depth (He et al., 2016). However, the additive effect of the residual branch causes the magnitudes of the activations to increase with the number of residual connections. Fixup is a recently proposed initialization method that mitigates the exploding residual branch problem without using BatchNorm (Zhang et al., 2019a). We use a simplified version of Fixup by initializing the scales of the final layer in each residual branch to a small constant.

heading: 4.3 Weight normalization

text:
We use weight normalization (Salimans & Kingma, 2016) to separate the direction and scale of each linear feature vector
θ(i) = φ(i)s φ (i) d∥∥∥φ(i)d ∥∥∥ , (11)
where φ (i) d is the direction vector and φ (i) s is the magnitude of a feature vector θ(i). Weight normalization does not depend on batch statistics and is compatible with MCMC methods.
The scale of the direction vector does not affect the outputs of the model. It does however affect the effective step size (Wu et al., 2018). Therefore the prior on the direction vector φ (i) d is chosen such that it is forced to unit length
p(φ (i) d ) ∝ exp
[ −d
2
(∥∥∥φ(i)d ∥∥∥2 − 1)2 ] . (12)
The prior on the scales p(φs) is problem-specific and can for example be chosen to encode a preference for structurally sparse models.

heading: 5 Experiments

text:
The experiments presented here aim to demonstrate that the ATMC sampler is competitive with a well-tuned optimization baseline for large-scale datasets and models. We use the TensorFlow official implementation of ResNet-56 and ResNet-50 on Cifar10 and ImageNet, respectively. We compare our ATMC sampler to an optimization baseline with and without BatchNorm. For the optimization baseline without BatchNorm we use the ResNet++ architecture as described in Sec. 4. For the baseline with BatchNorm we found standard ResNet with Xavier initialization and the ReLU non-linearity to work better.
For the ATMC sampler we report both the performance of a single sample and the estimated posterior predictive based on a finite number of samples. Similar to earlier work (Zhang et al., 2019b) we found that many fewer samples are needed when a cyclic step size ht = h0 ∗ 12 [1 + cos(πmod[t, n])] with cycle length n is used. The final sample in each cycle is used to estimate the posterior predictive.
For ResNet++ we further use a group Laplace prior p(θi) ∝ exp(−‖θi‖ /b) with b = 5 to regularize the scales of each linear feature in ResNet++. The momentum noise is chosen as Dc = − log(0.9)/h0 such that the friction applied to the momentum is at least 0.9.

heading: 5.1 Cifar 10

text:
For Cifar10 we choose the step size h0 = 0.001 and the cycle length is set to 50 epochs. The momentum hyper-parameters are m = (0.0003/h0)
−2 and c = 0.001/h0 such that the average speed and maximum speed per step are 0.0003 and 0.001, respectively. The number of convolution filters is doubled to 32 compared to the original ResNet-56 implementation. We use a single V100 GPU with a batch size of 128. The sampler runs for 1000 epochs and we start collecting samples for the posterior predictive after 150 epochs. The optimization baseline converges in 180 epochs. We also report the results of sampling with a sampler based Nosé-Hoover thermostats (SGNHT) (Ding et al., 2014; Lu et al., 2016) applied to the ResNet++ architecture.
Table 1 lists the test set performance for Cifar10. A single sample from the posterior already outperforms the baseline without BatchNorm by a significant margin in both test accuracy and log-likelihood. Using BatchNorm significantly improves the generalization of the optimization baseline. It outperforms the estimate of the posterior predictive in accuracy yet it does not have a better test log-likelihood.
To further analyze the quality of the uncertainty estimates, we group each model’s prediction in 8 equally sized bins based on the confidence p(ω̂i|xi) where ω̂i is the maximum probability class for example xi. If the probabilities are well-calibrated, the average confidence should be close to the average accuracy. Figure 2 shows the calibration of the uncertainty estimates for the posterior predictive and optimization baselines. The posterior predictive is calibrated for the least confident predictions p(ω̂i|xi) < 0.9 and shows less bias towards overconfidence compared to the models trained with SGD.

heading: 5.2 ImageNet

text:
For the ImageNet experiments we use an initial step size h0 = 0.0005 and a cycle length of 20 epochs. The other hyper-parameters for the sampler are the same as for the Cifar10 experiments. We use a a single Google Cloud TPUv3 with a batch size of 1024. We did not observe a significant difference in wall clock time per training step between SGD and ATMC on the same model. Each training step using ResNet+BatchNorm model takes 20% longer in wall clock time compared to a single train step using ResNet++. Samples for the posterior predictive are collected after 150 epochs and the sampler runs for 1000 epochs. The optimization baseline converges in 90 epochs.
Table 2 lists the results for ImageNet classification. A single sample from the posterior outperforms the optimization baseline without BatchNorm. The posterior predictive based on ATMC outperforms the optimizer with BatchNorm by a wide margin in both accuracy and test log-likelihood. We note that the sampler runs significantly longer (10x) compared to the optimization baseline because it takes a long time for the posterior predictive estimate to converge. However, the posterior predictive of ATMC matches the accuracy of the optimization baseline with BatchNorm (76.2%) after 240 epochs.
Figure 3 shows the quality of the uncertainty for various levels of confidence. Again, the ATMC based posterior predictive produces much better calibrated predictions and is almost perfectly calibrated for low confidence predictions p(ω̂i|xi) < 0.9 and shows less bias towards overconfidence compared to the optimization baseline.

heading: 6 Discussion

text:
The empirical results show it is possible to sample the posterior distribution of neural networks on large scale image classification problems like ImageNet. A major obstacle for sampling the posterior of ResNets in particular is the lack of compatibility with BatchNorm. Using recent advances in initialization and the SELU activation function we are able to stabilize and speed up training of ResNets without resorting to BatchNorm. Nonetheless, we observe that BatchNorm still offers a unique advantage in terms of generalization per-
formance. We hope that future work will allow the implicit inductive bias that BatchNorm has to be transferred into an explicit prior that is compatible with sampling methods.
Multiple posterior samples provide a much more accurate estimate of the posterior predictive, and consequently much better accuracy and uncertainty estimates. For inference, making predictions using a large ensemble of models sampled from the posterior can be costly. Variational Inference methods can be used to quickly characterize a local mode of the posterior (Blundell et al., 2015). More recent work shows that a running estimate of the mean and variance of the parameters during training can also be used to approximate a mode of the posterior (Maddox et al., 2019). Methods like distillation could potentially be used to compress a high-quality ensemble into a single network with a limited computational budget (Balan et al., 2015).
Although the form in (4) is very general, alternative methods for dealing with stochastic gradients have been proposed in the literature. One approach is to estimate the covariance of the stochastic gradient noise B explicitly and use it correct and pre-condition the sampling dynamics (Ahn et al., 2012; Li et al., 2016).
Other sampling methods are not based on an SDE that converges to the target distribution. Under some conditions stochastic optimization methods can be interpreted as such a biased sampling method (Mandt et al., 2017). Predictions based on multiple samples from the trajectory of SGD have been used successfully for obtaining uncertainty estimates in large scale Deep Learning (Maddox et al., 2019). However, these methods rely on tuning hyperparameters in such a way that just the right amount of noise is inserted.

heading: 7 Conclusion

text: This work introduces the ATMC sampler, a robust posterior sampling method that scales to large deep learning problems. To the best of our knowledge, we are the first to successfully train neural networks using MCMC on ImageNet. In a BatchNorm free setting, a single sample from the posterior generated by ATMC outperforms the optimization baseline. A posterior predictive estimate outperforms the optimization baseline with BatchNorm on ImageNet. Based on these empirical results we hope the ATMC sampler will enable new applications of Bayesian inference in deep learning.

title: Bayesian Inference for Large Scale Image Classification