U NIVERSAL A PPROXIMATION U NDER C ONSTRAINTS IS P OSSIBLE WITH T RANSFORMERS
A BSTRACT
Many practical problems need the output of a machine learning model to satisfy a set of constraints, K .
There are, however, no known guarantees that classical neural networks can exactly encode constraints while simultaneously achieving universality.
We provide a quantitative constrained universal approximation theorem which guarantees that for any convex or non-convex compact set K and any continuous function f : R n !
K , there is a probabilistic transformer ˆ F whose randomized outputs all lie in K and whose expected output uniformly approximates f .
Our second main result is a “deep neural version” of Berge (1963)’s
Maximum Theorem.
The result guarantees that given an objective function L , a constraint set K , and a family of soft constraint sets, there is a probabilistic transformer ˆ F that approximately minimizes L and whose outputs belong to K ; moreover, ˆ F approximately satisﬁes the soft constraints.
Our results imply the ﬁrst universal approximation theorem for classical transformers with exact convex constraint satisfaction, and a chart-free universal approximation theorem for Riemannian manifold-valued functions subject to geodesically-convex constraints.
I NTRODUCTION
In supervised learning we select a parameterized model ˆ f : R n !
R m by optimizing a real-valued loss function 1 L over training data from an input-output domain X ⇥ Y ✓
R n ⇥ R m .
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 R m ) 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 ✓ R m and the condition ˆ f ( X ) ✓ K .
For example, in classiﬁcation K = { y 2 [0 , 1] m :
P mi =
y
i
= 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 Z m (Conforti et al., 2014), in ﬁnancial 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 ✓ R m ⇥ m is the set of m ⇥ m matrices which are symmetric and positive semideﬁnite (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?
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.
; Cybenko (1989); Pinkus (1999); G¨uhring et al.
; 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 inf y 2 K L ( x, y ) and
L ( x, ˆ f ( x )) are uniformly at most ✏ apart.
Written in terms of the optimality set,
[Equation]
where the distance of a point y 2 R m to a set A ✓ R m is deﬁned by k y
A k , inf a 2 A k y  a k .
Since argmin y 2 K L ( x, y ) ✓ K , then (1) only implies that k ˆ f ( x )  K k  ✏ and there is no reason to believe that the constraint ˆ f ( x ) 2 K is exactly satisﬁed, 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-ﬁnancial crisis risk management, international regulatory bodies mandate that any ﬁnancial actor should maintain solubility proportional to the risk of their investments (Basel Committee on Banking Supervision, 2015; 2019).
To prevent future ﬁnancial crises, any violation of these risk constraints, no matter the size, incurs large and immediate ﬁnes.
(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) conﬁrms 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 = Z m ) 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:
We derive the ﬁrst 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 constraints.
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 R m ) in the special case of L ( x, y ) =
k f ( x )  y k for a suitable target function f : R n !
R m .
Our probabilistic perspective on transformer networks ﬁts 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.
Gonon et al.
(2020a;b).
We look at the same problem focusing on constraint satisfaction instead of training efﬁciency.
Finally, our proof methods are novel, and build on contemporary tools from metric geometry (Ambrosio & Puglisi, 2020; Bru`e et al., 2021).
1. T HE P ROBABILISTIC A TTENTION M ECHANISM
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 deﬁne 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 (deﬁned below) is applied row-wise to QK > .
Just as the authors of (Petersen & Voigtlaender, 2020; Zhou, 2020) focus on the simpliﬁed versions of practically implementable ConvNets in the study of approximation theory of deep ConvNets (e.g. omitting pooling layers), we ﬁnd it sufﬁcient to study the following simpliﬁed attention mechanism to obtain universal approximation results:
[Equation]
where w 2 R N , Softmax N : R N 3 w 7!
( e wk
P
Nj =1 e wj )
Nk =1 , and Y is an N ⇥ m matrix.
The attention mechanism (2) can be interpreted as “paying attention” to a set of particles Y 1 , . .
Y N 2 R m deﬁned by Y ’s rows.
This simpliﬁed form of attention is sufﬁcient 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 deﬁning (1) are regular then, for any ✏ 2 (0 , 1] , there is a feedforward network ˆ f , an X ✏ ⇢ R n of probability 1 - ✏ , and a matrix Y
such that the transformer Attention( ˆ f ( x ) , Y ) satisﬁes:
(i) Exact Constraint Satisfaction: For each x 2 R n , Attention( ˆ f ( x
[Equation]
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 Y 1 , Y 2 , Y 3 ; thus, any transformer network’s predictions must belong to these particles’ convex hull.
[Figure]
In Figures 1 and 2, Y ’s columns, i.e. the particles Y 1 , Y 2 , and Y 3 , are each illustrated by a • at the constraint set ( K ) vertices.
The bubble around each each Y i illustrates the predicted probability, for a given input, that f ( x ) is nearest to that Y i .
The ⇥ is the transformer’s prediction which is, by construction, a convex combination of the Y i 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.
and using the introduced terminology, we note that the role of the Softmax N layer is to rank the importance of the particles { Y n } Nn =1 when optimizing L , at any given input: the weights [Softmax N ( w )] n in (2) can be interpreted as charging their respective point masses {  Y n } 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):
[Equation]
Crudely put, P-attention( · , Y ) “pays relative attention to the particles” Y 1 , . .
, Y n 2 R m .
A simple computation shows that the mean prediction of our probabilistic attention mechanism, exactly implements “classical” Attention of Vaswani et al.
(2017), as deﬁned in (2),
Attention( w, Y ) = E X ⇠ P-attention( w,Y )
[ X ] , (4) where E X ⇠ 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 Y n 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 .
1. Q UALITATIVE R ESULTS : D EEP M AXIMUM T HEOREM
Probabilistic attention (3) yields the following non-convex generalization of Informal Theorem 1.1.
The result is a qualitative universal approximation theorem as well as a deep neural version of the
Maximum Theorem 4 (Berge, 1963), which states that under mild regularity conditions, given any well-behaved family of input dependent “soft constraint sets” { C x } x 2 R n compatible with K , there is a measurable function mapping each x 2 R n to a minimizer of L ( x, y ) on K \ C x .
We use W 1 to denote the Wasserstein-1 distance between probability measures on K .
The results also give the ﬂexibility to the user to enforce an input-dependent family of “soft constraints” { C x } x 2 R n which only need to hold approximately; deﬁnitions are provided in Section 1.4.
Informal Theorem 1.2 (Deep Maximum Theorem: Non-Convex Case) .
If the quantities deﬁningare regular, K is a compact set of “exact constraints”, and { C x } x 2 R n a set of “soft constraints”, then, for any approximation quality 0 < ✏  1 , there is a deep feedforward network ˆ f and a matrix
Y satisfying:
(i) Exact Constraint Satisfaction: For each x 2 R n , P-attention( ˆ f ( x ) , Y ) is supported in K , (ii) Universal Approximation: P ( W 1 (P-attention( ˆ f ( x ) , Y ) , argmin y ?
[Equation]
where for a probability measure P on R m and a B ✓ R m we deﬁne W 1 ( P , B ) , inf b 2 B W 1 ( P ,  b ) .
Example 1.3 (Reduction to Classical Point-to-Set Distance) .
In particular, when P is a point-mass
P =  y for some y 2 R m , then one recovers the familiar Euclidean distance to the set B via:
[Equation]
where the ﬁrst and second equality follows from (Villani, 2009, (5) - page 99), and the last equality is the deﬁnition of k y  B k (as in (Aubin & Frankowska, 2009, Deﬁnition 1.1.1)).
Another important class of non-convex constraints arising from geometric deep learning where K is a non-Euclidean ball in a Riemannian submanifold of R m .
In this broad case, we may extract mean predictions from P-attention( ˆ f, Y ) , by applying the Fr´echet mean introduced in Fr´echet (1948).
Such “geometric means” are well-understood theoretically (Bhattacharya & Patrangenaru, 2003) and easily handled numerically Miolane et al.
Lou et al.
1. Q UANTITATIVE R ESULTS : C ONSTRAINED U NIVERSAL A PPROXIMATION T HEOREM
In its current form, the objective function L is too general to derive quantitative approximation rates 5 .
Nevertheless, as with most universal approximation theorems (Hornik et al., 1989; Pinkus, 1999; Kidger & Lyons, 2020), if each soft constraint C x is set to R m and L quantiﬁes the uniform distance to an unknown continuous function f : R n !
K in the Euclidean sense,
[Equation]
then, Informal Theorem 1.2 reduces to a (qualitative) universal approximation for transformer networks with exact constraint satisfaction.
In fact, this additional structure is enough for us to derive quantitative versions of the aforementioned results.
We permit ourselves the general situation, where
K is contained in an unknown d -dimensional submanifold (where d 2 ⇥ ( m 1 s ) for some s > 0 ).
Our approximation rates scale favourably in the ratio s ⇡ log( m ) log( d ) ; i.e., we avoid the curse of dimensionality for low-dimensional constraint sets.
This additional structure translates into the familiar encoder-decoder structure deployed in most transformer network implementations.
Figure 3 illustrates the encoder network ˆ E : R n !
R m , whose role is to perform a (classical) unconstrained approximation of the target function, f .
Since ˆ E is a classical feedforward network then its approximation of the target function can be arbitrarily close to the constraint set K
but it need not lie in it.
The next step is to “map the encoder network ’s output onto K with low distortion.”
The role of the decoder network ˆ D is to correct any constraint violation made by encoder network by “projecting them back on to K ”.
However, such a projection does not exist if K is not convex since there must be more than one closest point in K to some y 2 R m (Motzkin, 1935).
Nevertheless, if the “projection” were capable of mapping any y 2 R m to multiple points on K , ranked by their proximity to y , then there would be no trouble.
The decoder network accomplishes precisely this, as illustrated in Figure 4, where the bubbles illustrate the probability of any particle in K being closest to y , illustrated by the size of the bubbles in Figure 4.
Mathematically, 6 ˆ D : R m !
P 1 ( K ) approximates a (non-afﬁne) random projection , in the sense of Ohta (2009);
Ambrosio & Puglisi (2020);
Bru`e et al.
(2021); i.e.: a 1 -Lipschitz map ⇧ : R m !
P 1 ( K ) satisfying the random projection property : for all y 2 K
[Equation]
Thus, ⇧ ’s random projection property means that it ﬁxes any output already satisfying the constraint
K , and its Lipschitz regularity implies that it is stable.
Thus, sampling from ⇧ ( y 1 ) is similar to sampling from ⇧ ( y 2 ) whenever the points y 1 , y 2 2 R m are near to one another.
Remark 1.4.
Random projections are closely tied to the (random) partitions of unity of Lee & Naor (2005) (see (Ambrosio & Puglisi, 2020, Theorem 2.8)).
These random projections generalize the random projections of Johnson & Lindenstrauss (1984), beyond the case where K is afﬁne.
Remark 1.5.
The special case of random projections onto afﬁne spaces has recently been used when constructing universal neural models (Cuchiero et al., 2021; Puthawala et al., 2020).
We record the complexity of both the decoder and encoder networks constructed in our quantitative results in Table 1.
Here A, B, C, D  0 are constants independent of ✏ and k , where k 2 N + is the number of continuous derivatives which f admits (when viewed as a function into R m ).
From
Network
[Equation]
[Table]
[Equation]
[Table]
Table 1, we see that if m 1 s ⌧ m then, s > 0 is large; hence, ✏ ms , (1  4 ✏  1 ) 2 ms , and N ms are small.
1. N OTATION AND B ACKGROUND
Optimal Transportd Given any non-empty subset K ✓ R m , the set of all Borel probability measures P on K with a ﬁnite mean; i.e.: E X ⇠ P [ k X k ] < 1 , is denoted by P 1 ( K ) .
Wasserstein distance W 1 is deﬁned for any P , Q 2 P 1 ( K ) by the minimal energy needed to transport all mass from P to Q .
Following Villani (2009), W 1 ( P , Q ) is deﬁned by:
[Equation]
where the inﬁmum is taken over all Borel probability measures ⇡ on K 2 with marginals P and Q .
The metric space ( P 1 ( K ) , W 1 ) is named the Wasserstein space over K ; we abbreviate it by P 1 ( K ) .
Smooth Function Spaces The set of real-valued continuous functions on R n is denoted by C ( R n ) .
Let k 2 N + and X ✓ [0 , 1] n be non-empty.
The set of functions f : X ! K for which there is a k -times continuously differentiable f : R n !
R m extending f ; i.e.: f | X = f , is denoted by
C ktr ( X , K ) .
Our interest in C ktr
( X , K ) does not stem from the fact that it contains all smooth functions mapping [0 , 1] n to K , but rather that it allows us to speak about the uniform approximation of discontinuous K -valued functions on regions in [0 , 1] n where they are “regular”.
This is noteworthy for pathological constraint sets, such as integer constraints 7 .
For details on C ktr ( X , K ) , see (Brudnyi & Brudnyi, 2012a;b) and the extension theorems of Whitney (1934); Fefferman (2005).
Neural Networks It has recently been observed that deep feedforward networks with multiple activation functions, or more generally parametric families of activation functions, achieved significantly more efﬁcient approximation rates than classical feedforward networks with a single activation function (Yarotsky & Zhevnerchuk, 2020; Yarotsky, 2021;
Shen et al., 2021a;b).
Practically deployed examples of parametric activation functions are the PReLU activation function of He et al.
(2015), the Sigmoid-weighted Linear Unit (SiLU) of Elfwing et al.
(2018), and the Swish activation function of Ramachandran et al.
We also observe a similar phenomenon, and therefore our quantitative results consider deep feedforward networks whose activation functions belongs to a 1parameter family  ? , {  t } t 2 [0 , 1] ✓ C ( R ) .
The set of all such networks is denoted by NN  ? n,N and it includes all ˆ f : R n !
R N with iterative representation:
[Equation]
where x 2 R n , j = 1 , . .
, J  1 , each A ( j ) is a d j ⇥ d j +1 -matrix, each b ( j ) 2 R d j +1 , d
J +1 =
N , d 1 = 0 , t 1 , 1 , . . .
, t J,N J 2 [0 , 1] , for each j .
The integer J is ˆ f ’s depth and max j =1 ,...,
[Equation]
is
[Equation]
width .
Example 1.6 (Networks with Untrainable Nonlinearity) .
Denote  ,  0 .
The subset of classical feedforward networks consisting of all ˆ f 2 NN  ? n,N with each  t ij =  in (5) is denoted NN  n,N .
It is approximation theoretically advantageous to generalize the proposed deﬁnition of probabilistic attention in the introduction (3) by replacing Y with a 3 -dimensional array (elementary 3 -tensor).
Deﬁnition 1.7 (Probabilistic Attention) .
Let N, Q, m 2 N + , and Y be an N ⇥ Q ⇥ m -array with Y n,q 2 K for n = 1 , . . .
, N , q = 1 , . .
, Q .
Probabilistic attention is the function:
[Equation]
= Softmax N ( w ) n  Y n,q 2 P 1 ( K ) . If Y is an N ⇥ m -matrix, as in (3), then we identify Y as the N ⇥ m ⇥ 1 -array in the obvious manner.
Set-Valued Analysis: A family of non-empty subsets { C x } x 2 R n of K
is said to be a weakly measurable correspondence , denoted C : R n ◆ R m , if for every open subset 8 U ✓ K , { x 2 R n : C x \ U 6 = ; } is a non-empty Borel subset of R n (Aliprantis & Border, 2006, pages 557, 592).
M AIN R ESULTS
We now present our main results in detail.
All proofs are relegated to the paper’s appendix.
2. Q UALITATIVE A PPROXIMATION : D EEP M AXIMUM T HEOREM
Our main qualitative result is the following deep neural version of Berge (1963)’s Maximum Theorem where, the measurable selector is approximately implemented by a probabilistic transformer Z
Since K is equipped with its subspace topology, then an open subset U of K is any subset
R m form U
= U 1 \ K where U 1 is an open subset of R
m (see (Munkres, Lemma 16.1) for further details). network. We ﬁrst present the general qualitative result which gives a concrete description of a measurable selector of (1), with high-probability, which has the key property that all its predictions satisfy the required constraints deﬁned by K .
Assumption 2.1 (Kidger & Lyons (2020)) .
: R !
R is continuous,  is differentiable at some x 0 2 R , and its derivative satisﬁes  0 ( x 0 ) 6 = 0 .
Theorem 2.2 (Deep Maximum Theorem) .
Let  satisfy Assumption 2.1.
Let K ✓ R n be a nonempty compact set, C : R n ◆ R m be a weakly-measurable correspondence with closed values such that C x \ K 6 = ; for each x 2 R n , L 2 C ( R m ) , and P be a Borel probability measure on R n .
For each 0 < ✏  1 , there is an N 2 N + , an ˆ f 2 NN  n,N of width at most 2 + n + N , and an N ⇥ m -matrix Y such that:
[Equation]
satisﬁes the following:
(i) Exact Constrain Satisfaction: [ x 2 R n supp( ˆ F ( x )) ✓ K , (ii) Probably Approximately Optimality: There is a compact X ✏ ✓ R n satisfying: (a) max x 2 X ✏ W 1 ( ˆ F ( x ) , argmin y 2 C x \ K L ( x, y ))
 ✏ , (b) 1  P ( X ✏ )  ✏ .
Theorem 2.2 implies that for any random ﬁeld ( Y x ) x 2 R n on R m (i.e. a family of R m -valued random vectors indexed by R n ) with Y x ⇠ ˆ F ( x ) : 1.
samples drawn from Y x are in K (by (i)) and 2.
samples drawn from each Y x are near to the optimality set argmin y 2 C x \ K L ( x, y ) (by (ii)).
Corollary 2.3 ( ˆ F ’s Mean Prediction) .
Assume the setting of Theorem 2.2.
Let { Y x } x 2 R n be a
K -valued random ﬁeld with Y x
⇠ ˆ F ( x ) for each x 2 R n then, 1  P ( X ✏ )  ✏ and
[Equation]
Appendix 8 contains additional consequences of the Deep Maximum Theorem, such as the special case of classical transformers when K is convex.
Next, we complement our qualitative results by their quantitative analogues, within the context of universal approximation under constraints .
2. Q UANTITATIVE A PPROXIMATION : C ONSTRAINED U NIVERSAL A PPROXIMATION
In order to derive a quantitative constrained universal approximation theorem, we require the loss function to be tied to the Euclidean norm in the following manner.
Assumption 2.4 (Norm-Controllable Loss) .
There is a continuous f : R n !
R m with f ( R n ) ✓ K and a continuous l : [0 , 1 ) !
[0 , 1 ) with l (0) = 0 , satisfying: L ( x, y )  l ( k f ( x )  y k ) .
Just as with transformer networks, our “constrained universal approximation theorem” approximates a suitably regular function f : R n !
K ✓ R m while exactly respecting the constraints K by implementing an encoder-decoder network architecture.
Thus, our model is a composition of an encoder network ˆ E : R n !
R d whose role is to approximate f in a classical “unconstrained fashion” and a decoder network (with probabilistic attention layers at its output) ˆ D : R d !
P 1 ( K ) whose role is to enforce the constraints K while preserving the approximation performed by ˆ E , where d n m .
To take advantage of the encoder-decoder framework present in most transformer networks, we formalize what is often called a “latent low-dimensional manifold” hypothesis.
Brieﬂy, this means that, the hard constraints in set K are contained in a “low dimensional” subspace.
Assumption 2.5 (Low-Dimensional Manifold) .
There is an 0 < s and a smooth bijection  from
R n to itself with smooth inverse 9 , such that  ( K ) ✓ R d ; where 2  d and d 2 ⇥ ( m 1 s ) .
Assumption 2.5 does not postulate that K is itself a single-chart low-dimensional manifold, or even a manifold.
Rather, K need only be contained in a low-dimensional manifold.
For the fast rates we use activation functions generalizing the swish function (Ramachandran et al., 2018) as follows.
Assumption 2.6 (Swish-Like Activation Function) .
The map  : [0 , 1] ⇥ R 3 ( ↵ , t ) 7!
↵ ( t ) 2 R is continuous;  0 is non-afﬁne and piecewise-linear, and  1 is smooth 10 and non-polynomial.
Theorem 2.7 (Constrained Universal Approximation) .
Let k 2 N + and X ✓ [0 , 1] n be non-empty.
Suppose that  satisﬁes 2.6, L satisﬁes Assumption 2.4, K ✓ R n is non-empty, compact and satisﬁes Assumption 2.5.
For any f 2 C ktr ( X , K ) , every constraining quality ✏ K > 0 , and every approximation error ✏ f > 0 , there exist N, Q 2 N + , an encoder ˆ E 2 NN  n,d , and a decoder:
[Equation]
where ˆ D 2 NN  d,N and Y is an N ⇥ Q ⇥ m -array with Y 1 , 1 , . .
, Y N,Q 2 K such that:
(i) Exact Constrain Satisfaction:
For each x 2 R n : supp( ˆ D  ˆ E ( x )) ✓ K, (ii) Universal Approximation: The estimate holds 11 :
[Equation]
where, 0 < k is an absolute constant independent of n , m , d , f , and of ✏ and Lip(   1 ) denotes the Lipschitz constant of   1 on the compact set { z 2 R d : k z   ( K ) k
 ✏ f } .
Furthermore, the “complexities” of ˆ D and ˆ E are recorded in Table 12 1 for ✏ 2 = ✏ k = ✏ f .
In practice, we can only sample from each measure ˆ D  ˆ E ( x ) .
In this case, we may ask how the typical sample drawn from a random-vector Y x distributed according to our learned measure ˆ D
ˆ E ( x ) performs when minimizing L ( x, y ) .
The next result relates the estimates in Theorem 2.7 (ii) to the typical (in Y x ) worst-case (in x ) gap between a sample from Y x and f ( x ) , as quantiﬁed by L ( x, · ) .
Corollary 2.8 (Average Worst-Case Loss) .
Assume the setting of Theorem 2.7 and suppose that the “modulus” l in Assumption 2.4 is strictly increasing and concave.
Let ˆ D and ˆ E be as in Theorem 2.
and let { Y x } x 2 X be an R m -valued random ﬁeld with Y x ⇠ ˆ D
ˆ E ( x ) , for each x 2 R n .
Then:
[Equation]
Corollary 2.8 quantiﬁes the expected performance of a sample from our probabilistic transformer model, as expressed by L , whereas Theorem 2.7 (ii) quantiﬁes the difference from the transformer’s prediction to the optimal prediction value.
Next, we consider implications of our main results.
2. A PPLICATIONS
We apply our theory to obtain a universal approximation theorem for classical transformer networks with exact convex constraint satisfaction and to derive a version of the non-Euclidean universal approximation theorems of Kratsios & Bilokopytov (2020); Kratsios & Papon (2021) for Riemannianmanifold valued functions which does not need explicit charts.
As with most quantitative (uniform) universal approximation theorems (G¨uhring et al., 2020; Kidger & Lyons, 2020; Shen et al., 2021a), we henceforth consider L ( x, y ) =
k f ( x )  y k .
We also ﬁx f 2 C ktr
([0 , 1] n , K ) .
2.3. T RANSFORMERS ARE C ONVEX -C ONSTRAINED U NIVERSAL A PPROXIMATORS
We return to the familiar transformer networks of Vaswani et al.
The next result shows that transformer networks can balance universal approximation and exact convex constraint satisfaction.
This is because when K is convex, then the mean of the random ﬁeld { Y x } x 2 R n of Corollary 2. must belong to K .
Consequently, the identity (4) implies that Attention( ˆ D  ˆ E ( · ) , Y ) ⇡ f .
Corollary 2.9 (Constrained Universal Approximation: Convex Constraints) .
Consider the setting and notation of Corollary 2.8.
Suppose that K is convex and let L ( x, y ) =
k f ( x )  y k .
Then:
[Equation]
(i) Exact Constraint Satisfaction: E Y x
⇠ ˆ D  ˆ E ( x ) [ Y x ] 2 K , for each x 2 R n , (ii) Universal Approximation: sup [0 , 1] n k f ( x )
E Y x
⇠ ˆ D  ˆ E ( x )
[ Y x ] k < ✏ K + kd ✏ f .
The “complexities” of the networks ˆ D and ˆ E are recorded in Table 13 1 for ✏ 2 = ✏ k = ✏ f .
[Equation]
2.3. C HART -F REE R
M U A
We explore how additional non-convex structure of the constraint set K can be encoded by the probabilistic transformer networks of Theorems 2.2 and 2.7 and be used to build new types of (deterministic) transformer networks.
These results highlight that the standard transformer networks of (8) are specialized for convex constraints and that by instead using an intrinsic variant of expectation, we build can new types of “geometric transformer networks” customized to K ’s geometry.
This section makes use of Riemannian geometry; for an overview see Jost (2017)
Let ( M, g ) be a connected d -dimensional Riemannian submanifold of R m with distance function by d g .
We only require the following mild assumption introduced in Afsari (2011).
We recall that the injectivity radius at y 0 , denoted by inf g ( y 0 ) , (see (Jost, 2017, Deﬁnition 1.4.6)) is the minimum length of a geodesic (or minimal length curve) in M with starting point y 0 .
We also recall that the sectional curvature (see (Jost, 2017, Deﬁnition 4.3.2) for a formal statement) quantiﬁes the curvature of ( M, g ) as compared the geometry of its ﬂat counterpart R d .
We focus on a broad class of nonconvex constrains, namely geodesically convex constraints , which generalize convex constraint and have received recent attention in the optimization literature (Zhang & Sra, 2016; Liu et al., 2017).
Assumption 2.10 (Geodesically Convex Constraints) .
The Riemannian manifold ( M, g ) is connected, it is complete as a metric space, and all its sectional curvatures of ( M, g ) are all bounded above by a constant C  0 .
The non-empty constrain set K satisﬁes:
K is contained in the geodesic ball B ( y 0 , ⇢ ) , { y 2 M : d g ( y 0 , y ) < ⇢ } for some point y 0 2 M and some radius ⇢ satisfying 14 : 0 < ⇢ < 2  1 min { inj g ( y 0 ) , ⇡ p C } , 2.
For each y 0 , y 1 2 K there exists a unique geodesic  : [0 , 1] !
K joining y 0 to y 1 .
Our latent probabilistic representation grants us the ﬂexibility of replacing the usual “extrinsic mean” used in (8) to extract deterministic predictions from our probabilistic transformer networks via an additional Fr´echet mean layer at their readout.
This intrinsic notion of a mean, was introduced independently in Fr´echet (1948) and in Karcher (1977), and is deﬁned on any P 2 P 1 ( K ) by:
[Equation]
With this “geometric readout layer” added to our model, we obtain the following variants of our main results in this non-convex, but geometrically regular, setting.
Corollary 2.11 (Constrained Universal Approximation: Riemannian Case) .
Consider the setting and notation of Corollary 2.8.
Let L ( x, y ) =
k f ( x )  y k .
If Assumption 2.10 holds then:
[Equation]
is a well-deﬁned Lipschitz-continuous function, and the following hold:
(i) Exact Constraint Satisfaction: ˆ D
ˆ E ( x ) 2 K , for each x 2 X , (ii) Universal Approximation: sup X d g ( f ( x ) , ˆ D  ˆ E ( x ))
< ✏ K + kd ✏ f .
The “complexities” of ˆ D and ˆ E are recorded in Table 15 1 for ✏ 2 = ✏ k = ✏ f .
D ISCUSSION
In this paper, we derived the ﬁrst constrained universal approximation theorems using probabilistic reformation of Vaswani et al.
(2017)’s transformer networks.
The results assumed both a quantitative form (Theorem 2.7) and a qualitative form in the more general case of an arbitrary loss functions L and additional compatible soft constraints in (Theorem 2.2).
Our results provide (generic) direction to end-users designing deep learning models processing non-vectorial structures and constraints.
As this is the ﬁrst approximation theoretic result in this direction, there are naturally as many questions raised as have been answered.
In particular, it is natural to ask: “ Are the probabilistic transformer networks trainable in practice; especially when K is non-convex?” .
In Appendix 5, we show that the answer is indeed: “Yes!” , by proposing a training algorithm in that direction and showing that we outperform an MLP model and a classical transformer network in terms of a joint MSE and distance to the constraint set.
The evaluation is performed on a large number of randomly generated experiments, whose objective is to reduce the MSE to a randomly generated function mapping a high-dimensional Euclidean space to there sphere R 3 with outputs constrained to the sphere. A CKNOWLEDGMENTS
Anastasis Kratsios and Ivan Dokmani´c were supported by the European Research Council (ERC)
Starting Grant 852821—SWING.
The authors thank Wahid Khosrawi-Sardroudi, Phillip Casgrain, and Hanna Sophia Wutte from ETH Z¨urich, Valentin Debarnot from the University of Basel for their helpful feedback and Sven Seuken from the University of Z¨urich for his helpful feedback in the rebuttal phase.