Universal Approximation Under Constraints is Possible with Transformers
Authors: Anastasis Kratsios
Behnoosh Zamanlooy
Tianlin Liu
Ivan Dokmanić
Venue: ICLR 2022 Spotlight (2022)
Keywords: Constrained Universal Approximation
Probabilistic Attention
Transformer Networks
Geometric Deep Learning
Measurable Maximum Theorem
Non-Affine Random Projections
Optimal Transport.

Abstract: Many practical problems need the output of a machine learning model to satisfy a set of constraints, $K$.  Nevertheless, there is no known guarantee that classical neural network architectures can exactly encode constraints while simultaneously achieving universality.  We provide a quantitative constrained universal approximation theorem which guarantees that for any non-convex compact set $K$ and any continuous function $f:\mathbb{R}^n\rightarrow K$, there is a probabilistic transformer $\hat{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's Maximum Theorem (1963).  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 $\hat{F}$ that approximately minimizes $L$ and whose outputs belong to $K$; moreover, $\hat{F}$ approximately satisfies the soft constraints.  Our results imply the first universal approximation theorem for classical transformers with exact convex constraint satisfaction.  They also yield that a chart-free universal approximation theorem for Riemannian manifold-valued functions subject to suitable geodesically convex constraints.

Published as a conference paper at ICLR 2022
UNIVERSAL APPROXIMATION UNDER CONSTRAINTS IS POSSIBLE WITH TRANSFORMERS
Anastasis Kratsios, Tianlin Liu & Ivan Dokmanié Universität Basel,
Behnoosh Zamanlooy" Universität Zürich, Department of Informatics
Departement Mathematik und Informatik
(firstname. Lastname] ) aanibas...
Saamanlooye i fi, wwh. ch
ABSTRACT
Many practical problems need the output of a machine learning model to satisfy set of constraints, K. There are, however, no known guarantees that classical neural networks can exactly encode constraints while simultaneously achieving universality. We provide quantitative constrained universal approximation theo- rem which guarantees that for any convex or non-convex compact set K and any continuous function f R" - K, there is a probabilistic transformer F whose randomized outputs all lie in K and whose expected output uniformly approxi- mates f. Our second main result is a "deep neural version" of Berge (1963)'s Maximum Theorem. The result guarantees that given 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, approximately satisfies the soft constraints. Our results imply the first universal approximation theorem for classical transformers with exact convex constraint satisfaction, and chart-free universal approximation theorem for Rie- mannian manifold-valle functions subject to geodesically-comee constraints.
Keywords: Constrained Universal Approximation, Probabilistic Attention, Transformer Networks, Geometric Deep Learning, Mcasurable Maximum Theorem, Non-Affine Random Projections.
INTRODUCTION
In supervised learning, we select parameterized model f: R -> Rm by optimizing real-valued loss function's L. over training data from an input-output domain X' x y S R" x 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 (R")) is present in learning problem and this structure must be encoded into the trained model ff to obtain meaningful or feasible predictions. This additional structure is typically descrihed by a constraint set K C RR and the condition f(x) C K. For example, in classification K = { E [0, 1|" EE Y: = 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 atility-maximinn programming K the integer lattice 7 (Conforti et al., 2014), in financial risk-management K set of positions meeting the minimum solvency requirements imposed by international regularity bodies (Basel Committee on Banking Supervision, 2015; 2019; McNeil et al., 2015), covariance matrix prediction K Rm is the set of m matrices which are symmetric and positive semi- definite (Bonnabel et al., 2013; Bonnabel & Sepulchre, 2009; Baes et al., 2021), in geometric deep learning K is typically manifold (e.g. 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:
actions of an opponent, in integer
Is exact constraint satisfaction possible with universal deep learning models?
Corresponning authors.
For example, in regression problem one set T.x,,, 3) ll((z)) bll for unknown function f or in regression problems T( = ([))) (og( for unknown classifier C.

---

Published as a conference paper at ICLR 2022
The answer to this question begins by examining the classical universal approximation theorems for deep feedforward networks. L and K are mildly regular, the universal approximation theorems of Hornik et al. (1989); Cybenko (1989); Pinkus (1999); Gühring et al. (2020); Kidger & Lyons (2020); Park et al. (2021) guarantee that for any "good activation function o" and for every tolerance level >> O, there is a deep feedforward network with activation function a, such that inf yek T(x,, y) and L(x, (x)) are uniformly at most e apart. Written in terms of the optimality set,
sup llf(x) mEX
argmin Lax, y)|ls e,
(1)
NER
where the distance of a point y E R to a set A C RT defined by |1y - All _ inf lly - all- Since argmingex [zz v) S K, then (1) only implies that [lf(xz) - K| < and there is no reason to believe that the constraint f(x) E K is exactly satisfied, for every æ E 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-vectur vectoril 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) & K closest point in K. Thus, we are faced with the dilemma: either make an ad-hoc choice of a k in K with k = (xx) (ex.: an arbitrary choice scheme when ZT) or have meaningless predictions (ex: non-integer values to integer programs, or symmetry breaking (Weinberg, 1976)²).
Constrained learning was recognized as an effective framework for fairness and robustness by Cha- mon & 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 model- based domain generalization (Robey et al., 2021). Despite the importance of (deep) learning with constraints, there are no related ixproximation-thhoree 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 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 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 non- Euclidean 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 RT)) in the special case of L(r, y)
(f(x)-// for a suitable target function f RR" > RT 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-theorriie counterpart to the constrained statistical learning theory of Chamon & Ribeiro (2020). Further, they put forward perspective on random- ness 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 effi- ciency. Finally, our proof methods are novel, and build on contemporary tools from metric geometry (Ambrosio & Puglisi, 2020; Bruè al., 2021).
2 As discussed in Rosset al. (2021) this is problematic since respecting symmetries can often massively reduce the computational hurden of a learning task.
2

---

Published as a conference paper at ICLR 2022
1.1 THE PROBABILISTIC ATTENTION MECHANISM
We now give high-level explanation of our results; the detailed formulations are in Section 2.
Introduced in (Bahdanau al., 2015) and later used define the transformer architecture (Vaswani et al., 2017), in the NLP context, attention maps matrix of queries Q, a matrix of keys K, and matrix of values V 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) ف Softmaxx (w) T Y=
(w)
(2)
where w RR, Softmaxx : RN e w + - k=1 and Y an X m matrix. The attention
mechanism (2) can be interpreted as "paying attention" to a set of particles Yi.... ......,Ye Yy E R" 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 quan- tities defining (1) are regular then, for any e E (0, 1]. there a feedforward network f. an X. C R" of probability 1-c, and matrix Y such that the transformer Attention( (xx), Y) satisfies:
(i) Exact Constraint Satisfaction: For each E RR, Attention([(()). Y)ε K, argmin Laz Ss y'ek
(ii) Universal Approximation: supzex, |tttettion(((t)
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 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 convex combinations of the particles Y1, Y2, Y3; thus, any transformer network's 's predictions must belong to these particles' convex hull.
X
X
Figure 1: Convex Constraints
Figure 2: Non-Convex Constraints
In Figures and 2, Y's columns, i.e. the particles Y1, Y2, and Y3, are each illustrated by a the constraint set (K) vertices. The bubble around each each Y, illustrates the predicted probability, for a given input, that f(x) is nearest to that Y. The X is the transformer's prediction which is, by construction, a convex combination of the Y, weighted by the aforementioned probabilities and therefore they lie the K if it convex (Figure 1) but not if K is non-convex (Figure 2).
Naturally, arrive the question: How (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 Softmax, layer is to rank the importance of the particles (Y,, 1.N when optimizing L, at any given input: the weights [Softmaxx (w)],, in (2) can interpreted as charging their respective point masses {ōy, 1 with probabilities of being optimal for L (relative to the other particles)'. This suggests the following probabilistic reinterpretation of attention (which we denote by p-attention):
P-attention Y)
Σ
Following Villani (2009), ōy, is the Borel probability measure on RT assigning full probability to any Borel subset of RM containing the particle Y, and 0 otherwise.
3

---

Published as a conference paper at ICLR 2022
Crudely put, P-attention(; Y) "pays relative attention to the particles" Yı,.... Yn E RT
A simple computation shows that the mean prediction of our probabilistic attention mechanism, exactly implements "classical" Attention of Vaswani al. (2017), as defined in (2),
Attention(w: w. Y) = [xx--aateemtioon(:.)]] [X],
(4)
where IE X [X] denotes the (vector-valued) expectation of a random-vector X dis- tributed according to P-attention(w: Y). Hence, (3) is no less general than (2). The advantage of (3) is that, if each particle Y, belongs to K (even if K is non-convex) then, any sample drawn from the probability measure P-attention( ww Y) necessarily belongs to K.
1.2 QUALITATIVE RESULTS: DEEP MAXIMUM THEOREM
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" (Berge, 1963), which states that under mild regularity conditions, given any well-behaved family of input dependent "soft constraint sets" {Cz) compatible with K, there measurable function mapping each R" to minimizer of L(x, y) on K C...
We use W1 to denote the Wasserstein-1 distance between probability measures on K. The re- sults also give the flexibility the user to enforce an input-dependent family of "soft constraints" (ER" which only need to hold approximately; definitions are provided in Section 1.4.
Non-Convex Case). If the quantities defining regular, K is compact set of "exact constraints", and {CZ)zeR> set of "soft constraints", then, for any approximation quality 0 < E 1, there is a deep feedforward network f and matrix Y satisfying:
Informal Theorem 1.2 (Deep Maximum Theorem:
(i) Exact Constraint Satisfaction: For each x E R", P-attention(f (z)), Y) is supported in K, (ii) Universal Approximation: (P (P-ateention(( (z)), Y), argmin L(x, )*)) S <) 1 y ,CC,n
where for a probability measure on R" and a B C RT we define Wi(P B) & infreB Wi(P, 86). Example 1.3 (Reduction Classical Point-to-Set Distance). In particular, when P is point-mass
F P= by for some E R"" then one recovers the familiar Euclidean distance to the set B via:
Wi (6 B) (def)
1 bll (def)
inf Wi (8y, inf bEB
lly IBB;;
BEB
where the first and second equality follows from (Villani, 2009, (5)-page 99), and the last equality is the definition of lyy - B|| (as in (Aubin & Frankowska, 2009, Definition 1.1.1)).
Another important class of non-convex constraints arising from geometric deep learning where K is non-Euclidean ball in a Riemannian submanifold of R In this broad case, we may extract mean predictions from P-attention Y), by applying the Fréchet mean introduced in Fréchet (1948). Such "geometric means" are well-understood theoretically (Bhattacharya & Patrangenaru, 2003) and easily handled numerically Miolane et al. (2020); Lou al. (2020).
1.3 QUANTITATIVE RESULTS: CONSTRAINED UNIVERSAL A APPROXIMATION THEOREM
In its current form, the objective function L is too general to derive quantitative approximation rates Nevertheless, as with most universal approximation theorems (Hornik et al., 1989; Pinkus, 1999; Kidger & Lyons, 2020), if each soft constraint C, is set to RM and L quantifies the uniform distance to an unknown continuous function f RD - K in the Euclidean sense,
L(a, v) - 1f(x)) - vll-
then, Informal Theorem 1.2 reduces (qualitative) universal approximation for transformer net- works 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
"More precisely, our result a deep neural version the measure-theoretic counterpart to Berge's 's Maxi- mum Theorem; see (Aliprantis & Border, 2006, (Measurable Maximum Theorem) Theorem 18.19). For instance, can describe anything from a regression, clustering problem.
4

---

Published as a conference paper at ICLR 2022
K' is contained in unknown d-dimensional submanifold (where d ((m²) for some 0). log(m), Our approximation rates scale favourably in the ratio R3 log(d] i.e., we avoid the curse of dimen- sionality for low-dimensional constraint sets. This additional structure translates into the familiar encoder-decoder structure deployed in most transformer network implementations.
Figure 3: Encoder f
Figure 4: Decoder: Random Projection K
Figure 3 illustrates the encoder network & R + R" whose role is to perform (classical) unconstrained approximation of the target function, f. Since & is classical feedforward network then approximation of the target function can be arbitrarily close to the constraint set K need not lie in it. The next step is "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 projection does not exist if K is not convex since there must be more than one closest point in to some y 1Rm (Motzkin, 1935). Nevertheless, if the "projection" were capable of mapping any y IRTH to multiple points K, ranked by their proximity to y. then there would be no trouble. The decoder network accomplishes preciscly this, as illustrated in Figure 4, where the bubbles illustrate the probability of any particle in K being closest y. illustrated by the size of the bubbles in Figure 4. Mathematically, D Rm - PI(K) approx- imates (non-affine) random projection, in the sense of Ohta (2009); Ambrosio & Puglisi (2020); Bruè et al. (2021); i.e.: 1-Lipschitz map II RR" + Pi PI(K) satisfying the random projection property: for all y E K
11 8y.
Thus, II's random projection property means that it fixes any output already satisfying the constraint K, and its Lipschitz regularity implies that it is stable. Thus, sampling from I(()1) is similar to sampling from II())) whenever the points y1, Y2 E R" are near to one another.
Remark 1.4. Random projections are closely lied 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 affine.
Remark 1.5. The special case of random projections onto affine 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, L,B,C.I. D Z are constants independent of and k; where k N4 is the number of continuous derivatives which f admits (when viewed as function into RT"). From
Network
D
Depth
Om + 3ck kn+1))
2n 2n
0 ((N+
(IN+AA+ (A)2 1)))))
Width
m
In + 10)
m² NNNNNN +2 0 ((e A + B)
N
Q
Table Complexity of simple transformer network f= D approximating f. Table 1, we see that if m < m then, > 0 is large; hence, " 11-46 46 1)) ) and N are small.
1.4 NOTATION AND BACKGROUND
Optimal Transportd Given any non-empty subset K R", the set of all Borel probability measures P on K with finite mean; i.e.: Ex-r co, is denoted by P (K). Wasserstein
(P((K) denotes the Wasserstein space K, and is defined formally below.
5

---

Published as a conference paper at ICLR 2022
distance Wi is defined for any Q.Q E Pi(K) by the minimal energy needed to transport all mass from P to Q. Following Villani (2009), (I(P, Q) is defined by:
(I Q) inf [a) - X21l.:
where the infimum is taken over all Borel probability measures a K22 with marginals PP and Q- The metric space (Pi (K), W)i is named the Wasserstein space over K; we abbreviate it by Pi (K). Smooth Function Spaces The set of real-valued continuous functions R" is denoted by C(R"). Let k E IN+ and X' C [0. 11"d be non-empty. The set of functions f : X' + K for which there is k-times continuously differentiable f R + Rm extending f: i.e.: flx = ʃ. is denoted by Ch K). Our interest in C (X,, K) does not stem from the fact that it contains all smooth func- tions mapping [0, 1]"h to K, but rather that it allows us to speak about the uniform approximation of discontinuous K- K-valued functions on regions in [0, 1]" where they are "regular". This is noteworthy for pathological constraint sets, such as integer constraints'. For details on C₁2, (X, K). see (Brudnyi & Brudnyi, 20;2a; and the extension theorems of Whitney (1934); Fefferman (2005).
Neural Networks has recently been observed that deep feedforward networks with multiple activation functions, or more generally parametric families of activation functions, achieved signif- icantly more efficient approximation rates than classical feedforward networks with a single acti- vation function (Yarotsky& & Zhevnerchuk, 2020; Yarotsky, 2021; Shen et al., 202 la: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. (2018). We also observe a similar phenomenon, and therefore our quantitative results consider deep feedforward networks whose activation functions belongs to a 1- parameter family a, {or (ej.1 S C(R). The set of all such networks is denoted by N n.N and
it includes all f Rm - RRN with iterative representation:
f(x) - (1)
(++1) A otz, ((A4))) x),, + (O))
xe (0)
(5)
where x E R 1, ...... . J 1, each A dj x dj+1-matrix, each b60 E Rdst,, dب41 N,
dj = 004,,1,...,,14 IIN,, E [o, 1). for each j. The integer J is f's depth and max dj is I's width. j=1....J41
Example 1.6 (Networks with Untrainable Nonlinearity). Denote a - 06- The subset of classical feedforward networks consisting of all NN nN with each Ot., = a in (5) is denoted NN N.NN
It is approximation theoretically advantageous to generalize the proposed definition of probabilistic attention in the introduction (3) by replacing Y with a 3-dimensional array (elementary 3-tensor).
Definition 1.7 (Probabilistic Attention). Let N, Q.m E N+, and Y be an N x Q x m-array with 1,..., Q. Probabilistic attention is the function:
Yn E K for 1,.... 1=
N Q
R" 3 w + P-attention(w, Y) & EE Coftmax((w),bb E PK)K).. n=19=1
Y is an X m-matrix, as in (3), then we identify as the N x X I-array in the obvious manner.
Set- Velued Analysis: A family of non-empty subsets {CZ)ERR- of K is said to be weakly measurable correspondence, denoted C RD = RmM if for every open subset"U K K, {a E IR" C2 UU + 0} is non-empty Borel subset of R (Aliprantis & Border, 2006, pages 557, 592).
2
MAIN RESULTS
We now present our main results in detail. All proofs are relegated to the paper's appendix.
2.1
QUALITATIVE APPROXIMATION: DEEP MAXIMUM THEOREM
Our main qualitative result is the following deep neural version of Berge (1963)'s Maximum The- orem where, the measurable selector is approximately implemented by probabilistic transformer
For example, there is no non-constant continuous function f [0, 1] + 7. However, for any À E (0. }) and any 3/, E 2. = y1 |oo .I + Velp+g.a ja belongs to C& (10, AJ U [A + 2 1], 2) for each & E N4- SSince K is equipped with its subspace topology, then open subset U of K is any subset of Rm of the form U U2 n K where U, is an open subset of R (see (Munkres, 2000, Lemma 16.1 for further details).

---

Published as a conference paper at ICLR 2022
network. We first present the general qualitative result which gives concrete description of mea- surable selector (I), with high-prober probbabiib,,, which has the key property that all its predictions satisfy the required constraints defined by K.
Assumption 2.1 (Kidger& & Lyons (2020)). : R + R is continuous, a is differentiable at some Fo R, and its derivative satisfies o' (x()) &
Theorem 2.2 (Deep Maximum Theorem). Let satisfy Assumption 2.1. Let K S R" be non- empty compact set, C: R" = RT be a weakly- correspondence with closed values such that Oz O K + 0 for each E IR", L E C(RT",, and IP be Borel probubility measure on R". For each 0 < S 1, there is N E N4. an f NNN, NN of width most N X m-matrix Y such that:
N, and
FR 3 + P-attention (f(x),)) E (P(")"))
(6)
satisfies the following:
(i) Exact Constrain Satisfaction: UseR" supp(F(z)) S K.
(ii) Probably Approximately Optimality: There is a compact X. S RR satisfying:
(a) max (F(),, argmin Læ, v)) S, X
yEC,nK
(b) 1 P(X.) S.
Theorem 2.2 implies that for any random field (2E)) R" (i.e. family of -allue random
vectors indexed by R") with y= F(x): 1. samples drawn from Y are in K (by (i)) and 2. samples drawn from each Y are near to the optimality set argmin yeC, L(x, v) (by (ii).
Corollary 2.3 (F's Mean Prediction). Assume the setting of Theorem 2.2. Let {Y*) be K-valued random field with y ~ F(x) for each x E R then, - P(X.) S c and
max Elj - argmin L(x, y yI] S. x
y*EC,nk
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.2 QUANTITATIVE APPROXIMATION. CONSTRAINED UNIVERSAL APPROXIMATION
In order to derive a quantitative constrained universal approximation theorem, require the loss function to be tied to the Euclidean norm in the following manner.
Assumption 2.4 (Norn-Controllable Loss). There is a continuous f IR" Rm with f(R") S K and a continuous / [0, oc) + [0, oc) with ((0) 0, satisfying: L(x, v) S ((1f(())(//))) yll).
Just as with transformer networks, our "constrained universal approximation theorem" approximates suitably regular function f Ra + K S [RM while exactly respecting the constraints K by implementing an encoder-decoder network architecture. Thus, our model is a composition of an encoder network & R" + Rd whose role is to approximate f in classical "unconstrained fashion" and decoder network (with probabilistic attention layers at its output) D Rd - Pi PKK) whose role to enforce the constraints K while preserving the approximation performed by É, where d <<
To take advantage of the encoder-decoder framework present most transformer networks, we formalize what is often called "latent low-dimensional manifold" hypothesis. Briefly, this means that, the hard constraints in set K are contained in a "low dimensional" subspace.
Assumption 2.5 Low--Diensiional Manifold), There is an O and a smooth bijection & from R" to itself with smooth inverse", such that & ((K) S Rd; where s d and E O(m+).
Assumption 2.5 does not postulate that K is itself a single-chart now-dimensional manifold, or even 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) follows.
Assumption 2.6 Swish-Liie Activation Function). The map a [o, 1] x R 3 (a, + oa(() E is continuous; do non non-affin -ffine and sercewissellnedde and is smooth and non non-polynomial.
Here smooth means that is continuously differentiahle any number of times. NE, smooth bijections with smooth inverses are called diffeomorphisms in the differential geometry and differential topology literature. Following Jost (2017), function RR R called smooth class ( if a exists for each l e N4-

---

Published as a conference paper at ICLR 2022
Theorem 2.7 (Constrained Universal Approximation). Let l E N+ and X C [0, 11"h be non-empry. Suppose that a satisfies 2.6, satisfies Assumption 2.4. K S RR non-empry, compact and sat- isfies Assumption 2.5. For any f C (i, K) every constraining quality сK O, and every approximation error Ef O, there exist N, Q N4. an encoder & ENN n, and a decoder:
D: D:R
P-altention (D(x), E Pi(KK
where ò NN d. and Y an N Q x m-array with Y1. YN.Q E K such that:
(i) Exact Constrain Satisfaction: For each x E R": supp E(x) SK,
(ii) Universal Approximation: The estimate holds" !1:
WD o (x), argmin L(z, u)) S <K + LLip( )dyy; yEK
sup
==(0,1]"
where, 0 < k is an absolute constant independent n., m, d, and of and Lip o denotes
the Lipschitz constant of ④h on the compact set {2 E Rd ||2 4 - cy).
Furthermore, the "complexities" of D and & are recorded in Table for c ct-
In practice, we can only sample from each measure (z). In this case, we may ask how the typical sample drawn from random-vector Y distributed according to our learned measure D (xx) performs when minimizing [xx, y). The next result relates the estimates in Theorem 2.7 (ii) to the typical (in Y") worst-case (in z) gap between sample from Y and f(z). as quantified by L -). Corollary 2.8 (Average Worst-Case Loss). Assume the setting of Theorem 2.7 and suppose that the "modulus" " 1 in Assumption 2.4 is strictly increasing and concave. Let D and & be as in Theorem 2.7 and let {Y' JzEx be an Rr"-valued random field with Y= ~ DD (x), for each x E R". Then:
max Ey=nDoe(s) [I) S I (ex + Lip(@@'>)da
#EX
Corollary 2.8 quantífies the expected performance of sample from our probabilistic transformer model, expressed by I., whereas Theorem 2.7 (ii) quantifies the difference from the transformer's prediction to the optimal prediction value. Next, we consider implications of our main results.
2.3 APPLICATIONS
We apply theory obtain universal approximation theorem for classical transformer networks with exact convex constraint satisfaction and to derive a version of the non-Euclidean universal ap- proximation theorems of Kratsios Bilokopytov (2020); Kratsios & Papon (2021) for Riemannian- manifold valued functions which does not need explicit charts. As with most quantitative (uniform) universal approximation theorems (Gühring et al., 2020; Kidger & Lyons, 2020; Shen et al., 2021a), we henceforth consider L(x, y) [z](z) y/|. We also fix f E C (10, 1J" K).
2.3.1 TRANSFORMERS ARE CONVEX-CONSTR.CN CONSTRAINED UNIVERSAL APPROXIMATORS
We return to the familiar transformer networks of Vaswani et (2017). 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 field {Y") JoER" of Corollary 2.3 must belong to K. Consequently, the identity (4) implies that Attention D o {(), 1 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, v) = [lf((z) - vll. Then:
R 3 x + E[y Attention (x), E K;
(8)
(i) Exact Constraint Satisfaction: EyonDoš(e) [ for each x E RR"
(ii) Universal Approximation: suP(o.1]" (I[(x) If (===D((())) [Y* < kdeg-
The "complexities" of the networks D and & are recorded in Table 1 for ; = =f..
"In fact, we actually prove that slightly stronger statement: sup=ejo.1j" W1 (D ((x)) S ck +
A: Lip (app($=")dee yde.. Both formulations align when has a unique minimum as is the when L(z, v) R
[] 2||- and || ||. any 12 Explicit constants are recorded in Table 2 within the paper's appendix; there, ex and ey may differ. Expliii constants are recorded in Table 2 within the paper's appendix; there, eg and ey may differ.
8

---

Published as a conference paper at ICLR 2022
2.3.2 CHART-FREE RIEMANNIAN MANIFOLD- VALUED UNIVERSAL APPROXIMATION
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 (de- terministic) 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 expec- tation, 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 TR with distance function by dg. We only require the following mild assumption introduced in Afsari (2011). We recall that the injectivity radius at Vo. denoted by inf i (36). (see (Jost, 2017, Definition 1.4.6)) is the minimum length of a geodesic (or minimal length curve) in M with starting point yoo. We also recall that the sectionul curvature (see (Jost, 2017, Definition 4.3.2) for a formal statement) quantifies the curvature of (M, .g) as compared the geometry of its flat counterpart Rd. We focus on a broad class of non- convex 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 con- nected, is complete as a metric space, and all its sectional curvatures of (M,9) are all bounded above by constant C 2 0. The non-empty constrain set K satisfies:
1. K is contained in the geodesic ball B p) - {y E M dg (g(00,3) y) < p} for some point yo E M and some radius p satisfying 0 < p < 2-1 min (20), Jch,
2. For each you 91 E K there exists a unique geodesic y [0.1] + K joining yo to y1-
Our latent probabilistic representation grants the flexibility of replacing the usual "extrinsic mean" used in (8) to extract deterministic predictions from our probabilistic transformer networks via an additional Fréchet mean layer their readout. This intrinsic notion of mean, was introduced independently in Fréchet (1948) and in Karcher (1977), and is defined on any IP E (K) by:
/ dk u) P(du).
P argmin KEK
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) = (z) - vll. If Assumption 2. 10 holds then:
R" 3 x + D ((r E K, is well-defined Lipschitz function, and the following hold:
(10)
(i) Exact Constraint Satisfaction: D &x)( E for each x E X, (ii) Universal Approximation: supx dx(f(x), (o((x) < EK + kdes.
The "complexities" " D and E are recorded in T5 1 for = EN = Ef.
DISCUSSION
In this paper, we derived the first constrained universal approximation theorems using probabilistic reformation of Vaswani et al. (2017)'s 's transformer networks. The results assumed both a quantitative form (Theorem 2.7) and qualitative form in the more general case of 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 first approximation theoretic result in this direction, there are naturally as many ques- tions raised have been answered. In particular, it is natural to ask: "Are the probabilistic trans- former networks trainable in practice; especially when K is non-convex?". In Appendix 5, we show that the answer is indeed: "Yes!", by proposing 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 performed on a large number of randomly generated experiments, whose objective is to reduce the MSE to a randomly generated function mapping high-dimensional Euclidean space to there sphere R with outputs constrained to the sphere.
"Following Afsari (2011), 1), we make the convention that if C S 0 then, v/ is interpreted as co. 15 Explicit constants are recorded in Table 2 within the paper's appendix; there, eg and ey may differ.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
Anastasis Kratsios and Ivan Dokmanić were supported by the European Research Council (ERC) Starting Grant 85282I-SWING. The authors thank Wahid Khostawi-Sardroudi, Phillip Casgrain, and Hanna Sophia Wutte from ETH Zürich, Valentin Debarot from the University of Basel for their helpful feedback, and Sven Seuken from the University of Zürich for his helpful feedback in the rebuttal phase.
REFERENCES
Bijan Afsari. Riemannian L" center of mass: existence, uniqueness, and convexity. Proceedings of the American Mathematical Society, 139(():655-673, 2011.
Charalambos D. Aliprantis and Kim C. Border. Infinite dimensional analysis: A hitchhiker's guide. Springer, Berlin, third edition, 2006.
Luigi Ambrosio and Daniele Puglisi. Linear extension operators between spaces of Lipschitz maps and optimal transport. Journal für die Reine und Angewandte Mathematik. 764:1-21, 2020.
Anonymized. Pytorch implementation of attend-to-constrainee 2021. URL https: /drive.
com/ file/ /ttp://e/////yytyttttttttt...eee Ctt.://UUUUuttttoooo je 2swppgk
Jean-Pierre Aubin and Hélène Frankowska. Ser-valued analysis. Modern Birkhäuser Classics.
Birkhäuser Boston, Inc., Boston, MA, 2009.
Michel Baes, Calypso Herrera, Ariel Neufeld, and Pierre Ruyssen. Low-rank plus sparse decompo- sition of covariance matrices using neural network parametrization. IEEE Transaction on Neural Networks and Learning Systems, 2021.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In Proceedings of the International Conference Learning Rep- resentations (ICLR), 2015.
Yair Bartal. On approximating arbitrary metrices by tree metrics. In Proceedings of the Thirtieth Annual ACM Symposium on the Theory of Computing, pp. 161-168. ACM, New York, 1999.
Basel Committee Banking Supervision. Fundamental review of the trading book: outstanding issues, February 2015. https: www. bis www...rr......oo........o.ooo..t
wwwooorgbbbbb.. ppbb1/3355 pdf.
Basel Committee on Banking Supervision. Minimum capital requirements for market risk, February 2019. https: wwwp///wwwwwwo/ bis wtgpbbbs//pbbbb//555555 pdf..
Heinz H. Bauschke and Patrick L. Combettes. Convex analysis and monotone operator theory in Hilbert spaces. CMS Books in Mathematics/Ouvragee de Mathématiques de la SMC. Springer, 1001/007997--1 1-4419-9467-77 URL https: With foreword by Hédy Attouch.
New York, 2011. ISBN 978-1-4441-9946-00
doi: 10.
/ //doi. org 10. 1001/ 000000007777810 1-4---1-419999 -7
Claude Berge. Éspaces Topologiques (Topological Spaces). Dunod, 1963.
Rabi Bhattacharya and Vic Patrangenaru. Large sample theory of intrinsic and extrinsic sample means on manifolds. The Annals of Statistics, (1(1):1-29, 2003.
Silvère Bonnabel and Rodolphe Sepulchre. Riemannian metric and geometric mean for positive semidefinite matrices of fixed rank. SIAM Journal on Matrix Analysis and Applications, 31)) 1055-1070, 2009.
Silvère Bonnabel, Anne Collard, and Rodolphe Sepulchre. Rank-preserving geometric means of positive semi-definite matrices. Linear Algebra and its Applications, 438(8)33202-3211 2013.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geomet- ric deep learning: going beyond Euclidean data. IEEE Signal Processing Magazine, 34(4):18-42, 2017.
10

---

Published as a conference paper at ICLR 2022
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learn- ing: Grids, Groups, Graphs, Geodesics, and Gauges. arXiv:2104.08708, 2021. URL http: wtw.v..rrgaabs//1444
Bernard Bru, Henri Heinich, and Jean-Claude Lootgieter. Distances de Lévy et extensions des théoremès de la limite centrale et de Glivenko-Cantelli, Publ. Inst. Statist. Univ. Paris, 37(3-4): 29-42, 1993.
Alexander Brudnyi and Yuri Brudnyi. Methods of geometric analysis in extension and trace prob- lems. Volume volume 102 of Monographs in Mathematics. Birkhääseer/prrnger Basel AG, Basel, 2012a.
Alexander Brudnyi and Yuri Brudnyi. Methods of geometric analysis in extension and trace prob- lems. Volume 2, volume 103 of Monographs in Mathematics. Birkhäuser/Springer Basel AG, Basel, 2012b.
Elia Bruè, Simone Di Marino, and Federico Stra. Linear Lipschitz and C* extension operators through random projection. Journal of Functional Analysis, 288(4): 2021.
Luiz Chamon and Alejandro Ribeiro. Probably approximately correct constrained learning. In Proceedings of Advances in Neural Information Processing Systems (NeurlPS), 2020.
Michele Conforti, Gérard Cornuéjols, and Giacomo Zambelli. Integer Programming, volume 271
of Graduate Texts in Mathemutics. Springer, Cham, 2014.
Christa Cuchiero, Lukas Gonon, Lyudmila Grigoryeva, Juan-Pablo Ortega, and Josef Teichmann. Discrete-ting time signatures and randomness in reservoir computing. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-10, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport, In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), pp. 2292-2300, 2013.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Con- trol, Signals, and Systems, 204):303-3,,, 1989.
Meng Ding and Guoliang Fan. Multilayer joint gait- sait- manifolds for human gait motion modeling. IEEE Transactions Cybernetics, 45(11):2413 24222244442 2014.
Ivan Dokmanic, Reza Parhizkar, Juri Ranieri, and Martin Vetterli. Euclidean distance matrices:
Essential theory, algorithms, and applications. JEEE Signal Processing Magazine, 32(6):12-33 12-30. 2015. doi: 1091
Richard M. Dudley. Real analysis and probability, volume 74 of Cambridge Studies in Advanced Mathematics. Cambridge University Press, Cambridge, 2002.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Charles Fefferman. sharp form of Whitney's extension theorem. Annals of Mathematics, 161 (1):509-577, 2005.
Thomas Fletcher. Geodesic regression and the theory of least squares on Riemannian manifolds. International Journal of Computer Vision, 105(2):171-885 2013.
Maurice Fréchet. Les éléments aléatoires de nature quelconque dans un espace distancié. Annales de I' l'Institut Henri Poincaré, 10:215-310, 1948.
Lukas Gonon, Lyudmila Grigoryeva, and Juan-Pablo Ortega. Risk bounds for reservoir comput- ing. Journal of Machine Learning Research, 21(240):1-61, 2020a. URL papers/v21/2 19- 902 html.
org/
Lukas Gonon, Lyudmila Grigoryeva, and Juan-Pablo Ortega. Approximation bounds for random neural networks and reservoir systems. arXiv preprint arXiv:2002.05933. 2020b.
11

---

Published as a conference paper at ICLR 2022
Lyudmila Grigoryeva and Juan-Pablo Ortega. Universal discrete-time reservoir computers with stochastic inputs and linear readouts using non-homogeneous state-affine systems. J. Mach. Learn. Res., 19:Paper No. 24, 40, 2018.
Lyudmila Grigoryeva and Juan-Pablo Ortega. Differentiable reservoir computing. J. Mach. Learn. Res., 20:Paper No. 179, 62, 2019.
Ingo Gühring, Gitta Kutyniok, and Philipp Petersen. Error bounds for approximations with deep ReLU neural networks in W*P norms. Analysis and Applications, 18(5):803-859, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance ImageNet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Juha Heinonen. Lectures on analysis on metric spaces. Universitext. Springer-Verlag New York, 2001.
Ludger Holters, Björn Bahl, Maike Hennen, and André Bardow. Playing Stackelberg games for minimal cost for production and utilities. In ECOS Pell-Poocesdings of the 31st International Conference on Efficiency, Cost, Optimisation, Simulation and Environmental Impact of Energy Systems, pp. 36-36,36 University of Minho, 2018
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural Network, 2(5):359-366, July 1989.
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What is local optimality in nonconvex-nonccocave minimax optimization? In Proceedings of the International Conference on Machine Learning (ICML). 2020.
William B. Johnson and Joram Lindenstrauss. Extensions of Lipschitz mappings into Hilbert space. In Conference in modern analysis and probability, volume 26 of Contemp. Math., pp. 189-206. American Mathematical Society, RI, 1984.
Jürgen Jost. Riemannian geometry and geometric analysis. Universitext. Springer, Heidelberg, seventh edition, 2017.
Heinrich W. E. Jung. Über die Cremonasche Transformation der Ebene. J. Reine Angew. Math., 088:255318, 1910. ISSN 0075-4102. doi: 19111
Olav Kallenberg. Foundations of modern probability, volume 99 of Probability Theory and Stochas- tic Modelling. Springer, Cham, third edition, 2021.
Hermann Karcher. Riemannian center mass and mollifier smoothing. Communications on Pure and Applied Mathematics, 30(5):509-541 1977.
Patrick Kidger and Terry Lyons. Universal approximation with deep narrow networks. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Machine Learning Research, volume 125, pp. 2327. PMLR, 09-12 Jul 2020.
Achim Klenke. Probability theory: comprehensive course. Universitext. Springer, London, sec- ond edition, 2014.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced Wasserstein distances. In Proceedings of Advances in Neural Information Processing Sys- tems (NeurIPS), 2019.
Anastasis Kratsios. Universal regular conditional distributions. arXiv preprin::2105.07743, 2021.
URL https::/ http:://axxvvvoo org/ abs/ /205. 07743.
Anastasis Kratsios and Fugene Bilokopytov. Non-Fuclidean universal approximation. In Proceed- ings of Advances in Neural Information Processing Systems (NeurIPS), 2020.
Anastasis Kratsios and Leonie Papon. Universal approximation theorems for differentiable geomet- 2101 ric deep learning, 2021. URL https: http///aavvv/v.////ae//ee/e.aall 220...0330.........111
12

---

Published as a conference paper at ICLR 2022
James R. Lee and Assaf Naor. Extending Lipschitz functions via random metric partitions. Inven- tiones Mathematicae, 160(1):59-99) 2005.
Haochuan Li, Yi Tian, Jingzhao Zhang, and Ali Jadbabaic. Complexity lower bounds for nonconvex- strongly-coneave min- max optimization. arXiv:2104.08708. 2021. URL https://arxiv. org/aba hts//bb//t 2104 08708.
Yuanyuan Liu, Fanhua Shang. James Cheng, Hong Cheng, and Licheng Jiao. Accelerated first- order methods for geodesically convex optimization on riemannian manifolds. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2017.
Aaron Lou, Isay Katsman, Qingxuan Jiang, Serge Belongie, Ser-Nam Lim, and Christopher De Sa. Differentiating through the Fréchet mean. In Proceedings of the International Conference Machine Learning (ICML), 2020.
Cosme Louart, Zhenyu Liao, and Romain Couillet. A random matrix approach to neural networks. Ann. Appl. Probub., 28(2): 199--1248, 2018. ISSN 1050-5164. doi: 10.1214/1--AAP1328 URL btttps://don .acloooorroooooo .rgg 1211111-1111111111111
Mantas Lukoševičius and Herbert Jaeger. Reservoir computing approaches to recurrent neural net- work training. Computer Science Review, 3(3):177-14,, 2009, ISSN 1574-0137.
Alexander J. McNeil, Rüdiger Frey, and Paul Embrechts. Quantitative Risk Management: Concepts, Techniques and Tools. Princeton Series Finance. Princeton University Press, Princeton, NJ, 2015.
Nina Miolane, Nicolas Guigui, Alice Le Brigant, Johan Mathe, Benjamin Hou, Yann Thanwerdas, Stefan Heyder, Olivier Peltre, Niklas Koep, Hadi Zaatiti, Hatem Hajri, Yann Cabanes, Thomas Gerald, Paul Chauchat, Christian Shewmake, Daniel Brooks, Bernhard Kainz, Claire Donnat, Susan Holmes, and Xavier Pennec. Geomstats: A python package for riemannian geometry in machine learning. Journal of Machine Learning Research, 11 2020.
Theodore Samuel Motzkin. Sur quelques propriétés caractéristiques des ensembles bornés non convexes. Bardi, 1935.
James R. Munkres. Topology. Prentice Hall, Inc., Upper Saddle River, NJ, 2000. Second edition.
Shin-ichi Ohta. Extending Lipschitz and Hölder maps between metric spaces. Positivity, 13(2): 407-425, 2009.
Sejun Park, Chulhee Yun, Jaeho Lee, and Jinwoo Shin. Minimum width for universal approximation. International Conference on Learning Representations (ICIR), 2021.
Ofir Pele and Michael Werman. Fast and robust Earth Mover's distances. In Proceedings of the 12th IEEE International Conference on Computer Vision (ICCV), pp. 460-467, 2009.
Philipp Petersen and Felix Voigtlaender. Equivalence of approximation by convolutional neural networks and fully-connected networks. Proceedings of the American Mathematical Society, 148 (4):1567-1581, 2020.
Allan Pinkus. Approximation theory of the MLP model in neural networks. Acta Numerica, 1999, 8:1 1:143-195, 1999.
Chayne Planiden and Xianfu Wang. Most convex functions have unique minimizers. Journal of Convex Analysis, 23(3):877-892, 2016.
Pakize Simin Pulat. On the relation of max-flow min-cut for generalized networks. European Journal of Operational Research, 39(1):103- 107, 1989.
Michael Puthawala, Konik Kothari, Matti Lassas, Ivan Dokmanić, and Maarten de Hoop. Glob- ally injective ReLU networks. wwXiv:2006008466 2020. URL https:///exxi..rrr org/abs/ 2105.077 07743.
/
Prajit Ramachandran, Barret Zoph, and Quoc Le. Searching for activation functions. In Proceedings of the International Conference of Learning Representations (ICLR), 2018.
13

---

Published as a conference paper at ICLR 2022
Alexander Robey, George J Pappas, and Hamed Hassani. Model-based domain generalization. arXiv:2102.11436, 2021. URL https:/: wttp:/wwwwwwp/p////.///.2.2..
James C. Robinson. Dimensions, embeddings, and attractors, volume 186 Cambridge Tracts in Mathematics. Cambridge University Press, Cambridge, 2011.
Denis Rosset, Felipe Montealegre-Mora. and Jean-Daniel Bancal. RepLAB: A computa- tional/numerical approach to representation theory. In Quantum Theory and Symmetries, pp. 643-653, Springer, 2021.
Shai Shalev-Shwarzz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, USA, 2014.
Zuowei Shen, Haizhao Yang, and Shijun Zhang. Neural network approximation: Three hidden layers are enough. Neural Networks, 141:160-173, 2021a.
Zuowei Shen, Haizhao Yang, and Shijun Zhang. Deep network with approximation error being reciprocal of width to power of square root of depth. Neural Computation, 33(4):1005-1036, 03 2021b.
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast proba- bilistic approximation with exact solvers. Journal of Machine Learning Research, 20(105):1-23, 2019.
Sturm. Probability measures on metric spaces of nonpositive curvature. In Heat kernels and analysis on manifolds, graphs, and metric spaces, volume 338 of Contemp. Math., pp. 357-390. Amer. Math. Soc., Providence, RI, 2003.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Proceedings of Advances in Neural Information Processing Systems, pp. 6908, 2017.
Cédric Villani. Optimal Transport: Old and New, volume 338. Springer, 2009.
James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. On the regularity of attention. arXiv:2102.05668, 2021. URL https::///// http//aavvvvv.////////ee//..aev. owtp/vivorrraaooooo org /ttgg/ab///aa aatp//s/....../ 05628.
Steven Weinberg. Implications of dynamical symmetry breaking. Physical Review D, 13(4):974, 1976.
Hassler Whitney. Analytic extensions of differentiable functions defined in closed sets. Transactions of the American Mathematical Society, 36(1):63-89) 1934.
Dmitry Yarotsky. Elementary superexpressive activations. In Proceedings of the 38th International Conference on Machine Learning (ICML), 2021.
Dmitry Yarotsky and Anton Zhevnerchuk. The phase diagram of approximation rates for deep neural networks. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sseuuenee functions? In Proceedings of the International Conference on Learning Representations (ICLR), 2020a.
Chulhee Yun, Yin- Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O(n) connections are expressive enough: Universal approximability of sparse transformers. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2020b.
Hongyi Zhang and Suvrit Sra. First-order methods for geodesically convex optimization. In Pro- ceedings of the 29th Conference Learning Theory (COLT), 2016.
Ding- Xuan Zhou Universality of deep convolutional neural networks. Applied and Computational Harmonic Analysis, 48(2):787-794, 2020.
14


Reviews:
[{"review_id": "srbIudF_uGZ", "reviewer": "Reviewer_w3Aa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: + Provides rigorous theory to support the utility of transformers for deep neural network based predictive purposes under constrained settings.\n+ Soft constraint set that is defined using available data points is an interesting aspect that the paper considers. In fact, the authors were able to extend their theory to show that transformers can be used to model any data dependent constraint sets as well.\n+ Results of this paper may be applicable to practical use cases.\nWeakness: - The paper is hard to follow at many places since only informal theorems are provided throughout the paper.\n- The main theorem 2.2 is suboptimal -- optimal network may require width greater than the number of input dimensions n. For image analysis, this corresponds to the number of pixels.\n-Confusing notations.\n- No experiments to verify any of the quantitative results.\n", "summary_of_the_paper": "The paper intends to provide a stronger type of universal approximation result of transformer models. However, more details on experiments is required with a major revision of presentation and clarity.", "summary_of_the_review": "Justification: The constraint examples are too unnatural. For example, in geometric deep learning, authors argue that when K is a disjoint set, then the closest point computation becomes an unmanageable dilemma, claiming that the linear relaxations are meaningless for integer program (IP) while LPs are routinely used to solve IPs in practice using a branch and bound scheme. While some pieces of notations are kept constant throughout the paper, there are many confusing notations at various places that can easily be avoided - \"Attention\" in eq (4) is different from \"attention\". Adding to this issue, another central issue with the presentation is that all the technical statements in the main paper are Informal theorems with no rigorous explanations provided. This makes it very hard to read the paper because while the statements themselves may be locally correct or correct in isolation, it sometimes does not even make sense when put together as a whole, and makes it possible to verify the technical correctness of the paper. Somewhat related to the previous point, some technically deep concepts (such as Wasserstein-1 distance) are introduced in passing and handled in a slightly cavalier manner. It could have been nice if the authors provided some explanations to the figures -- at the moment, the reader has to guess what the authors intend to say. \n\nGiven that the paper shows quantitative results in the infinite sample setting, it is not clear whether the results hold in the finite sample setting where transformers are being used in practice. It would have been nice if the authors can formulate and show some toy experiments to even check whether such universality may be possible for dataset in the hands of a practitioner.\n\nAfter response: Thanks for the detailed response. I agree with the modifications, revisions, additional sections, and references proposed in the response. These changes were really helpful in parsing the paper for me since I have adequate mathematical background, but not so much experience in theoretical computer science. With more clearer notations, and figure descriptions, I can now see that the paper indeed focuses on universal approximation rather than generalization, and now I feel that my criticism about finite sample settings has been answered by the toy experiments also.  I choose the option of \"minor issues\" in correctness only because some statements still read handwavy in Section 1, for example, there are many places where the authors use the word \"any\", as in, any loss function, any constraint set, any goood activation etc.. While it may be correct in some sense, I find it odd to use in an otherwise rigorous paper. As an aside, I would like to point out this certainly does not hold if one wants to train from scratch since we know that loss functions and activation functions have a significant effect on first order algorithms. Thanks for providing additional references for us to check, and while going through them I found that some of the theoretical papers have very relevant practical use cases and I believe that they may add significant value to the overall scope and presentation given here, for example, Meta-learning with implicit gradients. I have raised my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "k198GFjemD5", "reviewer": "Reviewer_tJ9y", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n**Examples.**  Throughout the paper, the authors give numerous examples to illustrate why and how one might have constraints when training neural networks.  These were quite helpful in giving intuition for the (rather technical) contents of the paper.  This is a model I wish more technical papers would follow, as it allows even non-technical readers to get an idea of what the theory is all about.  Further, the authors do a commendable job of relating this work to the statistical results which have recently been proposed.  \n\n**Problem statement & setting.**  The problem statement is really clearly articulated.  The idea of providing a universal approximation property under constraints seems quite natural, and the authors motivate it extremely well.  Indeed, it is true that there are many works that argue for constraints in deep learning without considering the question of whether satisfying such constraints is even possible.  \n\n**Technical quality.**  I was not able to read through all of the proofs, especially given that the appendix gets pretty involved.  However, from what I saw, the proofs are of high quality.  Indeed, the generalization to the case of nonconvex K seems quite involved.  The proof techniques used here should be of wider interest.\n\n**Informal theorem presentation.**  The choice to first state the main results informally greatly helped the readability of the paper.  The main results are also illustrated with useful diagrams.  \n\n**Novelty.**  This is the first result of this kind that I am aware of.  For this reason, I believe that this paper has a significant claim to novelty.\n\n**Experiments.**  I like the experiments at the end of the appendix as well.  This helps to cover another natural question regarding how one can use these results in practice.\n\n### Weaknesses\n\n**On the attention module.**  I think the paper could benefit from further explaining why this particular form of attention is natural.  It would also perhaps help to illustrate this with an example.  As the credibility of the paper relies on the reader accepting that this attention model is natural and general enough to form the basis of the approach to universal approximation, I think that it's worth spending some more time here.  In particular, I didn't get a good sense of the role of the matrix $Y$ in build up to Informal Theorem 1.  \n\n**Figures.**  Don't get me wrong, I love that there are figures here.  However, I think that they could be explained better.  For instance, it wasn't clear to me what the green region was in Figure 1.  Is this $K$?  Also, the blue dot seems too small, and thus quite hard to see.  ", "summary_of_the_paper": "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.", "summary_of_the_review": "My review for this paper is rather short because overall I think that this paper is excellent.  It's clear that the problem is well-motivated and there are few results in the area.   To this end, the authors provide highly technical results to specifically address this novel and important question.  Moreover, the authors make a great effort to explain their results to a non-technical audience.  This is a clear accept in my opinion, and should be highlighted at the conference.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "FDtHDPkVDOF", "reviewer": "Reviewer_VnD6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. The paper shows the first universal approximation with constraints. The general theory shows that the universal approximation can use the probabilistic transformer to project the outputs in the set $K$. In particular, the probabilistic transformer helps to deal with non-convex constraints in Euclidean space and the geodesic convex constraints in non-Euclidean space, I.e., Riemannian manifold.\n\n2. The geodesically convex constraint seems to be an interesting and promising idea to treat non-convex constraints in the Euclidian space. The assumption on the sectional curvature is uniformly bounded above. Is there any connection that can be said in this case under Ricci curvature tensor? \n\n3. The paper is technically well presented. However, there are still some typos. The authors need to improve their writing. \n\n4. It would be a good idea to provide more examples to demonstrate the general theory. \n", "summary_of_the_paper": "The paper under review studies the universal approximation theory with constraints. \nFor any convex or non-convex compact set, a universal approximation is proved through a probabilistic transformer with constraints. Furthermore, a chart-free universal approximation is established on the Riemannian manifold with geodesically-convex constraints.", "summary_of_the_review": "The paper seems to address a well-known important class of problems. The universal approximation theorem might make significant impacts in solving optimization problems with constraints. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies an interesting question of whether neural networks can approximate the target function while keep the output in the constraint set. The constraint set is quite natural for  e.g. multi-class classification, where the output has to stay on on the probability manifold. The challenge here is that traditional universal approximation theory only guarantees that $\hat{f}(x) \approx f(x)$, but can not guarantee that $\hat{f}(x)$ lies exactly in the same constraint set as $f(x)$.

The paper made a significant contribution in the theory of deep learning -- It is shown that the neural network can indeed approximate any regular functions while keep the output stay in the regular constraint set. This gives a solid backup in terms of the representation power of neural networks in practice, to represent target functions whose outputs are in certain constraint set (e.g. probabilities).
Decision: Accept (Spotlight)