Universal Approximation Under Constraints is Possible with Transformers

Keywords : Constrained Universal Approximation , Probabilistic Attention , Transformer Networks , Geometric Deep Learning , Measurable Maximum Theorem , Non-Affine Random Projections . 1 INTRODUCTION . In supervised learning , we select a parameterized model f̂ : Rn ! Rm by optimizing a real-valued loss function1L over training data from an input-output domain X ⇥ Y ✓ Rn ⇥ Rm . A necessary property for a model class to produce asymptotically optimal results , for any continuous loss L , is the universal approximation property . However , often more structure ( beyond vectorial Rm ) is present in a learning problem and this structure must be encoded into the trained model f̂ to obtain meaningful or feasible predictions . This additional structure is typically described by a constraint set K ✓ Rm and the condition f̂ ( X ) ✓ K. For example , in classification K = { y 2 [ 0 , 1 ] m : Pm i=1 yi = 1 } ( Shalev-Shwartz & Ben-David , 2014 ) , in Stackelberg games ( Holters et al. , 2018 ; Jin et al. , 2020 ; Li et al. , 2021 ) K is the set of utility-maximizing actions of an opponent , in integer programming K is the integer lattice Zm ( Conforti et al. , 2014 ) , in financial risk-management K is a set of positions meeting the minimum solvency requirements imposed by international regularity bodies ( Basel Committee on Banking Supervision , 2015 ; 2019 ; McNeil et al. , 2015 ) , in covariance matrix prediction K ✓ Rm⇥m is the set of m⇥m matrices which are symmetric and positive semidefinite ( Bonnabel et al. , 2013 ; Bonnabel & Sepulchre , 2009 ; Baes et al. , 2021 ) , in geometric deep learning K is typically a manifold ( e.g . a pose manifold in computer vision and robotics ( Ding & Fan , 2014 ) or a manifold of distance matrices ( Dokmanic et al. , 2015 ) ) , a graph , or an orbit of a group action ( Bronstein et al. , 2017 ; 2021 ; Kratsios & Bilokopytov , 2020 ) . Therefore , we ask : Is exact constraint satisfaction possible with universal deep learning models ? ⇤Corresponding authors . 1For example , in a regression problem one can set L ( x , y ) = kf ( x ) yk for an unknown function f or in regression problems one sets L ( x , y ) = Pm i=1 [ C ( x ) ] i log ( yi ) for an unknown classifier C. The answer to this question begins by examining the classical universal approximation theorems for deep feedforward networks . If L and K are mildly regular , the universal approximation theorems of Hornik et al . ( 1989 ) ; Cybenko ( 1989 ) ; Pinkus ( 1999 ) ; Gühring et al . ( 2020 ) ; Kidger & Lyons ( 2020 ) ; Park et al . ( 2021 ) guarantee that for any “ good activation function ” and for every tolerance level ✏ > 0 , there is a deep feedforward network with activation function , such that infy2K L ( x , y ) and L ( x , f̂ ( x ) ) are uniformly at most ✏ apart . Written in terms of the optimality set , sup x2X kf̂ ( x ) argmin y2K L ( x , y ) k ✏ , ( 1 ) where the distance of a point y 2 Rm to a set A ✓ Rm is defined by ky Ak , infa2A ky ak . Since argminy2K L ( x , y ) ✓ K , then ( 1 ) only implies that kf̂ ( x ) Kk ✏ and there is no reason to believe that the constraint f̂ ( x ) 2 K is exactly satisfied , for every x 2 X . This kind of approximate constraint satisfaction is not always appropriate . In the following examples constraint violation causes either practical or theoretical concerns : ( i ) In post-financial crisis risk management , international regulatory bodies mandate that any financial actor should maintain solubility proportional to the risk of their investments ( Basel Committee on Banking Supervision , 2015 ; 2019 ) . To prevent future financial crises , any violation of these risk constraints , no matter the size , incurs large and immediate fines . ( ii ) In geometric deep learning , we often need to encode complicated non-vectorial structure present in a dataset , by viewing it as a K valued function ( Fletcher , 2013 ; Bonnabel & Sepulchre , 2009 ; Baes et al. , 2021 ) . However , if K is non-convex then Motzkin ( 1935 ) confirms that there is no unique way to map predictions f̂ ( x ) 62 K to a closest point in K. Thus , we are faced with the dilemma : either make an ad-hoc choice of a k in K with k ⇡ f̂ ( x ) ( ex . : an arbitrary choice scheme when K = Zm ) or have meaningless predictions ( ex : non-integer values to integer programs , or symmetry breaking ( Weinberg , 1976 ) 2 ) . Constrained learning was recognized as an effective framework for fairness and robustness by Chamon & Ribeiro ( 2020 ) who study empirical risk minimization under constraints . Many emerging topics in machine learning lead to constrained learning formulations . A case in point is modelbased domain generalization ( Robey et al. , 2021 ) . Despite the importance of ( deep ) learning with constraints , there are no related approximation-theoretic results to the best of our knowledge . In this paper , we bridge this theoretical gap by showing that universal approximation with exact constraint satisfaction is always possible for deep ( probabilistic ) transformer networks with a single attention mechanism as output layer . Our contribution is three-fold : 1 . We derive the first universal approximation theorem with exact constraint satisfaction ; 2 . Our transformer network ’ s encoder and decoder adapt to the dimension of the constraint set and thus beat the curse of dimensionality for low-dimensional constraints ; 3 . Our models leverage a probabilistic attention mechanism that can encode non-convex con- straints . This probabilistic approach is key to bypass the topological obstructions to nonEuclidean universal approximation ( Kratsios & Papon , 2021 ) . Our analysis provides perspective on the empirical success of attention and adds to the recent line of work on approximation theory for transformer networks , ( Yun et al. , 2020a ; b ) , which roughly considers the unconstrained case ( with K in ( 1 ) replaced by Rm ) in the special case of L ( x , y ) = kf ( x ) yk for a suitable target function f : Rn ! Rm . Our probabilistic perspective on transformer networks fits with the representations of Vuckovic et al . ( 2021 ) and of Kratsios ( 2021 ) . Our results can be regarded as an approximation-theoretic counterpart to the constrained statistical learning theory of Chamon & Ribeiro ( 2020 ) . Further , they put forward a perspective on randomness in neural networks that is complementary to the work of Louart et al . ( 2018 ) ; Gonon et al . ( 2020a ; b ) . We look at the same problem focusing on constraint satisfaction instead of training efficiency . Finally , our proof methods are novel , and build on contemporary tools from metric geometry ( Ambrosio & Puglisi , 2020 ; Bruè et al. , 2021 ) . 2As discussed in Rosset et al . ( 2021 ) this is problematic since respecting symmetries can often massively reduce the computational burden of a learning task . 1.1 THE PROBABILISTIC ATTENTION MECHANISM . We now give a high-level explanation of our results ; the detailed formulations are in Section 2 . Introduced in ( Bahdanau et al. , 2015 ) and later used to define the transformer architecture ( Vaswani et al. , 2017 ) , in the NLP context , attention maps a matrix of queries Q , a matrix of keys K , and a matrix of values V to the quantity Softmax ( QK > ) V , where the softmax function ( defined below ) is applied row-wise to QK > . Just as the authors of ( Petersen & Voigtlaender , 2020 ; Zhou , 2020 ) focus on the simplified versions of practically implementable ConvNets in the study of approximation theory of deep ConvNets ( e.g . omitting pooling layers ) , we find it sufficient to study the following simplified attention mechanism to obtain universal approximation results : Attention ( w , Y ) , SoftmaxN ( w ) > Y = NX n=1 [ SoftmaxN ( w ) n ] Yn , ( 2 ) where w 2 RN , SoftmaxN : RN 3 w 7 ! ( e wkPN j=1 e wj ) N k=1 , and Y is an N ⇥m matrix . The attention mechanism ( 2 ) can be interpreted as “ paying attention ” to a set of particles Y1 , . . . , YN 2 Rm defined by Y ’ s rows . This simplified form of attention is sufficient to demonstrate that transformer networks can approximate a function while respecting a constraint set , K , whether convex or non-convex . Informal Theorem 1.1 ( Deep Maximum Theorem for Transformers ) . If K is convex and the quantities defining ( 1 ) are regular then , for any ✏ 2 ( 0 , 1 ] , there is a feedforward network f̂ , an X✏ ⇢ Rn of probability 1-✏ , and a matrix Y such that the transformer Attention ( f̂ ( x ) , Y ) satisfies : ( i ) Exact Constraint Satisfaction : For each x 2 Rn , Attention ( f̂ ( x ) , Y ) 2 K , ( ii ) Universal Approximation : supx2X✏ kAttention ( f̂ ( x ) , Y ) argmin y ? 2K L ( x , y ? ) k ✏ Informal Theorem 1.1 guarantees that simple transformer networks can minimize any loss function while exactly satisfying the set of convex constraints . As illustrated by Figure 1 and Figure 2 , K ’ s convexity is critical here , since without it the transformer ’ s prediction may fail to lie in K. This is because any transformer network ’ s output is a convex combinations of the particles Y1 , Y2 , Y3 ; thus , any transformer network ’ s predictions must belong to these particles ’ convex hull . In Figures 1 and 2 , Y ’ s columns , i.e . the particles Y1 , Y2 , and Y3 , are each illustrated by a • at the constraint set ( K ) vertices . The bubble around each each Yi illustrates the predicted probability , for a given input , that f ( x ) is nearest to that Yi . The ⇥ is the transformer ’ s prediction which is , by construction , a convex combination of the Yi weighted by the aforementioned probabilities and therefore they lie in the K if it is convex ( Figure 1 ) but not if K is non-convex ( Figure 2 ) . Naturally , we arrive at the question : How can ( i ) and ( ii ) simultaneously hold when K is non-convex ? Returning to Vaswani et al . ( 2017 ) and using the introduced terminology , we note that the role of the SoftmaxN layer is to rank the importance of the particles { Yn } Nn=1 when optimizing L , at any given input : the weights [ SoftmaxN ( w ) ] n in ( 2 ) can be interpreted as charging their respective point masses { Yn } Nn=1 with probabilities of being optimal for L ( relative to the other particles ) 3 . This suggests the following probabilistic reinterpretation of attention ( which we denote by p-attention ) : P-attention ( w , Y ) , NX n=1 [ SoftmaxN ( w ) ] n Yn . ( 3 ) 3Following Villani ( 2009 ) , Yn is the Borel probability measure on R m assigning full probability to any Borel subset of Rm containing the particle Yn and 0 otherwise . Crudely put , P-attention ( · , Y ) “ pays relative attention to the particles ” Y1 , . . . , Yn 2 Rm . A simple computation shows that the mean prediction of our probabilistic attention mechanism , exactly implements “ classical ” Attention of Vaswani et al . ( 2017 ) , as defined in ( 2 ) , Attention ( w , Y ) = EX⇠P-attention ( w , Y ) [ X ] , ( 4 ) where EX⇠P-attention ( w , Y ) [ X ] denotes the ( vector-valued ) expectation of a random-vector X distributed according to P-attention ( w , Y ) . Hence , ( 3 ) is no less general than ( 2 ) . The advantage of ( 3 ) is that , if each particle Yn belongs to K ( even if K is non-convex ) then , any sample drawn from the probability measure P-attention ( w , Y ) necessarily belongs to K .

Summary: This paper presents a family of constrained universal approximation results for probabilistic transformers.  The authors provide significant theoretical contributions for both convex and non-convex constraint sets.  In my opinion, this represents a significant advance in our understanding of universality in ML.