Year: 2022

Abstract:
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 : Rn ! 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 satisfies the soft constraints. Our results imply the first 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.

text: Keywords: Constrained Universal Approximation, Probabilistic Attention, Transformer Networks, Geometric Deep Learning, Measurable Maximum Theorem, Non-Affine Random Projections.

heading: 1 INTRODUCTION
text: In supervised learning, we select a parameterized model f̂ : Rn ! Rm by optimizing a real-valued loss function1L over training data from an input-output domain X ⇥ Y ✓ Rn ⇥ Rm. A necessary property for a model class to produce asymptotically optimal results, for any continuous loss L, is the universal approximation property. However, often more structure (beyond vectorial Rm) is present in a learning problem and this structure must be encoded into the trained model f̂ to obtain meaningful or feasible predictions. This additional structure is typically described by a constraint set K ✓ Rm and the condition f̂(X ) ✓ K. For example, in classification K = {y 2 [0, 1]m :Pm
i=1 yi = 1} (Shalev-Shwartz & Ben-David, 2014), in Stackelberg games (Holters et al., 2018; Jin et al., 2020; Li et al., 2021) K is the set of utility-maximizing actions of an opponent, in integer programming K is the integer lattice Zm (Conforti et al., 2014), in financial risk-management K is a set of positions meeting the minimum solvency requirements imposed by international regularity bodies (Basel Committee on Banking Supervision, 2015; 2019; McNeil et al., 2015), in covariance matrix prediction K ✓ Rm⇥m is the set of m⇥m matrices which are symmetric and positive semidefinite (Bonnabel et al., 2013; Bonnabel & Sepulchre, 2009; Baes et al., 2021), in geometric deep learning K is typically a manifold (e.g. a pose manifold in computer vision and robotics (Ding & Fan, 2014) or a manifold of distance matrices (Dokmanic et al., 2015)), a graph, or an orbit of a group action (Bronstein et al., 2017; 2021; Kratsios & Bilokopytov, 2020). Therefore, we ask:
Is exact constraint satisfaction possible with universal deep learning models? ⇤Corresponding authors. 1For example, in a regression problem one can set L(x, y) = kf(x) yk for an unknown function f or in
regression problems one sets L(x, y) = Pm
i=1[C(x)]i log(yi) for an unknown classifier C.
The answer to this question begins by examining the classical universal approximation theorems for deep feedforward networks. If L and K are mildly regular, the universal approximation theorems of Hornik et al. (1989); Cybenko (1989); Pinkus (1999); Gühring et al. (2020); Kidger & Lyons (2020); Park et al. (2021) guarantee that for any “good activation function ” and for every tolerance level ✏ > 0, there is a deep feedforward network with activation function , such that infy2K L(x, y) and L(x, f̂(x)) are uniformly at most ✏ apart. Written in terms of the optimality set,
sup x2X kf̂(x) argmin y2K L(x, y)k ✏, (1)
where the distance of a point y 2 Rm to a set A ✓ Rm is defined by ky Ak , infa2A ky ak. Since argminy2K L(x, y) ✓ K, then (1) only implies that kf̂(x) Kk  ✏ and there is no reason to believe that the constraint f̂(x) 2 K is exactly satisfied, for every x 2 X .
This kind of approximate constraint satisfaction is not always appropriate. In the following examples constraint violation causes either practical or theoretical concerns:
(i) In post-financial crisis risk management, international regulatory bodies mandate that any financial actor should maintain solubility proportional to the risk of their investments (Basel Committee on Banking Supervision, 2015; 2019). To prevent future financial crises, any violation of these risk constraints, no matter the size, incurs large and immediate fines.
(ii) In geometric deep learning, we often need to encode complicated non-vectorial structure present in a dataset, by viewing it as a K valued function (Fletcher, 2013; Bonnabel & Sepulchre, 2009; Baes et al., 2021). However, if K is non-convex then Motzkin (1935) confirms that there is no unique way to map predictions f̂(x) 62 K to a closest point in K. Thus, we are faced with the dilemma: either make an ad-hoc choice of a k in K with k ⇡ f̂(x) (ex.: an arbitrary choice scheme when K = Zm) or have meaningless predictions (ex: non-integer values to integer programs, or symmetry breaking (Weinberg, 1976)2).
Constrained learning was recognized as an effective framework for fairness and robustness by Chamon & Ribeiro (2020) who study empirical risk minimization under constraints. Many emerging topics in machine learning lead to constrained learning formulations. A case in point is modelbased domain generalization (Robey et al., 2021). Despite the importance of (deep) learning with constraints, there are no related approximation-theoretic results to the best of our knowledge.
In this paper, we bridge this theoretical gap by showing that universal approximation with exact constraint satisfaction is always possible for deep (probabilistic) transformer networks with a single attention mechanism as output layer. Our contribution is three-fold:
1. We derive the first universal approximation theorem with exact constraint satisfaction; 2. Our transformer network’s encoder and decoder adapt to the dimension of the constraint
set and thus beat the curse of dimensionality for low-dimensional constraints; 3. Our models leverage a probabilistic attention mechanism that can encode non-convex con-
straints. This probabilistic approach is key to bypass the topological obstructions to nonEuclidean universal approximation (Kratsios & Papon, 2021).
Our analysis provides perspective on the empirical success of attention and adds to the recent line of work on approximation theory for transformer networks, (Yun et al., 2020a;b), which roughly considers the unconstrained case (with K in (1) replaced by Rm) in the special case of L(x, y) = kf(x) yk for a suitable target function f : Rn ! Rm. Our probabilistic perspective on transformer networks fits with the representations of Vuckovic et al. (2021) and of Kratsios (2021).
Our results can be regarded as an approximation-theoretic counterpart to the constrained statistical learning theory of Chamon & Ribeiro (2020). Further, they put forward a perspective on randomness in neural networks that is complementary to the work of Louart et al. (2018); Gonon et al. (2020a;b). We look at the same problem focusing on constraint satisfaction instead of training efficiency. Finally, our proof methods are novel, and build on contemporary tools from metric geometry (Ambrosio & Puglisi, 2020; Bruè et al., 2021).
2As discussed in Rosset et al. (2021) this is problematic since respecting symmetries can often massively reduce the computational burden of a learning task.

heading: 1.1 THE PROBABILISTIC ATTENTION MECHANISM
text: We now give a high-level explanation of our results; the detailed formulations are in Section 2.
Introduced in (Bahdanau et al., 2015) and later used to define the transformer architecture (Vaswani et al., 2017), in the NLP context, attention maps a matrix of queries Q, a matrix of keys K, and a matrix of values V to the quantity Softmax(QK>)V , where the softmax function (defined below) is applied row-wise to QK>. Just as the authors of (Petersen & Voigtlaender, 2020; Zhou, 2020) focus on the simplified versions of practically implementable ConvNets in the study of approximation theory of deep ConvNets (e.g. omitting pooling layers), we find it sufficient to study the following simplified attention mechanism to obtain universal approximation results:
Attention (w, Y ) , SoftmaxN (w)> Y = NX
n=1
[SoftmaxN (w)n]Yn, (2)
where w 2 RN , SoftmaxN : RN 3 w 7! ( e wkPN
j=1 e wj ) N k=1, and Y is an N ⇥m matrix. The attention
mechanism (2) can be interpreted as “paying attention” to a set of particles Y1, . . . , YN 2 Rm defined by Y ’s rows. This simplified form of attention is sufficient to demonstrate that transformer networks can approximate a function while respecting a constraint set, K, whether convex or non-convex. Informal Theorem 1.1 (Deep Maximum Theorem for Transformers). If K is convex and the quantities defining (1) are regular then, for any ✏ 2 (0, 1], there is a feedforward network f̂ , an X✏ ⇢ Rn
of probability 1-✏, and a matrix Y such that the transformer Attention(f̂(x), Y ) satisfies:
(i) Exact Constraint Satisfaction: For each x 2 Rn, Attention(f̂(x), Y )2 K, (ii) Universal Approximation: supx2X✏ kAttention(f̂(x), Y ) argmin
y?2K L(x, y?)k ✏
Informal Theorem 1.1 guarantees that simple transformer networks can minimize any loss function while exactly satisfying the set of convex constraints. As illustrated by Figure 1 and Figure 2, K’s convexity is critical here, since without it the transformer’s prediction may fail to lie in K. This is because any transformer network’s output is a convex combinations of the particles Y1, Y2, Y3; thus, any transformer network’s predictions must belong to these particles’ convex hull.
In Figures 1 and 2, Y ’s columns, i.e. the particles Y1, Y2, and Y3, are each illustrated by a • at the constraint set (K) vertices. The bubble around each each Yi illustrates the predicted probability, for a given input, that f(x) is nearest to that Yi. The ⇥ is the transformer’s prediction which is, by construction, a convex combination of the Yi weighted by the aforementioned probabilities and therefore they lie in the K if it is convex (Figure 1) but not if K is non-convex (Figure 2).
Naturally, we arrive at the question: How can (i) and (ii) simultaneously hold when K is non-convex?
Returning to Vaswani et al. (2017) and using the introduced terminology, we note that the role of the SoftmaxN layer is to rank the importance of the particles {Yn}Nn=1 when optimizing L, at any given input: the weights [SoftmaxN (w)]n in (2) can be interpreted as charging their respective point masses { Yn}Nn=1 with probabilities of being optimal for L (relative to the other particles)3. This suggests the following probabilistic reinterpretation of attention (which we denote by p-attention):
P-attention(w, Y ) , NX
n=1
[SoftmaxN (w)]n Yn . (3)
3Following Villani (2009), Yn is the Borel probability measure on R m assigning full probability to any
Borel subset of Rm containing the particle Yn and 0 otherwise.
Crudely put, P-attention(·, Y ) “pays relative attention to the particles” Y1, . . . , Yn 2 Rm.
A simple computation shows that the mean prediction of our probabilistic attention mechanism, exactly implements “classical” Attention of Vaswani et al. (2017), as defined in (2),
Attention(w, Y ) = EX⇠P-attention(w,Y )[X], (4)
where EX⇠P-attention(w,Y )[X] denotes the (vector-valued) expectation of a random-vector X distributed according to P-attention(w, Y ). Hence, (3) is no less general than (2). The advantage of (3) is that, if each particle Yn belongs to K (even if K is non-convex) then, any sample drawn from the probability measure P-attention(w, Y ) necessarily belongs to K.

heading: 1.2 QUALITATIVE RESULTS: DEEP MAXIMUM THEOREM
text: 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 Theorem4 (Berge, 1963), which states that under mild regularity conditions, given any well-behaved family of input dependent “soft constraint sets” {Cx}x2Rn compatible with K, there is a measurable function mapping each x 2 Rn to a minimizer of L(x, y) on K \ Cx.
We use W1 to denote the Wasserstein-1 distance between probability measures on K. The results also give the flexibility to the user to enforce an input-dependent family of “soft constraints” {Cx}x2Rn which only need to hold approximately; definitions are provided in Section 1.4. Informal Theorem 1.2 (Deep Maximum Theorem: Non-Convex Case). If the quantities defining 1 are regular, K is a compact set of “exact constraints”, and {Cx}x2Rn 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 Rn, P-attention(f̂(x), Y ) is supported in K, (ii) Universal Approximation: P(W1(P-attention(f̂(x), Y ), argmin
y?2Cx\K L(x, y?))  ✏) 1 ✏;
where for a probability measure P on Rm and a B ✓ Rm we define W1(P, B) , infb2B W1(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 Rm, then one recovers the familiar Euclidean distance to the set B via:
W1( y, B) (def) = inf
b2B W1( y, b) = inf b2B ky bk
(def) = ky Bk;
where the first and second equality follows from (Villani, 2009, (5) - page 99), and the last equality is the definition of ky Bk (as in (Aubin & Frankowska, 2009, Definition 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 Rm. In this broad case, we may extract mean predictions from P-attention(f̂ , Y ), by applying the Fréchet mean introduced in Fréchet (1948). Such “geometric means” are well-understood theoretically (Bhattacharya & Patrangenaru, 2003) and easily handled numerically Miolane et al. (2020); Lou et al. (2020).

heading: 1.3 QUANTITATIVE RESULTS: CONSTRAINED UNIVERSAL APPROXIMATION THEOREM
text: In its current form, the objective function L is too general to derive quantitative approximation rates5. Nevertheless, as with most universal approximation theorems (Hornik et al., 1989; Pinkus, 1999; Kidger & Lyons, 2020), if each soft constraint Cx is set to Rm and L quantifies the uniform distance to an unknown continuous function f : Rn ! K in the Euclidean sense,
L(x, y) , kf(x) yk, 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
4More precisely, our result is a deep neural version of the measure-theoretic counterpart to Berge’s Maximum Theorem; see (Aliprantis & Border, 2006, (Measurable Maximum Theorem) - Theorem 18.19).
5For instance, L can describe anything from a regression, to a clustering problem.
K is contained in an unknown d-dimensional submanifold (where d 2 ⇥(m 1s ) 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̂ : Rn ! Rm, 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 Rm (Motzkin, 1935). Nevertheless, if the “projection” were capable of mapping any y 2 Rm 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̂ : Rm ! P1(K) approximates a (non-affine) random projection, in the sense of Ohta (2009); Ambrosio & Puglisi (2020); Bruè et al. (2021); i.e.: a 1-Lipschitz map ⇧ : Rm ! P1(K) satisfying the random projection property: for all y 2 K
⇧y = y.
Thus, ⇧’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 ⇧(y1) is similar to sampling from ⇧(y2) whenever the points y1, y2 2 Rm 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 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,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 Rm). From
Table 1, we see that if m 1s ⌧ m then, s > 0 is large; hence, ✏ms , (1 4✏ 1) 2ms , and N ms are small.

heading: 1.4 NOTATION AND BACKGROUND
text: Optimal Transportd Given any non-empty subset K ✓ Rm, the set of all Borel probability measures P on K with a finite mean; i.e.: EX⇠P[kXk] < 1, is denoted by P1(K). Wasserstein
6P1(K) denotes the Wasserstein space on K, and is defined formally below.
distance W1 is defined for any P,Q 2 P1(K) by the minimal energy needed to transport all mass from P to Q. Following Villani (2009), W1(P,Q) is defined by:
W1(P,Q) , inf ⇡ E(X1,X2)⇠⇡[kX1 X2k],
where the infimum is taken over all Borel probability measures ⇡ on K2 with marginals P and Q. The metric space (P1 (K) ,W1) is named the Wasserstein space over K; we abbreviate it by P1 (K). Smooth Function Spaces The set of real-valued continuous functions on Rn is denoted by C(Rn). 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 : Rn ! Rm extending f ; i.e.: f |X = f , is denoted by C
k tr(X ,K). Our interest in Cktr(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 constraints7. For details on Cktr(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 efficient 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. (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 ? , { t}t2[0,1] ✓ C(R). The set of all such networks is denoted by NN ?n,N and it includes all f̂ : Rn ! RN with iterative representation:
f̂(x) , A(J)x(J), x(j+1)ij , tij ((A (j) x)ij + b (j) ij ), x(0) , x, (5)
where x 2 Rn, j = 1, . . . , J 1, each A(j) is a dj ⇥ dj+1-matrix, each b(j) 2 Rdj+1 , dJ+1 = N , d1 = 0, t1,1, . . . , tJ,NJ 2 [0, 1], for each j. The integer J is f̂ ’s depth and max
j=1,...,J+1 dj is f̂ ’s 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 tij = in (5) is denoted NN n,N .
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 2 N+, and Y be an N ⇥ Q ⇥m-array with Yn,q 2 K for n = 1, . . . , N , q = 1, . . . , Q. Probabilistic attention is the function:
Rn 3 w 7! P-attention(w, Y ) , 1 Q
NX
n=1
QX
q=1
SoftmaxN (w)n Yn,q 2 P1 (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 {Cx}x2Rn of K is said to be a weakly measurable correspondence, denoted C : Rn ◆ Rm, if for every open subset8U ✓ K, {x 2 Rn : Cx \ U 6= ;} is a non-empty Borel subset of Rn (Aliprantis & Border, 2006, pages 557, 592).

heading: 2 MAIN RESULTS
text: We now present our main results in detail. All proofs are relegated to the paper’s appendix.

heading: 2.1 QUALITATIVE APPROXIMATION: DEEP MAXIMUM THEOREM
text: 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
7For example, there is no non-constant continuous function f : [0, 1] ! Z. However, for any 2 (0, 12 ) and any y1, y2 2 Z, f = y1I[0, ] + y2I[ + 12 ,1] belongs to C k tr([0, ] [ [ + 12 , 1],Z) for each k 2 N+.
8Since K is equipped with its subspace topology, then an open subset U of K is any subset of Rm of the form U = U1 \K where U1 is an open subset of Rm (see (Munkres, 2000, Lemma 16.1) for further details).
network. We first 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 defined by K. Assumption 2.1 (Kidger & Lyons (2020)). : R ! R is continuous, is differentiable at some x0 2 R, and its derivative satisfies 0(x0) 6= 0. Theorem 2.2 (Deep Maximum Theorem). Let satisfy Assumption 2.1. Let K ✓ Rn be a nonempty compact set, C : Rn ◆ Rm be a weakly-measurable correspondence with closed values such that Cx \ K 6= ; for each x 2 Rn, L 2 C(Rm), and P be a Borel probability measure on Rn. 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:
F̂ : Rn 3 x 7! P-attention ⇣ f̂(x), Y ⌘ 2 P1(Rm), (6)
satisfies the following:
(i) Exact Constrain Satisfaction: [x2Rn supp(F̂ (x)) ✓ K, (ii) Probably Approximately Optimality: There is a compact X✏ ✓ Rn satisfying:
(a) max x2X✏ W1(F̂ (x), argmin y2Cx\K L(x, y))  ✏,
(b) 1 P(X✏)  ✏. Theorem 2.2 implies that for any random field (Y x)x2Rn on Rm (i.e. a family of Rm-valued random vectors indexed by Rn) 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 argminy2Cx\K L(x, y) (by (ii)).
Corollary 2.3 (F̂ ’s Mean Prediction). Assume the setting of Theorem 2.2. Let {Y x}x2Rn be a K-valued random field with Y x ⇠ F̂ (x) for each x 2 Rn then, 1 P(X✏)  ✏ and
max x2X✏ E[kY x argmin y?2Cx\K L(x, y?)k]  ✏.
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.

heading: 2.2 QUANTITATIVE APPROXIMATION: CONSTRAINED UNIVERSAL APPROXIMATION
text: 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 : Rn ! Rm with f(Rn) ✓ K and a continuous l : [0,1) ! [0,1) with l(0) = 0, satisfying: L(x, y)  l(kf(x) yk). Just as with transformer networks, our “constrained universal approximation theorem” approximates a suitably regular function f : Rn ! K ✓ Rm while exactly respecting the constraints K by implementing an encoder-decoder network architecture. Thus, our model is a composition of an encoder network Ê : Rn ! Rd whose role is to approximate f in a classical “unconstrained fashion” and a decoder network (with probabilistic attention layers at its output) D̂ : Rd ! P1(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. Briefly, 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 Rn to itself with smooth inverse9, such that (K) ✓ Rd; 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-affine and piecewise-linear, and 1 is smooth10 and non-polynomial.
9Here smooth means that is continuously differentiable any number of times. NB, smooth bijections with smooth inverses are called diffeomorphisms in the differential geometry and differential topology literature.
10Following Jost (2017), a function : R ! R is called smooth (or class C1) if @k exists for each k 2 N+.
Theorem 2.7 (Constrained Universal Approximation). Let k 2 N+ and X ✓ [0, 1]n be non-empty. Suppose that satisfies 2.6, L satisfies Assumption 2.4, K ✓ Rn is non-empty, compact and satisfies Assumption 2.5. For any f 2 Cktr(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:
D̂ : Rd 3 x 7! NX
k=1
P-attention ⇣ D̂(x), Y ⌘ 2 P1(K) (7)
where D̂ 2 NN d,N and Y is an N ⇥Q⇥m-array with Y1,1, . . . , YN,Q 2 K such that:
(i) Exact Constrain Satisfaction: For each x 2 Rn: supp(D̂ Ê(x)) ✓ K, (ii) Universal Approximation: The estimate holds 11 :
sup x2[0,1]n W1(D̂ Ê(x), argmin y2K
L(x, y))  ✏K + k Lip( 1)d✏f ;
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 Rd : kz (K)k  ✏f}.
Furthermore, the “complexities” of D̂ and Ê are recorded in Table12 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 quantified 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.7 and let {Y x}x2X be an Rm-valued random field with Y x ⇠ D̂ Ê(x), for each x 2 Rn. Then:
max x2X
EY x⇠D̂ Ê(x) [L(x, Y x)]  l ✏K + k Lip( 1)d✏f .
Corollary 2.8 quantifies the expected performance of a sample from our probabilistic transformer model, as expressed by L, 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.

heading: 2.3 APPLICATIONS
text: 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 (Gühring et al., 2020; Kidger & Lyons, 2020; Shen et al., 2021a), we henceforth consider L(x, y) = kf(x) yk. We also fix f 2 Cktr([0, 1]n,K).

heading: 2.3.1 TRANSFORMERS ARE CONVEX-CONSTRAINED UNIVERSAL APPROXIMATORS
text: We return to the familiar transformer networks of Vaswani et al. (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 x}x2Rn of Corollary 2.3 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) = kf(x) yk. Then:
Rn 3 x 7! E[Y x] = Attention(D̂ Ê(x), Y ) 2 K; (8) (i) Exact Constraint Satisfaction: EY x⇠D̂ Ê(x)[Y
x] 2 K, for each x 2 Rn, (ii) Universal Approximation: sup[0,1]n kf(x) EY 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 . 11In fact, we actually prove that the slightly stronger statement: supx2[0,1]n W1 ⇣ D̂ Ê(x), f(x) ⌘  ✏K +
k Lip( 1)d✏f . Both formulations align when l has a unique minimum at 0, as is the case when L(x, y) = kf(x) yk? and k · k? is any norm on Rm.
12Explicit constants are recorded in Table 2 within the paper’s appendix; there, ✏K and ✏f may differ. 13Explicit constants are recorded in Table 2 within the paper’s appendix; there, ✏K and ✏f may differ.

heading: 2.3.2 CHART-FREE RIEMANNIAN MANIFOLD-VALUED UNIVERSAL APPROXIMATION
text: 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 Rm with distance function by dg . We only require the following mild assumption introduced in Afsari (2011). We recall that the injectivity radius at y0, denoted by infg(y0), (see (Jost, 2017, Definition 1.4.6)) is the minimum length of a geodesic (or minimal length curve) in M with starting point y0. We also recall that the sectional 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 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 satisfies:
1. K is contained in the geodesic ball B(y0, ⇢) , {y 2 M : dg(y0, y) < ⇢} for some point y0 2 M and some radius ⇢ satisfying14: 0 < ⇢ < 2 1 min{injg(y0), ⇡ p C },
2. For each y0, y1 2 K there exists a unique geodesic : [0, 1] ! K joining y0 to y1. Our latent probabilistic representation grants us the flexibility of replacing the usual “extrinsic mean” used in (8) to extract deterministic predictions from our probabilistic transformer networks via an additional Fréchet mean layer at their readout. This intrinsic notion of a mean, was introduced independently in Fréchet (1948) and in Karcher (1977), and is defined on any P 2 P1(K) by:
P̄ , argmin k2K
Z d 2 g(k, u)P(du). (9)
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) = kf(x) yk. If Assumption 2.10 holds then:
Rn 3 x 7! D̂ Ê(x) 2 K, (10) is a well-defined Lipschitz-continuous function, and the following hold:
(i) Exact Constraint Satisfaction: D̂ Ê(x) 2 K, for each x 2 X , (ii) Universal Approximation: supX dg(f(x), D̂ Ê(x)) < ✏K + kd✏f .
The “complexities” of D̂ and Ê are recorded in Table15 1 for ✏2 = ✏k = ✏f .

heading: 3 DISCUSSION
text: In this paper, we derived the first 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 first 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 R3 with outputs constrained to the sphere.
14Following Afsari (2011), we make the convention that if C  0 then, 1p C is interpreted as 1. 15Explicit constants are recorded in Table 2 within the paper’s appendix; there, ✏K and ✏f may differ.

heading: ACKNOWLEDGMENTS
text: Anastasis Kratsios and Ivan Dokmanić 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 Zürich, Valentin Debarnot from the University of Basel for their helpful feedback, and Sven Seuken from the University of Zürich for his helpful feedback in the rebuttal phase.

Review summaries:
The paper under review studies the universal approximation theory with constraints.  For 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.

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.