Learning Permutation Invariant Representations using Memory Networks

Deep^††${}^{\ast }$Authors have contributed equally. artificial neural networks have achieved impressive performance for representation learning. The majority of these deep architectures take a single instance as an input for learning tasks. However, we often need to learn representations for unordered sequential data, or exchangeable sequences. Exchangeable data arise in many practical scenarios, and can particularly be used in Multiple Instance Learning (MIL) where a label is associated with a set, instead of a single input. An example of MIL would be classification of high resolution histopathology images, known as whole-slide images (WSIs). These images are extremely high resolution $\approx$50,000 $\times$ 50,000 pixels, with labels for an entire image instead of patch-level annotations. Patches could be extracted from these WSIs using various histological features. A set of patches from a WSI can be used for set classification/regression tasks; one such example of LUAD/LUSC classification is explained in Figure 1. Recurrent Neural Networks (RNNs) are popular approach to learn representation from sequential ordered instances. However, the lack of permutation invariance renders RNNs ineffective for exchangeable sequences.

Exchangeable models are invariant to permutations of individual instances, i.e., the output of a network remains the same for all permutations of its input set. For example, finding a maximum value in a set of numbers constitutes such a case. Zaheer et al. [28] proposed an exchangeable model, called Deep Sets, for learning set functions. They proved that a pooling operation in a latent space can approximate any arbitrary set function. This implies that our method could an universal approximation of any set function.

In this paper, we propose a novel architecture for exchangeable sequences incorporating attention over the instances to learn inter-dependencies. Our main contribution is a sequence-to-sequence permutation invariant layer called the Memory Block. The proposed model uses series of connected memory block layers, to model complex dependencies within an input set using self attention mechanism. Our model achieves state-of-the-art performance for LUAD/LUSC classification of WSIs with 84.84% accuracy.

The paper is structured as follows: section 2 discusses the related and recent works. We cover the mathematical concepts for exchangeable data in section 3. section 4 and section 5 discuss the proposed model and experiments.

In statistics, exchangeability has been long studied. De Fenetti studied exchangeable random variables and showed that sequence of infinite exchangeable random variables can be factorised into some independent and identically distributed mixtures conditioned on some parameter $\theta$. Bayesian sets [6] introduced a method to model exchangeable sequences of binary random variables by analytically computing the integrals in de Fenitti’s theorem. Orbanz et. al. [16] used de Fenetti’s theorem for Bayesian modelling of graphs, matrices, and other data that can be modeled by random structures. Considerable work has also been done on partially exchangeable random variables [1].

Symmetry in neural networks was first proposed by Shawe et al. [20] under the name Symmetry Network. They proposed that invariance can be achieved by weight-preserving automorphisms of a neural network. Ravanbaksh et al. proposed a similar method for equivariance network through parameter sharing [18]. Bloem Reddy et al. [3] studied the concept of symmetry and exchangeability for neural networks in detail and established a link between functional and probabilistic symmetry, and obtained generative functional representations of joint and conditional probability distributions that are invariant or equivariant under the action of a compact group. Zhou et. al. [30] proposed treating instances in a set as non identical and independent samples for multi instance problem.

Most of the work published in recent years have focused on ordered sets. Vinyals et. al. introduced Order Matter: Sequence to Sequence for Sets in 2016 to learn a sequence to sequence mapping. Many related models and key contributions have been proposed that uses the idea of external memories like RNNSearch [2], Memory Networks [25, 23] and Neural Turing Machines [8]. Recent interest in exchangeable models was developed due to their application in MIL. Deep Symmetry Networks  [5] used kernel-based interpolation to tractably tie parameters and pool over symmetry spaces of any dimension. Deep Sets [28] by Zaheer et al. proposed a permutation invariant model. They proved that any pooling operation (mean, sum, max or similar) on individual features is a universal approximator for any set function. They also showed that any permutation invariant model follows de Fenitti’s theorem.Work has also been done on learning point cloud classification which is an example of MIL problem. Deep Learning with Sets and Point Cloud [17] used parameter sharing to get a equivariant layer. Another important paper on exchanegable model is Set Transformer. Set Transformer [14] by Lee et al. used results from Zaheer et al. [28] and proposed a Transformer [23] inspired permutation invariant neural network. The Set Transformer uses attention mechanisms to attend to inputs in order to invoke activation. Instead of using averaging over instances like in Deep Sets, the Set Transformer uses a parametric aggregating function pool which can adapt to the problem at hand. Another way to handle exchangeable data is to modify RNNs to operate on exchangeable data. BRUNO [13] is a model for exchangeable data and makes use of deep features learned from observations so as to model complex data types such as images. To achieve this, they constructed a bijective mapping between random variables $x_i$ $\in X$ in the observation space and features $z_i$ $\in Z$, and explicitly define an exchangeable model for the sequences $z_1,z_2,z_3,\dots ,z_n$. Deep Amortized Clustering [15] proposed using Set Transformers to cluster sets of points with only few forward passes. Deep Set Prediction Networks [29] introduced an interesting approach to predict sets from a feature vector which is in contrast to predicting an output using sets.

In this section, we explain the general concepts of exchangeability, its relation to de Fenetti’s theorem, and briefly discuss Memory Networks.

A model represented by a function $f:X\to Y$ where $X$ is a set, is said to be permutation equivariant if permutation of input instances permutes the output labels with the same permutation $\pi$. Mathematically, a permutation-equivariant model is represented as,

Similarly, a function is permutation invariant if permutation of input instances does not change the output of the model. Mathematically,

Deep Sets [28] incorporate a permutation-invariant model to learn arbitrary set functions by pooling in a latent space. The authors further showed that any pooling operation such as averaging and max on individual instances of a set can be used as a universal approximator for any arbitrary set function. The authors proved that exchangeability implies the following two theorems.

Theorem 1 is linked to de Finetti’s theorem, which states that a random infinitely exchangeable sequence can be factorised into mixture densities conditioned on some parameter $\theta$ which captures the underlying generative process i.e..

Taking an entire sequence and representing it as a single input offers the above-mentioned advantages with respect to memory, but of course, it has challenges. One of the main difficulties is that the order in the sequence is lost, as well as the local proximity of sequence elements, which is crucial for most sequential inputs.

This section discusses the motivations, components, and offers an analysis of the proposed Memory-based Exchangeable Model (MEM)—a neural network designed to work with infinitely exchangeable sequences.

We performed two series of experiments comparing MEM against the simple pooling operations. In the first series of experiments, we established the learning ability of the proposed model using toy datasets. For the second experiment series, we used our model for classification of subtypes of lung cancer against the largest public dataset of histopathology whole slide images (WSIs) [24].

We experimented with different choices of pooling operations—max, mean, dot product, and sum. Our model also has a special pooling $mb1$, i.e., a memory block with a single memory unit. Therefore, we tested 9 different models for each experiment—five configurations of our model, and four configurations with just pooling operations.

In this paper, we introduced Memory-based Exchangeable Model (MEM) for learning permutation invariant representations. The proposed method uses attention mechanisms over “memories” (higher order features) for modelling complicated interactions among elements of a set. Typically for MIL, instances are treated as independently and identically distributed. However, instances are rarely independent in real tasks, and we overcome this limitation using “attention” mechanism in memory units, that exploits relations among instances. We also provided some theoretical properties of our model, including the fact that it is a universal approximator for permutation invariant functions. We achieved good performance on all problems that requires exploiting instance relationships. Our model scales well on real world problems as well, achieving accuracy score of 84.84% on classifying lung cancer subtypes on the largest public repository of histopathology images.