Embedding models through the lens of Stable Coloring

1 INTRODUCTION . Embedding models are ubiquitous in wide range of real-world applications such as information retrieval ( Zuccon et al. , 2015 ) , natural language processing ( NLP ) ( Mikolov et al. , 2013a ; b ) , graph classification ( Grover & Leskovec , 2016 ; Hamilton et al. , 2017 ) and many more . These models map categorical entities to continuous dense representations ( typically in Rd ) which provide a continuous measure of semantic similarity across categorical entities . Nowadays , there is a heavy dependence on unsupervised pre-trained embedding models across domains like Transformers in NLP ( Devlin et al. , 2019 ) , Visual Transformers ( ViT ) in Computer Vision ( Dosovitskiy et al. , 2020 ) , Graph Neural Networks ( Hamilton et al. , 2017 ; Xu et al. , 2019b ) since they learn rich semantic representations of entities from massive amounts of unlabelled data . With little finetuning , these models achieve stateof-the-art results on most of the supervised downstream tasks like sentiment analysis ( Xu et al. , 2019a ) , object detection ( Beal et al. , 2020 ) , and graph classification ( Xu et al. , 2019b ) . Historically , embedding models were developed almost independently across structured domains such as NLP , graphs , images , and so on . These algorithms use the neighborhood structure around an entity to obtain the embedding for the entity . Interestingly , a popular hypothesis in NLP - Distributional Hypothesis states that the “ meaning ” of the word is determined by its context ( neighbors ) ( Harris , 1954 ; Sahlgren , 2008 ) . This hypothesis forms the basis of most unsupervised embedding learning models in NLP ( Mikolov et al. , 2013a ; b ; Pennington et al. , 2014b ; Bojanowski et al. , 2017 ; Sonkar et al. , 2020 ) . Similarly , non-local-means , a denoising algorithm in signal processing , tries to find pixels that should be the same based on the similarity of its neighborhood structure ( patch of the image around the pixel in this case ) ( Awate & Whitaker , 2006 ; Buades et al. , 2005 ) . Even Graph neural network ( GNN ) architectures ensure that information of the surrounding neighbors is systematically incorporated in the embedding of a node , even in supervised settings ( Hamilton et al. , 2017 ; Maron et al. , 2019 ; Xu et al. , 2019b ) . Thus somehow various communities working across diverse domains have narrowed down on an entity ’ s neighborhood structure to define the entity ’ s meaning . In this paper , we try to establish this common principle on mathematically robust grounds . Structured domains can be easily represented as graphs with relations between entities as edges in the graphs . For example in NLP words can be treated as nodes of some graph , and co-occurrence relation between words can be represented as an edge . As mentioned before that graph embedding architectures capture the topological structure around the node in the node embeddings , and if structured domains can be represented as graphs , this raises the question do embedding models from structured domains like NLP and Vision also operate on some domain induced graph and capture neighborhood structural properties in their entity embeddings since we have seen that embedding models across these domains tend to capture “ neighborhood ” information ? To answer this question , we dive into combinatorial graph theory to understand how to define the notion of structural equivalence a.k.a . isomorphic structures . Weisfeiler-Lehman ( 1-WL ) algorithm ( or color refinement algorithm ) is the most popular heuristic used to identify graph isomorphism ( Weisfeiler & Leman , 1968 ) and can distinguish a broad class of graphs ( Babai & Kucera , 1979 ) . The fixed point solution of 1-WL is called a stable-coloring and has the property that any two nodes with the same color have the same multi-set of colors in their neighborhood . In essence it means if two nodes have the same color , the graph looks structurally identical from these nodes . Finally to answer the question if embedding models from structured domains capture the domain induced graph ’ s topological properties in their entity embeddings , one can find links between the mechanics of these models and stable coloring / 1-WL algorithm . In this paper , we establish this connection by providing a general framework linking existing algorithms to stable coloring . We propose a more flexible version of stable coloring ( SC ) called continuous stable coloring ( CSC ) — a strict generalization of SC . CSC states that the similarity of their neighborhoods determines the similarity of two nodes . Based on this notion , we define a series of optimization problems to solve the problem of CSC . We show that various algorithms in NLP like word2vec , AWE , BERT ( Devlin et al. , 2019 ) , images processing like Visual Transformer ( Dosovitskiy et al. , 2020 ) , graphs like Node2Vec ( Grover & Leskovec , 2016 ) , etc . are essentially solving different instantiations of this common optimization problem . Current research already establishes the link between the 1-WL algorithm ( Grohe , 2020 ; Morris et al. , 2021 ; Shervashidze et al. , 2011 ; Morris et al. , 2017 ) and GNN architectures , which has sparked a new line of research in improving GNN architectures ( Hamilton et al. , 2017 ; Xu et al. , 2019b ; Maron et al. , 2019 ; Morris et al. , 2020a ; b ) . We hope that the new link we establish between stable coloring and unsupervised embedding algorithms will also stimulate new exciting research in embeddings for other structured domains of NLP and Vision . 2 BACKGROUND . In this section , we define a stable colored graph , provide an outline for 1-WL ( Weisfeiler-Lehman ) graph isomorphism test and General Aggregate and Update ( GAU ) for Graph Neural Networks ( GNNs ) . We also discuss how domains of NLP and images can be seen as graphs . 2.1 STABLE COLORING . Let a coloring function C be an overloaded function defined on the vertices as well on set of vertices of G = ( V , E ) , i.e. , C : V → N and C : 2V → { { N } } where N is a set of natural numbers representing colors , and { { . } } is a multiset with overloading defined as C ( A ) = { { C ( v ) |v ∈ A ⊂ V } } . We denote neighborhood of a node u ∈ V as N ( v ) = { u| ( v , u ) ∈ E } . Definition 2.1 ( Stable Coloring ) . An undirected graph G = ( V , E ) is stable colored w.r.t coloring function C if it holds that C ( u ) = C ( v ) if and only if C ( N ( u ) ) = C ( N ( v ) ) . We can extend the above definition to directed graphs and graphs with labels . A directed graph G = ( V , E ) is stable colored w.r.t a coloring function C : V → N if it holds that C ( u ) = C ( v ) if and only if C ( Nin ( u ) ) = C ( Nin ( v ) ) and C ( Nout ( u ) ) = C ( Nout ( v ) ) where Nin ( u ) = { w| ( w , u ) ∈ E } and Nout ( u ) = { w| ( u , w ) ∈ E } . Likewise , an edge-labelled undirected graph G = ( V , E ) is stable colored w.r.t a coloring function C : V → N if it holds that C ( u ) = C ( v ) if and only if ∀l , C ( Nl ( u ) ) = C ( Nl ( v ) ) where Nl ( u ) = { w| ( u , w ) ∈ El } where El ⊂ E for an edge label l. We also define a weak-stable coloring for a graph G = ( V , E ) . An undirected graph G = ( V , E ) is weak-stable colored w.r.t coloring function C if it holds that C ( u ) = C ( v ) if C ( N ( u ) ) = C ( N ( v ) ) where N ( u ) = { w| ( u , w ) ∈ E } . 2.2 1-WL ALGORITHM AND GENERAL AGGREGATE AND UPDATE FRAMEWORK ( GAU ) . 1-WL algorithm : 1-WL is an iterative algorithm to achieve a stable coloring C for G = ( V , E ) . Let Ci denote the coloring at iteration i . It starts with a coloring scheme C0 such that C0 ( v ) is same for all v ∈ V . In each iteration , it assigns a different color to u and v if Ci ( N ( u ) ) 6= Ci ( N ( u ) ) until a stable coloring C is reached . General Aggregate and Update Framework ( GAU ) For a general multi-layer GNN , the General Aggregate and Update framework to compute node/vertex embeddings ( corresponding to colors ) of G = ( V , E ) is given iteratively by : Ek ( u ) = f ( k ) update ( Ek−1 ( u ) , f ( k ) agg ( Ek−1 ( v ) : v ∈ { { N ( u ) } } ) ) , ( 1 ) where f ( k ) agg ( . ) and f ( k ) update ( . ) map vertex multiset embeddings to a metric space such as R d. In GraphSage ( Hamilton et al. , 2017 ) , authors showed that the iterative procedure in 1-WL algorithm is analogous to General Aggregate and Update procedure in GNNs . This connection has lead to the research direction where 1-WL is being used as a standard to which GNN architectures are being compared to ( Xu et al. , 2019b ; Maron et al. , 2019 ; Morris et al. , 2020b ) . Xu et al . ( 2019b ) prove that GAU is as powerful as 1-WL if the functions fupdate and fagg are injective . Various domains such as NLP , Images and Graphs can be viewed as a graph on their elementary tokens . How we construct graphs is explained in section 6.1 and examples are given in appendix B . 3 RELATED WORK . Word embeddings have been popular in NLP since decades ( Deerwester et al. , 1989 ; Morin & Bengio , 2005 ; Mikolov et al. , 2013b ; Bojanowski et al. , 2017 ) . A lot of work has been done to understand the mathematical underpinning of these models , for instance , relation of embedding models to co-occurrence statistics ( Levy & Goldberg , 2014 ; Hashimoto et al. , 2016 ; Allen et al. , 2019 ) . Study of empirical properties of these embedding models ( e.g . analogies ) has also attracted theoretical research ( Allen & Hospedales , 2019 ; Ethayarajh et al. , 2019 ) . Recently , graph community has also seen a surge in learning node and graph embeddings . The notion of capturing the structural neighborhood around a node inside the node embedding has been the driving principle of these semi-supervised node embedding algorithms like node2vec ( Grover & Leskovec , 2016 ) , and GraphSAGE ( Hamilton et al. , 2017 ) . Hamilton et al . ( 2017 ) pointed out that their GraphSAGE node embedding algorithm mimics the aggregate and update procedure of 1-WL algorithm . Subsequently , these neighborhood informative node embeddings found their applications in constructing graph embeddings , and thereby used for classification of structurally equivalent or isomorphic graphs ( Chen et al. , 2019 ) . Xu et al . ( 2019b ) in their GIN ( Graph Isomorphism Network ) model modified the aggregate procedure of GraphSAGE to construct graph embedders which were provably as powerful as 1-WL algorithm in distinguishing non-isomorphic graphs . This redirected the research into designing more powerful variants of graph embedders like PPGN ( Maron et al. , 2019 ) and k-GNN ( Morris et al. , 2019 ) which were provably as powerful as 3-WL and k-WL test respectively . Grohe ( 2020 ) discusses about these phenomenal works in increasing the expressivity of graph embeddings for supervised graph classification and regression tasks . While this line of research focuses on improving expressivity and generalizability of Graph networks based on its connection to 1-WL , we explore and formalize the unsupervised algorithms under the light of stable coloring / 1-WL algorithm and show that all the current models stem from the common principle that a tokens meaning is derived from its neighbours . The rest of the paper is organised as follows . We begin with a discussion of connections between discrete stable coloring and non-local means algorithm in Image processing and distributional hypothesis in NLP . We then define a continuous version of SC and develop optimization problems to solve for CSC in section 5 . In section 6 we show how current state-of-the art embedding models are solving the CSC problem in disguise .

Summary: In this paper, the authors investigate how the solution of the Weisfeiler-Lehman test of isomorphism (WL) is related to embedding algorithms from structured domains. The authors generalize the WL algorithm to the case where the colors assigned to the nodes correspond to feature vectors, and propose a series of optimization problems to solve this problem. Finally, they show that widely used algorithms in natural language processing, computer vision and graph representation learning such as word2vec, BERT, Visual Transformer and node2vec actually solve different instantiations of this common optimization problem.