O N THE L EARNING AND L EARNABILITY OF Q UASIMETRICS
A BSTRACT
Our world is full of asymmetries.
Gravity and wind can make reaching a place easier than coming back.
Social artifacts such as genealogy charts and citation graphs are inherently directed.
In reinforcement learning and control, optimal goal-reaching strategies are rarely reversible (symmetrical).
Distance functions supported on these asymmetrical structures are called quasimetrics .
Despite their common appearance, little research has been done on the learning of quasimetrics.
Our theoretical analysis reveals that a common class of learning algorithms, including unconstrained multilayer perceptrons (MLPs), provably fails to learn a quasimetric consistent with training data.
In contrast, our proposed Poisson
Quasimetric Embedding (PQE) is the ﬁrst quasimetric learning formulation that both is learnable with gradient-based optimization and enjoys strong performance guarantees.
Experiments on random graphs, social graphs, and ofﬂine Q-learning demonstrate its effectiveness over many common baselines.
Project Page: ssnl.github.io/quasimetric .
Code: github.com/SsnL/poisson_quasimetric_embedding .
I NTRODUCTION
Learned symmetrical metrics have been proven useful for innumerable tasks including dimensionality reduction (Tenenbaum et al., 2000), clustering (Xing et al., 2002), classiﬁcation (Weinberger et al., 2006; Hoffer & Ailon, 2015), and information retrieval (Wang et al., 2014).
However, the real world is largely asymmetrical , and symmetrical metrics can only capture a small fraction of it.
Generalizing metrics, quasimetrics (Defn. 2.1) allow for asymmetrical distances and can be found in a wide range of domains (see Fig. 1).
Ubiquitous physical forces, such as gravity and wind, as well as human-deﬁned rules, such as one-way roads, make the traveling time between places a quasimetric.
Furthermore, many of our social artifacts are directed graphs— genealogy charts, follow-relation on
Twitter (Leskovec & Krevl, 2014), citation graphs (Price, 2011), hyperlinks over the Internet, etc.
Shortest paths on these graphs naturally induce quasimetric spaces.
In fact, we can generalize to
Markov Decision Processes (MDPs) and observe that optimal goal-reaching plan costs (i.e., universal value/Q-functions (Schaul et al., 2015; Sutton et al., 2011)) always form a quasimetric (Bertsekas &
Tsitsiklis, 1991; Tian et al., 2020).
Moving onto more abstract structures, quasimetrics can also be found as expected hitting times in Markov chains, and as conditional Shannon entropy H ( · | · ) in information theory.
(See the appendix for proofs and discussions of these quasimetrics.)
In this work, we study the task of quasimetric learning .
Given a sampled training set of pairs and their quasimetric distances, we ask: how well can we learn a quasimetric that ﬁts the training data?
We deﬁne quasimetric learning in analogy to metric learning: whereas metric learning is the problem of learning a metric function, quasimetric learning is the problem of learning a quasimetric function.
This may involve searching over a hypothesis space constrained to only include quasimetric functions (which is what our method does) or it could involve searching for approximately quasimetric functions (we compare to and analyze such approaches).
Successful formulations have many potential applications, such as structural priors in reinforcement learning (Schaul et al., 2015; Tian et al., 2020), graph learning (Rizi et al., 2018) and causal relation learning (Balashankar & Subramanian, 2021).
Towards this goal, our contributions are
• We study the quasimetric learning task with two goals: (1) ﬁtting training data well and (2) respecting quasimetric constraints (Sec. 3);
[Figure] [Figure]
Unconstrained
Functions
[Figure]
• We prove that a large family of algorithms, including unconstrained networks trained in the Neural Tangent Kernel (NTK) regime (Jacot et al., 2018), fail at this task, while a learned embedding into a latent quasimetric space can potentially succeed (Sec. 4);
• We propose Poisson Quasimetric Embeddings (PQEs), the ﬁrst quasimetric embedding formulation learnable with gradient-based optimization that also enjoys strong theoretical guarantees on approximating arbitrary quasimetrics (Sec. 5); • Our experiments complement the theory and demonstrate the beneﬁts of PQEs on random graphs, social graphs and ofﬂine Q-learning (Sec. 6).
P RELIMINARIES ON Q UASIMETRICS AND P OISSON P ROCESSES
Quasimetric space is a generalization of metric space where all requirements of metrics are satisﬁed, except that the distances can be asymmetrical.
Deﬁnition 2.1 (Quasimetric Space).
A quasimetric space is a pair ( X , d ) , where X is a set of points and d : X × X
→ [0 , ∞ ] is the quasimetric, satisfying the following conditions:
[Equation]
Triangle
[Equation]
Being asymmetric, quasimetrics are often thought of as (shortest-path)
distances of some (possibly inﬁnite) weighted directed graph.
A natural way to quantify the complexity of a quasimetric is to consider that of its underlying graph.
Quasimetric treewidth is an instantiation of this idea.
Deﬁnition 2.2 (Treewidth of Quasimetric Spaces (Mémoli et al., 2018)).
Consider a quasimetric space M as shortest-path distances on a positively-weighted directed graph.
Treewidth of M is the minimum over all such graphs’ treewidths.
Poisson processes are commonly used to model events (or points) randomly occurring across a set A (Kingman, 2005) , e.g., raindrops hitting a windshield, photons captured by a camera.
The number of such events within a subset of A is modeled as a Poisson distribution, whose mean is given by a measure µ of A that determines how “frequently the events happen at each location”.
Deﬁnition 2.3 (Poisson Process).
For nonatomic measure µ on set A , a Poisson process on A with mean measure µ is a random countable subset P ⊂ A (i.e., the random events / points) such that
• for any disjoint measurable subsets A 1 , . . .
, A n of A , the random variables N ( A 1 ) , . . .
, N ( A n ) are independent, where N ( B ) (cid:44) # { P ∩ B } is the number of points of P in B , and • N ( B ) has the Poisson distribution with mean µ ( B ) , denoted as Pois( µ ( B ))
Fact 2.4 (Differentiability of P [ N ( A 1 ) ≤ N ( A 2 )] ).
For two measurable subsets A 1 , A 2 ,
[Equation]
Furthermore, for independent X ∼ Pois( µ 1 ) , Y ∼ Pois( µ 2 ) , the probability P [ X ≤ Y ] is differentiable w.r.t.
µ 1 and µ 2 .
In the special case where µ 1 or µ 2 is zero, we can simply compute
[Equation]
always
[Equation] [Equation]
where x + (cid:44) max(0 , x ) .
For general µ 1 , µ 2 , this probability and its gradients can be obtained via a connection to noncentral χ 2 distribution (Johnson, 1959).
We derive the formulas in the appendix.
Therefore, if A 1 and A 2 are parametrized by some θ such that µ ( A 1 \ A 2 ) and µ ( A 2 \ A 1 ) are differentiable w.r.t.
θ , so is P [ N ( A 1 ) ≤ N ( A 2 )] .
[Figure]
Q UASIMETRIC L EARNING
Consider a quasimetric space ( X , d ) .
The quasimetric learning task aims to infer a quasimetric from observing a training set { ( x i , y i , d ( x i , y i ))
} i ⊂ X
× X × [0 , ∞ ] .
Naturally, our goals for a learned predictor ˆ d : X × X
→ R are: respecting the quasimetric constraints and ﬁtting training distances.
Crucially, we are not simply aiming for the usual sense of generalization , i.e., low population error.
Knowing that true distances have a quasimetric structure, we can better evaluate predictors and desire ones that ﬁt the training data and are (approximately) quasimetrics.
These objectives also indirectly capture generalization because a predictor failing either requirement must have large error on some pairs, whose true distances follow quasimetric constraints.
We formalize this relation in Thm. 4.3.
3. L EARNING A LGORITHMS AND H YPOTHESIS S PACES
Ideally, the learning should scale well with data, potentially generalize to unseen samples, and support integration with other deep learning systems (e.g., via differentiation).
Relaxed hypothesis spaces.
One can simply learn a generic function approximator that maps the (concatenated) input pair to a scalar as the prediction of the pair’s distance, or its transformed version (e.g., log distance).
This approach has been adopted in learning graph distances (Rizi et al., 2018) and plan costs in MDPs (Tian et al., 2020).
When the function approximator is a deep neural network, we refer to such methods as unconstrained networks .
While they are known to ﬁt training data well (Jacot et al., 2018), in this paper we also investigate whether they learn to be (approximately) quasimetrics.
Restricted hypothesis spaces.
Alternatively, we can encode each input to a latent space Z , where a latent quasimetric d z gives the distance prediction.
This guarantees learning a quasimetric over data space X .
Often d z is restricted to a subset unable to approximate all quasimetrics, i.e., an overly restricted hypothesis space , such as metric embeddings and the recently proposed DeepNorm and
WideNorm (Pitis et al., 2020).
While our proposed Poisson Quasimetric Embedding (PQE) (speciﬁed in Sec. 5) is also a latent quasimetric, it can approximate arbitrary quasimetrics (and is differentiable).
PQE thus searches in a space that approximates all quasimetrics and only quasimetrics .
3. A T OY E XAMPLE
To build up intuition on how various algorithms perform according to our two goals, we consider a toy quasimetric space with only 3 elements in Fig.
The space has a total of 9 pairs, 8 of which form the training set.
Due to quasimetric requirements (esp.
triangle inequality), knowing distances of these 8 pairs restricts valid values for the heldout pair to a particular range (which is [28 , 31] in this case).
If a model approximates 8 training pairs well and respects quasimetric constraints well, its prediction on that heldout pair should fall into this range.
We train three models w.r.t.
mean squared error (MSE) over the training set using gradient descent:
• Unconstrained deep network that predicts distance, • Metric embedding into a latent Euclidean space with a deep encoder, • Quasimetric embedding into a latent PQE space with a deep encoder (our method from Sec. 5).
The three approaches exhibit interesting qualitative differences.
Euclidean embedding, unable to model asymmetries in training data, fails to attain a low training error.
While both other methods approximate training distances well, unconstrained networks greatly violate quasimetric constraints; only PQEs respect the constraints and consistently predicts within the valid range.
Here, the structural prior of embedding into a quasimetric latent space appears important to successful learning.
Without any such prior, unconstrained networks fail badly.
In the next section, we present a rigorous theoretical study of the quasimetric learning task, which conﬁrms this intuition.
T HEORETICAL A NALYSIS OF V ARIOUS L EARNING A LGORITHMS
In this section, we deﬁne concrete metrics for the two quasimetric learning objectives stated above, and present positive and negative theoretical ﬁndings for various learning algorithms.
Overview.
Our analysis focuses on data-agnostic bounds, which are often of great interests in machine learning (e.g., VC-dimension (Vapnik & Chervonenkis, 2015)).
We prove a strong negative result for a general family of learning algorithms (including unconstrained MLPs trained in NTK regime, k -nearest neighbor, and min-norm linear regression): they can arbitrarily badly fail to ﬁt training data or respect quasimetric constraints (Thm. 4.6).
Our informative construction reveals the core reason of their failure.
Quasimetric embeddings, however, enjoy nice properties as long as they can approximate arbitrary quasimetrics, which motivates searching for “universal quasimetrics”.
The next section presents PQEs as such universal approximators and states their theoretical guarantees.
Assumptions.
We consider quasimetric spaces ( X , d ) with X ⊂ R d , ﬁnite size n = | X | < ∞ , and ﬁnite distances (i.e., d has range [0 , ∞ ) ).
It allows discussing deep networks which can’t handle inﬁnities well.
This mild assumption can be satisﬁed by simply capping max distances in quasimetrics.
For training, m < n 2 pairs are uniformly sampled as training pairs S ⊂ X × X without replacement.
In the appendix, we provide all full proofs, further discussions of our assumptions and presented results, as well as additional results concerning speciﬁc learning algorithms and settings.
4. D ISTORTION AND V IOLATION M ETRICS FOR Q UASIMETRIC L EARNING
We use distortion as a measure of how well the distance is preserved, as is standard in embedding analyses (e.g., Bourgain (1985)).
In this work, we especially consider distortion over a subset of pairs , to quantify how well a predictor ˆ d approximates distances over the training subset S .
Deﬁnition 4.1 (Distortion).
Distortion of ˆ d over a subset of pairs S ⊂ X × X is dis S
( ˆ d ) (cid:44) (cid:0) max ( x,y ) ∈ S,x (cid:54) =
y ˆ d ( x,y ) d ( x,y ) (cid:1)(cid:0) max ( x,y ) ∈ S,x (cid:54) =
y d ( x,y )
ˆ d ( x,y ) (cid:1) , and its overall distortion is dis ( ˆ d ) (cid:44) dis X×X ( ˆ d ) .
For measuring consistency w.r.t.
quasimetric constraints, we deﬁne the (quasimetric) violation metric.
Violation focuses on triangle inequality , which can often be more complex (e.g., in Fig. 2), compared to the relatively simple non-negativity and Identity of Indiscernibles .
Deﬁnition 4.2 (Quasimetric Violation).
Quasimetric violation ( violation for short) of ˆ d is vio ( ˆ d ) (cid:44) max A 1 ,A 2 ,
A 3 ∈X ˆ d ( A 1 ,A 3 ) ˆ d ( A 1 ,A 2 )+ˆ d ( A 2 ,A 3 ) , where we deﬁne
= 1 for notation simplicity.
Both distortion and violation are nicely agnostic to scaling.
Furthermore, assuming non-negativity and Identity of Indiscernibles , vio ( ˆ d ) ≥ 1 always, with equality iff ˆ d is a quasimetric.
Distortion and violation also capture generalization.
Because the true distance d has optimal training distortion (on S ) and violation, a predictor ˆ d that does badly on either must also be far from truth.
Theorem 4.3 (Distortion and Violation Lower-Bound Generalization Error).
For non-negative ˆ d , dis ( ˆ d ) ≥ max( dis S ( ˆ d ) , (cid:112) vio ( ˆ d )) , where dis ( ˆ d ) captures generalization over the entire X space.
4. L EARNING A LGORITHMS E QUIVARIANT TO O RTHOGONAL T RANSFORMS
For quasimetric space ( X , d ) , X ⊂ R d , we consider applying general learning algorithms by concatenating pairs to form inputs ∈ R 2 d (e.g., unconstrained networks).
While straightforward, this approach means that the algorithms are generally unable to relate the same element appearing as 1st or 2nd input .
As we will show, this is sufﬁcient for a wide family of learning algorithms to fail badly– ones equivariant to orthogonal transforms, which we refer to as OrEq algorithms (Defn. 4.4).
For an OrEq algorithm, training on orthogonally transformed data does not affect its prediction, as long as test data is identically transformed.
Many standard learning algorithms are OrEq (Lemma 4.5).
y x z y (cid:48) w (cid:48)
w vio (ˆ d ) ≥ ˆ d ( x,z ) ˆ d ( x,y )
+ ˆ d ( y,z ) ≥ c dis S (ˆ d )( dis S (ˆ d ) + ˆ d ( y,z ))
Training ( ) :
d ( x,z ) =
c, d ( w,z ) = 1 , d ( x,y )
= 1 , d
( y,w (cid:48) )
= 1 .
Test ( ) : ˆ d ( y,z ) = ?
c   ?
y x z y (cid:48) w (cid:48)
w vio (ˆ d ) ≥ ˆ d ( y,z ) ˆ d ( y,w ) + ˆ d ( w,z ) ≥ ˆ d ( y,z ) 2 · dis S (ˆ d )
Training ( ) :
d ( x,z ) =
c, d ( w,z ) = 1 , d ( x,y (cid:48) )
= 1 , d
( y,w ) = 1 .
Test ( ) : ˆ d ( y,z ) = ?
c
With one-hot features, an orthogonal transform can exchange ( ∗ , y ) ↔ ( ∗ , y (cid:48) ) and ( ∗ , w ) ↔ ( ∗ , w (cid:48) ) , leaving the test pair ( y, z ) unchanged, but transforming the training pairs from one scenario to the other.
Given either training set, an OrEq algorithm must attain same training distortion and predict identically on ( y, z ) .
For appropriate c , this implies large distortion or violation in one of these cases.
Deﬁnition 4.4 (Equivariant Learning Algorithms).
Given training set D
=
( z i , y i ) } i ⊂ Z × Y , where z i are inputs and y
i are targets, a learning algorithm
Alg produces a function Alg
( D ): Z → Y such that Alg ( D )( z (cid:48) ) is the function’s prediction on sample z (cid:48) .
Consider T a set of transformations
Z → Z .
Alg is equivariant to T iff for all transform T ∈ T , training set D , Alg ( D ) =
Alg ( T D ) ◦ T , where T D = { ( Tz, y ): ( z, y ) ∈ D} is the training set with transformed inputs.
Lemma 4.5 (Examples of OrEq Algorithms).
k -nearest-neighbor with Euclidean distance, MLP trained with squared loss in NTK regime, and min-norm least-squares linear regression are OrEq .
Failure case.
The algorithms treats the concatenated inputs as generic vectors.
If a transform fundamentally changes the quasimetric structure but is not fully reﬂected in the learned function (e.g., due to equivariance), learning must fail.
The two training sets in Fig.
3 are sampled from two different quasimetrics over the same 6 elements An orthogonal transform links both training sets without affecting the test pair , which is constrained differently in two quasimetrics.
An OrEq algorithm, necessarily predicting the test pair identically seeing either training set, must thus fail on one.
In the appendix, we empirically verify that unconstrained MLPs indeed do fail on this construction.
Extending to larger quasimetric spaces, we consider graphs containing many copies of both patterns in
Fig.
With high probability, our sampled training set fails in the same way—the learning algorithm can not distinguish it from another training set with different quasimetric constraints.
Theorem 4.6 (Failure of OrEq Algorithms).
Let ( f n ) n be an arbitrary sequence of large values.
There is an inﬁnite sequence of quasimetric spaces (( X n , d n )) n with |X n | = n , X n ⊂ R n such that, over the random training set S of size m , any OrEq algorithm must output a predictor ˆ d that satisﬁes
• ˆ d fails non-negativity , or • max( dis S ( ˆ d ) , vio ( ˆ d )) ≥ f n (i.e., ˆ d approximates training S badly or is far
from a quasimetric), with probability 1 / 2 − o (1) , as long as S does not contain almost all pairs 1 − m/n 2 = ω ( n − 1 / 3 ) , and does not only include few pairs m/n 2 = ω ( n − 1 / 2 ) .
Furthermore, standard NTK results show that unconstrained MLPs trained in NTK regime converge to a function with zero training loss.
By the above theorem, the limiting function is not a quasimetric with nontrivial probability.
In the appendix, we formally state this result.
Despite their empirical usages, these results suggest that unconstrained networks are likely not suited for quasimetric learning.
4. Q UASIMETRIC E MBEDDINGS
A quasimetric embedding consists of a mapping f from data space X to a latent quasimetric space ( Z , d z ) , and predicts ˆ d ( x, y ) (cid:44)
d z ( f ( x ) , f ( y )) .
Therefore, they always respect all quasimetric constraints and attain optimal violation of value 1 , regardless of training data .
However, unlike deep networks, their distortion (approximation) properties depend on the speciﬁc latent quasimetrics.
If the latent quasimetric can generally approximate any quasimetric (with ﬂexible learned encoders such as deep networks), we have nice guarantees for both distortion and violation.
In the section below, we present
Poisson Quasimetric Embedding (PQE) as such a latent quasimetric, along with its theoretical distortion and violation guarantees.
P OISSON Q UASIMETRIC E MBEDDINGS (PQE S )
Motivated by above theoretical ﬁndings, we aim to ﬁnd a latent quasimetric space ( R d , d z ) with a deep network encoder f : X → R d , and a quasimetric d z that is both universal and differentiable : • for any data quasimetric ( X , d ) , there exists an encoder f such that d z ( f ( x ) , f ( y ))
≈
d ( x, y ) ; • d z is differentiable (for optimizing f and possible integration with other gradient-based systems).
Notation 5.1.
We use x, y for elements of the data space X , u, v for elements of the latent space R d , upper-case letters for random variables, and ( · ) z for indicating functions in latent space (e.g., d z ).
An existing line of machine learning research learns quasipartitions , or partial orders , via Order
Embeddings (Vendrov et al., 2015).
Quasipartitions are in fact special cases of quasimetrics whose distances are restricted to be binary, denoted as π .
An Order Embedding is a representation of a quasipartition, where π OE ( x, y ) =
(i.e., x is related to y )
iff f ( x ) ≤ f ( y ) coordinate-wise:
[Equation]
Order Embedding is universal and can model any quasipartition (see appendix and Hiraguchi (1951)).
Can we extend this discrete idea to general continuous quasimetrics?
Quite naïvely, one may attempt a straightforward soft modiﬁcation of Order Embedding:
[Equation]
which equals 0 if u ≤ v coordinate-wise, and increases to 1 as some coordinates violate this condition more.
However, it is unclear whether this gives a quasimetric.
A more principled way is to parametrize a (scaled) distribution of latent quasipartitions Π z , whose expectation naturally gives a continuous-valued quasimetric:
[Equation]
Poisson Quasimetric Embedding (PQE) gives a general recipe for constructing such Π z distributions so that d z is universal and differentiable .
Within this framework, we will see that π SoftOE z is actually a quasimetric based on such a distribution and is (almost) sufﬁcient for our needs.
5. D ISTRIBUTIONS OF L ATENT Q UASIPARTITIONS
A random latent quasipartition π z : R d × R d
→ { 0 , 1 } is a difﬁcult object to model, due to complicated quasipartition constraints.
Fortunately, the Order Embedding representation (Eq.
(3)) is without such constraints.
If, instead of ﬁxed latents u, v , we have random latents R ( u ) , R ( v ) , we can compute:
[Equation]
In this view, we represent a random π z via a joint distribution of random vectors 1 { R ( u ) } u ∈ R d , i.e., a stochastic process .
To easily compute the probability of this coordinate-wise event, we assume that each dimension of random vectors is from an independent process, and obtain
[Equation]
The choice of stochastic process is ﬂexible.
Using Poisson processes (with Lebesgue mean measure;
Defn. 2.3) that count random points on half-lines 2 ( −∞ , a ] , we can have R j ( u ) =
N j (( ∞ , u j ]) , the (random) count of events in ( ∞ , u j ] from j -th Poisson process:
[Equation]
where we used Fact 2.4 and the observation that one half-line is either subset or superset of another.
Indeed, π SoftOE z is an expected quasipartition (and thus a quasimetric), and is differentiable .
Considering a mixture of such distributions for expressiveness, the full latent quasimetric formula is
[Equation]
where we slightly abuse notation and consider latents u and v as (reshaped to) 2-dimensional.
We will see that this is a special PQE case with L ebesgue measure and h alf-lines, and thus denoted PQE-LH.
5. G ENERAL PQE F ORMULATION
We can easily generalize the above idea to independent Poisson processes of general mean measures µ j and (sub)set parametrizations u
→ A j ( u ) , and obtain an expected quasipartition as:
[Equation]
which is differentiable as long as the measures and set parametrizations are (after set differences).
Similarly, considering a mixture gives us an expressive latent quasimetric.
A general PQE latent quasimetric is deﬁned with { ( µ i,j , A
i,j ) } i,j and weights α i ≥
as: d PQE z
[Equation]
whose optimizable parameters include { α i } i , possible ones from { ( µ i,j , A i,j ) } i,j (and encoder f ).
This general recipe can be instantiated in many ways.
Setting A i,j ( u ) →
( −∞ , u i,j ] and Lebesgue µ i,j , recovers PQE-LH.
In the appendix, we consider a form with G aussian-based measures and
G aussian-shapes, denoted as PQE-GG.
Unlike PQE-LH, PQE-GG always gives nonzero gradients.
The appendix also includes several implementation techniques that empirically improve stability, including learning α
i ’s with deep linear networks, a formulation that outputs discounted distance, etc.
5. C ONTINUOUS - VALUED S TOCHASTIC P ROCESSES
But why Poisson processes over more common choices such as Gaussian processes?
It turns out that common continuous-value processes fail to give a differentiable formula.
Consider a non-degenerate process { R ( u ) } u , where ( R ( u ) , R ( v )) has bounded density if u (cid:54) = v .
Perturbing u
→
u + δ leaves P
[ R ( u ) =
R ( u + δ )]
= 0 .
Then one of P (cid:2) R ( u ) ≤ R ( u + δ ) (cid:3) and
P (cid:2) R ( u + δ ) ≤ R ( u ) (cid:3) must be far away from 1 (as they sum to 1 ), breaking differentiability at
P [ R ( u ) ≤ R ( u )] = 1 .
(This argument is formalized in the appendix.)
Discrete-valued processes, however, can leave most probability mass on R ( u )
=
R ( u + δ ) and thus remain differentiable.
5. T HEORETICAL G UARANTEES
Our PQEs bear similarity with the algorithmic quasimetric embedding construction in Mémoli et al.
Extending their analysis to PQEs, we obtain the following distortion and violation guarantees.
Theorem 5.2 (Distortion and violation of PQEs).
Under the assumptions of Sec. 4, any quasimetric space with size n and treewidth t admits a PQE-LH and a PQE-GG with distortion O ( t log 2 n ) and violation 1 , with an expressive encoder (e.g., a ReLU network with ≥ 3 layers and polynomial width).
In fact, these guarantees apply to any PQE formulation that satisﬁes a mild condition.
Informally, any
PQE with h × k Poisson processes (i.e., h mixtures) enjoys the above guarantees if it can approximate the discrete counterpart: mixtures of h Order Embeddings, each speciﬁed with k dimensions.
In the appendix, we make this condition precise and provide a full proof of the above theorem.
[Figure]
E XPERIMENTS
Our experiments are designed to (1) conﬁrm our theoretical ﬁndings and (2) compare PQEs against a wider range of baselines, across different types of tasks.
In all experiments, we optimize γ -discounted distances (with γ ∈ { 0 . 9 , 0 . 95 } ), and compare the following ﬁve families of methods:
• PQEs (2 formulations): PQE-LH and PQE-GG with techniques mentioned in Sec. 5.2.
• Unconstrained networks (20 formulations): Predict raw distance (directly, with exp transform, and with ( · ) 2 transform) or γ -discounted distance (directly, and with a sigmoid-transform).
Each variant is run with a possible triangle inequality regularizer E x,y,z (cid:2) max(0 , γ ˆ d ( x,y )+ˆ d ( y,z )
− γ ˆ d ( x,z ) ) 2 (cid:3) for each of 4 weights ∈ { 0 , 0 . 3 , 1 , 3 } .
• Asymmetrical dot products (20 formulations): On input pair ( x, y ) , encode each into a feature vector with a different network, and take the dot product.
Identical to unconstrained networks, the output is used in the same 5 ways, with the same 4 triangle inequality regularizer options.
• Metric encoders (4 formulations): Embed into Euclidean space, (cid:96) 1 space, hypersphere with (scaled) spherical distance, or a mixture of all three.
• DeepNorm (2 formulations) and WideNorm (3 formulations):
Quasimetric embedding methods that often require signiﬁcantly more parameters than PQEs (often on the order of 10 6 ∼ 10more effective parameters; see the appendix for detailed comparisons) but can only approximate a subset of all possible quasimetrics (Pitis et al., 2020).
We show average results from 5 runs.
The appendix provides experimental details, full results (including standard deviations), additional experiments, and ablation studies.
Random directed graphs.
We start with randomly generated directed graphs of 300 nodes, with 64 -dimensional node features given by randomly initialized neural networks.
After training with MSE on discounted distances, we test the models’ prediction error on the unseen pairs (i.e., generalization), measured also by MSE on discounted distances.
On three graphs with distinct structures, PQEs signiﬁcantly outperform baselines across almost all training set sizes (see Fig. 4).
Notably, while
DeepNorm and WideNorm do well on the dense graph quasimetric, they struggle on the other two, attaining both high test MSE (Fig. 4) and train MSE (not shown).
This is consistent with the fact that they can only approximate a subset of all quasimetrics, while PQEs can approximate all quasimetrics.
Large-scale social graph.
We choose the Berkeley - Stanford Web Graph (Leskovec & Krevl, 2014) as the real-wold social graph for evaluation.
This graph consists of 685 , 230 pages as nodes, and 7 , 600 , 595 hyperlinks as directed edges.
We use 128 -dimensional node2vec features (Grover &
Leskovec, 2016) and the landmark method (Rizi et al., 2018) to construct a training set of 2 , 500 ,pairs, and a test set of 150 , 000 pairs.
PQEs generally perform better than other methods, accurately predicting ﬁnite distances while predicting high values for inﬁnite distances (see Table 1).
DeepNorms and WideNorms learn ﬁnite distances less accurately here, and also do much worse than PQEs on learning the (quasi)metric of an undirected social graph (shown in the appendix).
Ofﬂine Q-learning.
Optimal goal-reaching plan costs in MDPs are quasimetrics (Bertsekas &
Tsitsiklis, 1991; Tian et al., 2020) (see also the appendix).
In practice, optimizing deep Q-functions often suffers from stability and sample efﬁciency issues (Henderson et al., 2018; Fujimoto et al., 2018).
As a proof of concept, we use PQEs as goal-conditional Q-functions in ofﬂine Q-learning, on the grid-world environment with one-way doors built upon gym-minigrid (Chevalier-Boisvert et al., 2018)
(see Fig.
1 right), following the algorithm and data sampling procedure described in Tian et al.
Adding strong quasimetric structures greatly improves sample efﬁciency and greedy planning success rates over popular existing approaches such as unconstrained networks used in Tian et al.
and asymmetrical dot products used in Schaul et al.
(2015) (see Fig. 5).
As an interesting observation, some metric embedding formulations work comparably well.
[Figure] [Figure]
R ELATED W ORK
Metric learning.
Metric learning aims to approximate a target metric/similarity function, often via a learned embedding into a metric space.
This idea has successful applications in dimensionality reduction (Tenenbaum et al., 2000), information retrieval (Wang et al., 2014), clustering (Xing et al., 2002), classiﬁcation (Weinberger et al., 2006; Hoffer & Ailon, 2015), etc.
While asymmetrical formulations have been explored, they either ignore quasimetric constraints (Oord et al., 2018;
Logeswaran & Lee, 2018; Schaul et al., 2015), or are not general enough to approximate arbitrary quasimetric (Balashankar & Subramanian, 2021), which is the focus of the present paper.
Isometric embeddings
Isometric (distance-preserving) embeddings is a highly inﬂuential and well-studied topic in mathematics and statistics.
Fundamental results, such as Bourgain’s random embedding theorem (Bourgain, 1985), laid important ground work in understanding and constructing (approximately) isometric embeddings.
While most such researches concern metric spaces, Mémoli et al.
(2018) study an algorithmic construction of a quasimetric embedding via basic blocks called quasipartitions .
Their approach requires knowledge of quasimetric distances between all pairs and thus is not suitable for learning.
Our formulation takes inspiration from the form of their embedding, but is fully learnable with gradient-based optimization over a training subset.
Quasimetrics and partial orders.
Partial orders (quasipartitions) are special cases of quasimetrics (see Sec. 5).
A line of machine learning research studies embedding partial order structures into latent spaces for tasks such as relation discovery and information retrieval (Vendrov et al., 2015; Suzuki et al., 2019;
Hata et al., 2020; Ganea et al., 2018).
Unfortunately, unlike PQEs, such formulations do not straightforwardly generalize to arbitrary quasimetrics, which are more than binary relations.
Similar to PQEs, DeepNorm and WideNorm are quasimetric embedding approaches learnable with gradient-based optimization (Pitis et al., 2020).
Theoreically, they universally approximates a subset of quasimetrics (ones induced by asymmetrical norms).
Despite often using many more parameters, they are restricted to this subset and unable to approximate general quasimetrics like PQEs do (Fig. 4).
I MPLICATIONS
In this work, we study quasimetric learning via both theoretical analysis and empirical evaluations.
Theoretically, we show strong negative results for a common family of learning algorithms, and positive guarantees for our proposed Poisson Quasimetric Embedding (PQE).
Our results introduce the novel concept of equivariant learning algorithms, which may potentially be used for other learnability analyses with algorithms such as deep neural networks.
Additionally, a thorough average-case or data-dependent analysis would nicely complement our results, and may shed light on conditions where algorithms like deep networks can learn decent approximations to quasimetrics in practice.
PQEs are the ﬁrst quasimetric embedding formulation that can be learned via gradient-based optimization.
Empirically, PQEs show promising performance in various tasks.
Furthermore, PQEs are fully differentiable, and (implicitly) enforce a quasimetric structure in any latent space.
They are particularly suited for integration in large deep learning systems, as we explore in the Q-learning experiments.
This can potentially open the gate to many practical applications such as better embedding for planning with MDPs, efﬁcient shortest path ﬁnding via learned quasimetric heuristics, representation learning with quasimetric similarities, causal relation learning, etc.
A PPENDIX A Discussions for Sec. 2: Preliminaries on Quasimetrics and Poisson Processes
B Proofs, Discussions and Additional Results for Sec. 4: Theoretical Analysis of Various Learning Algorithms
D Experiment Settings and Additional Results
A D ISCUSSIONS FOR S EC . 2: P RELIMINARIES ON Q UASIMETRICS AND P OISSON P ROCESSES A. Q UASIMETRIC S PACES
Deﬁnition 2.1 (Quasimetric Space).
A quasimetric space is a pair ( X , d ) , where X is a set of points and d : X × X
→ [0 , ∞ ] is the quasimetric, satisfying the following conditions:
[Equation]
of
[Equation]
Triangle Inequality
[Equation]
Deﬁnition A.1 (Quasipseudometric Space).
As a further generalization, we say ( X , d ) is a quasipseudometric space if the Identity of Indiscernibles requirement is only satisﬁed in one direction:
[Equation]
(Identity of
[Equation] [Equation]
(Triangle Inequality)
A.1. E XAMPLES OF Q UASIMETRIC S PACES
Proposition A.2 (Expected Hitting Time of a Markov Chain).
Let random variables ( X t ) t be a
Markov Chain with support X .
Then ( X , d hitting ) is a quasimetric space, where
[Equation]
[Equation] [Equation]
14) where we deﬁne the hitting time of s starting from s to be 0 .
Proof of Proposition A.2.
Obviously d hitting is non-negative.
We then verify the following quasimetric space properties:
• Identity of Indiscernibles.
By deﬁnition, we have, ∀ x, y ∈ X , x
[Equation]
• Triangle Inequality.
For any x, y, z ∈ X , we have
[Equation]
Hence, ( X , d hitting ) is a quasimetric space.
Proposition A.3 (Conditional Shannon Entropy).
Let X be the set of random variables (of some probability space).
Then ( X , d H ) is a quasipseudometric space, where
[Equation]
If for all distinct ( X, Y ) ∈
X × X , X can not be written as (almost surely) a deterministic function of Y , then ( X , d H ) is a quasimetric space.
Proof of Proposition A.3.
Obviously d H is non-negative.
We then verify the following quasipseudometric space properties:
• Identity of Indiscernibles.
By deﬁnition, we have, ∀ X, Y ∈ X ,
[Equation]
where ≤ is =
iff Y is a (almost surely) deterministic function of X .
• Triangle Inequality.
For any X, Y, Z ∈ X , we have
[Equation]
Hence, ( X , d H ) is a quasipseudometric space, and a quasimetric space when the last condition is satisﬁed.
Conditional Kolmogorov Complexity From algorithmic information theory, the conditional Kolmogorov complexity K ( y | x ) also similarly measures “the bits needed to create y given x as input” (Kolmogorov, 1963).
It is also almost a quasimetric, but the exact deﬁnition affects some constant/log terms that may make the quasimetric constraints non-exact.
For instance, when deﬁned with the preﬁx-free version, conditional Kolmogorov complexity is always strictly positive, even for K ( x | x ) >
(Li et al., 2008).
One may remedy this with a deﬁnition using a universal Turing machine (UTM) that simply outputs the input on empty program.
But to make triangle inequality work, one needs to reason about how the input and output parts work on the tape(s) of the UTM.
Nonetheless, regardless of the deﬁnition details, conditional Kolmogorov complexity do satisfy a triangle inequality up to log terms (Grunwald & Vitányi, 2004).
So intuitively, it behaves roughly like a quasimetric deﬁned on the space of binary strings.
Optimal Goal-Reaching Plan Costs in Markov Decision Processes (MDPs)
We deﬁne MDPs in the standard manner:
M = ( S , A , R , P , γ ) (Puterman, 1994), where S is the state space,
A is the action space, R : S ×
A → R is the reward function, P : S ×
A → ∆( S ) is the transition function (where ∆( S ) is the set of all distributions over S ), and γ ∈ (0 , 1) is the discount factor.
We deﬁne Π as the collection of all stationary policies π : S → ∆( A ) on M .
For a particular policy π ∈ Π , it induces random trajectories :
• Trajectory starting from state s ∈ S is the random variable
[Equation]
distributed as
[Equation]
• Trajectory starting from state-action pair ( s,
a ) ∈ S
× A is the random variable
[Equation]
distributed as
[Equation] [Equation]
Proposition A.
(Optimal Goal-Reaching Plan Costs in MDPs).
Consider an MDP M =
( S , A , R , P , γ ) .
WLOG
[Equation]
assume that R : S × A → ( −∞ , 0] has only non-positive rewards (i
[Equation]
e., negated costs).
Let X
= S ∪ ( S × A ) .
Then ( X
[Equation]
d sum ) and ( X
[Equation]
d γ ) are quasipseudometric spaces, where
[Equation] [Equation]
not reached
[Equation]
and performed a (cid:48) yet
[Equation]
goal is a state-action pair
[Equation]
and
[Equation] [Equation]
is deﬁned similarly.
If the reward function is always negative , ( X , d sum ) and ( X , d γ ) are quasimetric spaces.
Proof of Proposition A.4.
Obviously both d sum and d γ are non-negative, and satisfy Identity of
Indiscernibles (for quasipseudometric spaces).
For triangle inequality, note that for each y , we can instead consider alternative MDPs:
• If y = s (cid:48) ∈ S , modify the original MDP to make s (cid:48) a sink state, where performing any action yields 0 reward (i.e., 0 cost); • If y =
( s (cid:48) , a (cid:48) )
∈ S × A , modify the original MDP such that performing action a (cid:48) in state s (cid:48) surely transitions to a new sink state
where
performing
any action yields 0 reward (i
e., 0 cost).
Obviously, both are Markovian.
Furthermore, they are Stochastic Shortest Path problems with no negative costs (Guillot & Stauffer, 2020), implying that there are Markovian (i.e., stationary) optimal policies (respectively w.r.t.
either minimizing expected total cost or maximizing expected γ total cost ).
Thus optimizing over the set of stationary policies, Π , gives the optimal quantity over all possible policies, including concatenation of two stationary policies.
Thus the triangle inequality is satisﬁed by both.
Hence, ( X , d sum ) and ( X , d γ ) are quasipseudometric spaces.
Finally, if the reward function is always negative , x (cid:54) =
y =
⇒ d sum ( x, y ) > 0 and d γ ( x, y ) > 0 , so ( X , d sum ) and ( X , d γ ) are quasimetric spaces.
Remark A.5.
We make a couple remarks:
• Any MDP with a bounded reward function can be modiﬁed to have only non-positive rewards by subtracting the maximum reward (or larger); • We have
[Equation]
• When the dynamics is deterministic, d sum ≡
[Equation]
• Unless y is reachable from x with probability 1 under some policy, d sum ( x, y )
= ∞ .
• Unless y is unreachable from x with probability 1 under all policies, d sum ( x, y ) < ∞ .
Therefore, it is often favorable to consider d γ types.
• In certain MDP formulations, the reward is stochastic and/or dependent on the reached next state.
The above deﬁnitions readily extend to those cases
• γ d γ (( s,a ) ,y ) is very similar to Q-functions except that Q-function applies discount based on time, and γ d γ (( s,a ) ,y ) applies discount based on costs.
We note that a Q-learning-like recurrence can also be found for γ d γ (( s,a ) ,y ) .
If the cost is constant in the sense for some ﬁxed c < 0 , R ( s, a )
=
c , ∀ ( s, a ) ∈ S
×
A
then time and cost are equivalent up to a scale.
Therefore, γ d γ (( s,a ) ,y ) coincides with the optimal Q-functions for the MDPs described in proof, and γ d γ ( s,y ) coincides with the optimal value functions for the respective MDPs.
A.1. Q UASIMETRIC T REEWIDTH AND G RAPH T REEWIDTH
Deﬁnition 2.2 (Treewidth of Quasimetric Spaces (Mémoli et al., 2018)).
Consider representations of a quasimetric space M as shortest-path distances on a positively-weighted directed graph.
Treewidth of M is the minimum over all such graphs’ treewidths.
(Recall that the treewidth of a graph (after replacing directed edges with undirected ones) is a measure of its complexity.)
Graph treewidth is a standard complexity measure of how “similar” a graph is to a tree (Robertson & Seymour, 1984).
Informally speaking, if a graph has low treewidth, we can represent it as a tree, preserving all connected paths between vertices, except that in each tree node, we store a small number of vertices (from the original graph) rather than just 1 .
Graph treewidth is widely used by the Theoretical Computer Science and Graph Theory communities, since many NP problems are solvable in polynomial time for graphs with bounded treewidth (Bertele & Brioschi, 1973).
A. P OISSON P ROCESSES
Deﬁnition 2.3 (Poisson Process).
For nonatomic measure µ on set A , a Poisson process on A with mean measure µ is a random countable subset P ⊂ A (i.e., the random events / points) such that
• for any disjoint measurable subsets A 1 , . . .
, A n of A , the random variables N ( A 1 ) , . . .
, N ( A n ) are independent, where N ( B ) (cid:44) # { P ∩ B } is the number of points of P in B , and • N ( B ) has the Poisson distribution with mean µ ( B ) , denoted as Pois( µ ( B )) .
Poisson processes are usually used to model events that randomly happens “with no clear pattern”, e.g., visible stars in a patch of the sky, arrival times of Internet packages to a data center.
These events may randomly happen all over the sky / time.
To an extent, we can say that their characteristic feature is a property of statistical independence (Kingman, 2005).
To understand this, imagine raindrops hitting the windshield of a car.
Suppose that we already know that the rain is heavy, knowing the exact pattern of the raindrops hitting on the left side of the windshield tells you little about the hitting pattern on the right side.
Then, we may assume that, as long as we look at regions that are disjoint on the windshield, the number of raindrops in each region are independent.
This is the fundamental motivation of Poisson processes.
In a sense, from this characterization,
Poisson processes are inevitable (see Sec.
1.4 of (Kingman, 2005)).
A.2. P OISSON R ACE P ROBABILITY P [Pois( µ 1 ) ≤ Pois( µ 2 )] AND I TS G RADIENT F ORMULAS
In Fact 2.4 we made several remarks on the Poisson race probability, i.e., for independent X ∼
Pois( µ 1 ) , Y ∼ Pois( µ 2 ) , the quantity P [ X ≤ Y ] .
In this section, we detailedly describe how we arrived at those conclusions, and provide the exact gradient formulas for differentiating P [ X ≤ Y ] w.r.t.
µ 1 and µ 2 .
From Skellam distribution CDF to Non-Central χ 2 distribution CDF.
Distribution of the difference of two independent Poisson random variables is called the Skellam distribution (Skellam, 1946), with its parameter being the rate of the two Poissons.
That is, X − Y ∼ Skellam( µ 1 , µ 2 ) .
Therefore,
P [ X ≤ Y ] is essentially the cumulative distribution function (CDF) of this Skellam at 0 .
In Eq.
(4) of (Johnson, 1959), a connection is made between the CDF of Skellam( µ 1 , µ 2 ) distribution, and the
CDF of a non-central χ 2 distribution (which is a non-centered generalization of χ 2 distribution) with two parameters k > 0 degree(s) of freedom and non-centrality parameter λ ≥ 0 ): for integer n > 0 ,
[Equation]
which can be evaluated using statistical computing packages such as SciPy (Virtanen et al., 2020) and CDFLIB (Burkardt, 2021; Brown et al., 1994).
Marcum-Q-Function and gradient formulas.
To differentiate through Eq.
(41), we consider representing the non-central χ 2 CDF as a Marcum-Q-function (Marcum, 1950).
One deﬁnition of the
Marcum-Q-function Q M : R × R → R in statistics is
[Equation]
where I M − 1 is the modiﬁed Bessel function of order M − 1 .
(When M is non-integer, we refer readers to (Brychkov, 2012; Marcum, 1950) for deﬁnitions, which are not relevant to the discussion below.)
When used in CDF of non-central χ 2 , we have
[Equation]
Combining with Eq.
(41), and using the symmetry Skellam( µ 1 , µ 2 ) d = − Skellam( µ 2 , µ 1 ) , we have, for integer n ,
[Equation]
Prior work (Brychkov, 2012) provides several derivative formula for the Marcum-Q-Function: • For n < 0 , we have
[Equation]
where I ( e ) v ( x ) (cid:44) e −| x | I v ( x ) is the exponentially-scaled version of I v that computing libraries often provide due to its superior numerical precision (e.g., SciPy (Virtanen et al., 2020)),
[Equation]
• For n ≥ 0 , we have
[Equation]
52) and, ∂
[Equation]
Setting n = 0 gives the proper forward and backward formulas for P [ X ≤ Y ] .
B P ROOFS , D ISCUSSIONS AND A DDITIONAL R ESULTS FOR S EC . 4: T HEORETICAL A NALYSIS OF V ARIOUS L EARNING A LGORITHMS
Assumptions.
Recall that we assumed a quasimetric space, which is stronger than a quasipseudometric space (Defn.
A.1), with ﬁnite distances.
These are rather mild assumptions, since any quasipseudometric with inﬁnities can always be modiﬁed to obey these assumptions by (1) adding a small metric (e.g., d (cid:15) ( x, y ) (cid:44) (cid:15)
1 x (cid:54) =
y with small (cid:15) > 0 ) and (2) capping the inﬁnite distances to a large value higher than any ﬁnite distance.
Worst-case analysis.
In this work we focus on the worst-case scenario, as is common in standard (quasi)metric embedding analyses (Bourgain, 1985; Johnson & Lindenstrauss, 1984; Indyk, 2001;
Mémoli et al., 2018).
Such results are important because embeddings are often used as heuristics in downstream tasks (e.g., planning) which are sensitive to any error.
While our negative result readily extends to the average-case scenario (since the error (distortion or violation) is arbitrary), we leave a thorough average-case analysis as future work.
Data-independent bounds.
We analyze possible data-independent bounds for various algorithms.
In this sense, the positive result for PQEs (Thm.
C.4) is really strong, showing good guarantees regardless data quasimetric .
The negative result (Thm. 4.6) is also revealing, indicating that a family of algorithms should probably not be used, unless we know something more about data.
Dataindependent bounds are often of great interest in machine learning (e.g., concepts of VC-dimension (Vapnik & Chervonenkis, 2015) and PAC learning (Valiant, 1984)).
An important future work is to explore data-dependent results, possibly via deﬁning a quasimetric complexity metric that is both friendly for machine learning analysis, and connects well with combinatorics measures such as quasimetric treewidth.
Violation and distortion metrics.
The optimal violation has value 1 .
Speciﬁcally, it is 1 iff ˆ d is a quasimetric on X (assuming non-negativity ).
Distortion (over training set) and violation together quantify how well ˆ d learns a quasimetric consistent with the training data.
A predictor can ﬁt training data well (low distortion), but ignores basic quasimetric constraints on heldout data (high violation).
Conversely, a predictor can perfectly obey the training data constraints (low violation), but doesn’t actually ﬁt training data well (high distortion).
Indeed, (assuming non-negativity and Identity of
Indiscernibles ), perfect distortion (value 1 ) and violation (value 1 ) imply that ˆ d is a quasimetric consistent with training data.
Relation with classical in-distribution generalization studies.
Classical generalization studies the prediction error over the underlying data distribution, and often involves complexity of the hypothesis class and/or training data (Vapnik & Chervonenkis, 2015; McAllester, 1999).
Our focus on quasimetric constraints violation is, in fact, not an orthogonal problem, but potentially a core part of in-distribution generalization for this setting.
Here, the underlying distribution is supported on all pairs of X × X .
Indeed, if a learning algorithm has large distortion, it must attain large prediction error on S ⊂ X
× X ; if it has large violation, it must violates the quasimetric constraints and necessarily admits bad prediction error on some pairs (whose true distances obey the quasimetric constraints).
Thm. 4.3 (proved below) formalizes this idea, where we characterize generalization with the distortion over all possible pairs in X × X .
B. T HM . 4.3: D ISTORTION AND V IOLATION L OWER -B OUND G ENERALIZATION E RROR
Theorem 4.3 (Distortion and Violation Lower-Bound Generalization Error).
For non-negative ˆ d , dis ( ˆ d ) ≥ max( dis S ( ˆ d ) , (cid:112) vio ( ˆ d )) , where dis ( ˆ d ) captures generalization over the entire X space.
B.1. P ROOF
Proof of Thm.
It is obvious that
[Equation] [Equation]
Therefore, it remains to show that dis ( ˆ d ) ≥ (cid:113) vio ( ˆ d ) .
WLOG, say vio ( ˆ d ) > 1 .
Otherwise, the statement is trivially true.
By the deﬁnition of violation (see Defn. 4.2), we have, for some x, y, z ∈ X , with ˆ d ( x, z ) > 0 , ˆ d ( x, z ) ˆ d ( x,
y
[Equation] [Equation]
If ˆ d ( x, y ) + ˆ d ( y, z ) = 0 , then we must have one of the following two cases:
• If d ( x, y ) > 0 or d ( y, z ) > 0 , the statement is true because dis ( ˆ d ) = ∞ .
• If d ( x, y ) =
d ( y, z ) = 0 , then d ( x, z ) =
and the statement is true since dis ( ˆ d ) ≥ ˆ d ( x,z ) d ( x,z ) = ∞ .
It is sufﬁcient to prove the case that ˆ d ( x, y )
+ ˆ d ( y, z ) > 0 .
We can derive
[Equation] [Equation]
If d ( x, z ) = 0 , then dis ( ˆ d ) = ∞ and the statement is trivially true.
If d ( x, z ) > 0 , above Eq.
(59) implies
[Equation]
Combining Eqs.
(55) and (60) gives the desired statement.
B. L EMMA 4.5: E XAMPLES OF O R E Q A LGORITHMS
Lemma 4.5 (Examples of OrEq Algorithms).
k -nearest-neighbor with Euclidean distance, MLP trained with squared loss in NTK regime, and min-norm least-squares linear regression are OrEq .
Recall the deﬁnition of Equivariant Learning Transforms.
Deﬁnition 4.4 (Equivariant Learning Algorithms).
Given training set D
=
( z i , y i ) } i ⊂ Z × Y , where z i are inputs and y
i are targets, a learning algorithm
Alg produces a function Alg
( D ): Z → Y such that Alg ( D )( z (cid:48) ) is the function’s prediction on sample z (cid:48) .
Consider T a set of transformations
Z → Z .
Alg is equivariant to T iff for all transform T ∈ T , training set D , Alg ( D ) =
Alg ( T D ) ◦ T , where T D = { ( Tz, y ): ( z, y ) ∈ D} is the training set with transformed inputs.
B.2. P ROOF
Proof of Lemma 4.5.
We consider the three algorithms individually:
• k -nearest neighbor with Euclidean distance.
It is evident that if a learning algorithm only depend on pairwise dot products (or distances), it is
equivariant to
orthogonal transforms, which
preserve dot products
and distances)
k -nearestneighbor with Euclidean distance only depends on pairwise distances, which can be written in terms of dot products
[Equation]
Therefore, it is equivariant to orthogonal transforms.
• Min-norm least-squares linear regression.
Recall that the solution to min-norm least-squares linear regression Ax = b is given by
Moore–Penrose pseudo-inverse x =
A + b .
For any matrix A ∈ R m × n with SVD U Σ V ∗ =
A , and T ∈ O ( n ) (where O ( n ) is the orthogonal group in dimension n ), we have
[Equation]
where we used T ∗
= T T for T ∈ O ( n ) .
The solution for the transformed data AT T and b is thus
[Equation]
Thus, for any new data point ˜ x ∈ R n and its transformed version T ˜ x ∈ R n ,
[Equation]
Hence, min-norm least-squares linear regression is equivariant to orthogonal transforms.
• MLP trained with squared loss in NTK regime.
We ﬁrst recall the NTK recursive formula from (Jacot et al., 2018).
Denote the NTK for a MLP with L layers with the scalar kernel Θ ( L ) : R d × R d → R .
Let β > 0 be the (ﬁxed) parameter for the bias strength in the network model, and σ be the activation function.
Given x, z ∈ R d , it can be recursively deﬁned as following.
For h ∈ [ L ] ,
[Equation] [Equation]
where
[Equation]
for some constant c .
It is evident from the recursive formula, that Θ ( h ) ( x, z ) only depends on x T x , z T z and x T z .
Therefore, the NTK is invariant to orthogonal transforms.
Furthermore, training an MLP in NTK regime is the same as kernel regression with the NTK (Jacot et al., 2018), which has a unique solution only depending on the kernel matrix on training set, denoted as K train ∈ R n × n , where n is the training set size.
Speciﬁcally, for training data { ( x i , y i ) } i ∈ [ n ] , the solution f ∗ NTK : R → R can be written as
[Equation]
where y =
( yy. . .
y n ) is the vector of training labels.
Consider any orthogonal transform T ∈ O ( d ) , and the NTK regression trained on the transformed data { ( Tx i , y i ) } i ∈ [ n ] .
Denote the solution as f ∗ NTK ,T : R → R .
As we have shown, K − 1 train is invariant to such transforms, and remains the same.
Therefore,
[Equation]
Hence, MLPs trained (with squared loss) in NTK regime is equivariant to orthogonal transforms.
Furthermore, we note that there are many variants of MLP NTK formulas depending on details such as the particular initialization scheme and bias settings.
However, they usually only lead to slight changes that do not affect our results.
For example, while the above recursive NTK formula are derived assuming that the bias terms are initialized with a normal distribution (Jacot et al., 2018), the formulas for initializing bias as zeros (Geifman et al., 2020) does not affect the dependency only on dot product, and thus our results still hold true.
These cases conclude the proof.
y x z y (cid:48) w (cid:48)
w vio (ˆ d ) ≥ ˆ d ( x,z ) ˆ d ( x,y )
+ ˆ d ( y,z ) ≥ c dis S (ˆ d )( dis S (ˆ d ) + ˆ d ( y,z ))
Training ( ) :
d ( x,z ) =
c, d ( w,z ) = 1 , d ( x,y )
= 1 , d
( y,w (cid:48) )
= 1 .
Test ( ) : ˆ d ( y,z ) = ?
c   ?
y x z y (cid:48) w (cid:48)
w vio (ˆ d ) ≥ ˆ d ( y,z ) ˆ d ( y,w ) + ˆ d ( w,z ) ≥ ˆ d ( y,z ) 2 · dis S (ˆ d )
Training ( ) :
d ( x,z ) =
c, d ( w,z ) = 1 , d ( x,y (cid:48) )
= 1 , d
( y,w ) = 1 .
Test ( ) : ˆ d ( y,z ) = ?
c
With transform can
(cid:48) leaving the test pair ( y, z ) unchanged, but transforming the training pairs from one scenario to the other.
Given either training set, an OrEq algorithm must attain same training distortion and predict identically on ( y, z ) .
For appropriate c , this implies large distortion or violation in one of these cases.
B.
T HM . 4.6: F AILURE OF O R E Q A LGORITHMS
We start with a more precise statement of Thm. 4.6 that takes into consideration divergent m/n 2 :
Theorem 4.6 (Failure of OrEq Algorithms).
Let ( f n ) n be an arbitrary sequence of large values.
There is an inﬁnite sequence of quasimetric spaces (( X n , d n )) n with |X n | = n , X n ⊂ R n such that, over the random training set S of size m , any OrEq algorithm must output a predictor ˆ d that satisﬁes
• ˆ d fails non-negativity , or • max( dis S ( ˆ d ) , vio ( ˆ d )) ≥ f n (i
e., ˆ d approximates training S badly or is far from a quasimetric), with
probability 1 / 2 − o (1) , as long as S does not contain almost all pairs 1 − m/n 2 = ω (
n −
, and does not only include few pairs m/n 2 = ω ( n − 1 / 2 ) .
Recall that the little-Omega notation means f = ω ( g ) ⇐⇒ g
= o ( f ) .
B.3. P ROOF
Proof strategy.
In our proof below, we will extend the construction discussed in Sec.
4.2 to large quasimetric spaces (reproduced here as Fig. 6).
To do so,
we
Construct large quasimetric spaces containing many copies of the (potentially failing) structure in Fig.
6, where we can consider training sets of certain properties such that • we can pair up such training sets, • an algorithm equivariant to orthogonal transforms must fail on one of them, • for each pair, the two training sets has equal probability of being sampled; Then, it remains to show that with probability 1 − o (1) we end up with a training set of such properties.
Consider sampling training set as individually collecting each pair with a certain probability p , and carefully analyze the conditions to sample a training set with the special properties with high probability 1 − o (1) .
Extend to ﬁxed-size training sets and show that, under similar conditions, we sample a training set with the special properties with high probability 1 − o (1) .
In the discussion below and the proof, we will freely speak of inﬁnite distances between two elements of X , but really mean a very large value (possibly ﬁnite).
This allows us to make the argument clearer and less verbose.
Therefore, we are not restricting the applicable settings of Thm.
4.6 to quasimetrics with (or without) inﬁnite distances.
In Sec.
4.2, we showed how orthogonal-transform-equivariant algorithms can not predict ˆ d ( y, z ) differently for the two particular quasimetric spaces and their training sets shown in Fig. 6.
But are these the only bad training sets?
Before the proof, let us consider what kinds of training sets are bad for these two quasimetric spaces.
Consider the quasimetrics d left and d right over X (cid:44)
{ x, y, y (cid:48) , z, w, w (cid:48) } , with distances as shown in the left and right parts of Fig.
6, where we assume that the unlabeled pairs have inﬁnite distances except in the left pattern d ( x, w (cid:48) ) ≤ 2 , and in the both patterns d ( y, z ) has some appropriate value consistent with the respective triangle inequality.
Speciﬁcally, we ask:
• For what training sets S left ⊂ X
×
X
can we interchange y ↔ y (cid:48) and w ↔ w (cid:48) on 2nd input to obtain a valid training set for d right , regardless of c ?
• For what training sets S right ⊂
X
× X can we interchange y ↔ y (cid:48) and w ↔ w (cid:48) on 2nd input to obtain a valid training set for d left , regardless of c ?
Note that if S left (or S right ) satisﬁes its condition, the predictor ˆ d from an algorithm equivariant to orthogonal transforms must (1) predict ˆ d ( y, z ) identically and (2) attain the same training set distortion on it and its transformed training set.
As we will see in the proof for Thm. 4.6, this implies large distortion or violation for appropriate c .
Intuitively, all we need is that the transformed data do not break quasimetric constraints.
However, its conditions are actually nontrivial as we want to set c to arbitrary:
• We can’t have ( x, w ) ∈ S
right because it would be transformed into ( x, w ) which has d left ( x, w ) ≤ 2 .
Then d right ( x, w ) ≤ 2 and then restricts the possible values of c due to triangle inequality with d right ( w, z ) = 1 .
For similar reasons, we can’t have ( x, w (cid:48) )
∈ S left .
In fact, we can’t have a path of ﬁnite total distance from x to w (or w (cid:48) ) in S right (or S left ).
• We can not have ( y (cid:48) , y (cid:48) )
∈ S ( · ) (which has distance 0 ), which would get transformed into ( y (cid:48) , y ) with distance 0 , which (on the right pattern) would restrict the possible values of c due to triangle inequality, and break our assumption of d ( · ) not being a quasipseudometric.
For similar reasons ( w (cid:48) , w (cid:48) ) , and cycles containing y (cid:48) or w (cid:48) with ﬁnite total distance, should be avoided.
We note that having ( y, y ) or ( w, w ) would also break the non-quasipseudometric assumptions, and thus should avoid them as well (although cycles are okay here since they do not restrict values of c ).
In fact, with metrics more friendly to zero distances (than distortion and violation, which are based on distance ratios), it might be possible to allow them and obtain better bounds in the second-moment argument below in the proof for Thm. 4.6.
• Similarly, we can’t have ( y, y (cid:48) ) , ( y (cid:48) , y ) , ( w, w (cid:48) ) , or ( w (cid:48) , w ) in S ( · ) , as they will be mapped to ( y, y ) , ( y (cid:48) , y (cid:48) ) , ( w, w ) , or ( w (cid:48) , w (cid:48) ) .
With these understandings of the pattern shown in Fig.
6, we are ready to discuss the constructed quasimetric space and training sets.
Proof of Thm. 4.6.
Our proof follows the outline listed above.
Construct large quasimetric spaces containing many copies of the (potentially failing) structure in Fig. 6.
For any n > 0 , consider the following quasimetric space ( X n , d n ) of size n , with one-hot features.
WLOG, assume n = 12 k is a multiple of 12 .
If it is not, set at most 11 elements to have inﬁnite distance with every other node.
This won’t affect the asymptotics.
Let the n = 12 k elements of the space be
[Equation]
[Table] [Table]
[Table] [Table]
[Table] [Table]
z
[Table] [Table]
z k +1 , . . .
, z 2 k } , (74) with quasimetric distances, ∀ i, j ,
[Equation]
where subscripts are colored to better show when they are the same (or different), unlisted distances are inﬁnite (except that d n ( u, u ) = 0 , ∀ u ∈ X ).
Essentially, we equally divide the 12 k nodes into 6 “types”, { x, y, w, z, w (cid:48) , y (cid:48) } , corresponding to the 6 nodes from Fig.
6, where each type has half of its nodes corresponding to the left pattern (of Fig. 6), and the other half corresponding to the right pattern, except for the z types.
Furthermore,
• Among the left-pattern nodes, each set with the same subscript are bundled together in the sense that x left i only has ﬁnite distance to y left i which only has ﬁnite distance to w (cid:48) left i (instead of other y left j ’s or w (cid:48) left k ’s).
However, since distance to/from y left i and w left i are inﬁnite anyways, we can pair
[Equation]
for any i, j, l, h , to obtain a left pattern.
• Among the right-pattern nodes, each set with the same subscript are bundled together in the sense that x right i only has ﬁnite distance to y (cid:48) right i , and y right j which only has ﬁnite distance to w right j (instead of other y (cid:48)
right j ’s or w right k ’s).
However, since are distances are inﬁnite anyways, we can pair
[Equation]
for any i, j, l, h , to obtain a right pattern.
We can see that ( X , d ) indeed satisﬁes all quasimetric space requirements (Defn. 2.1), including triangle inequalities (e.g., by, for each ( a, b ) with ﬁnite distance d n ( a, b ) < ∞ , enumerating ﬁnite-length paths from a to b ).
Now consider the sampled training set S .
• We say S is bad on a left pattern speciﬁed by i left , j left , l left , h left , if
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation]
• We say S is bad on a right pattern speciﬁed by i right ,
j right , l right , h right , if S ⊃ { ( x right
i right , z h right ) , ( x right i right , y (cid:48)
right i right ) , ( y (cid:48)
right j right , w right j right ) ,
( w right j right ,
z h right
[Table]
[Equation] [Equation]
∅
[Equation]
∩
[Equation]
[Table]
y right
j right
[Table]
z h right )
[Equation]
[Table]
y right
[Table]
right
[Table]
y right j right ) , ( w right j right , w right j right ) , ( y
[Table]
right i right , y
[Table]
right i right
[Table] [Table]
w
[Table]
right l right , w
[Table]
right
[Table]
w
[Table]
right
[Table] [Table]
y
[Table]
right
[Table]
y
[Table]
right
[Table]
right
[Table] [Table]
w
[Table]
j
[Table]
, w
[Table]
right l right
[Table]
[Table]
y
[Table]
right
[Table]
right , y right j
[Table]
[Table]
right
[Table]
w right
[Table]
[Equation] [Equation]
Most importantly,
• If S is bad on a left pattern speciﬁed by i left , j left , l left , h left , consider the orthogonal transform that interchanges y left i left ↔ y (cid:48) left j left and w left l left ↔ w (cid:48) left i left on 2nd input.
In S , the
possible transformed pairs are
[Equation]
The crucial observation is that the transformed training set just look like one sampled from a quasimetric space where
– the quasimetric space has one less set of left-pattern elements, – the quasimetric space has one more set of right-pattern elements, and – transformed training set is bad on that extra right pattern (given by the extra set of right-pattern elements), which can be easily veriﬁed by comparing the transformed training set with the requirements in Eqs.
(86) and (87).
• Similarly, if S is bad on a right pattern speciﬁed by i right , j right , l right , h right , consider the orthogonal transform that interchanges y right
j right ↔ y (cid:48)
right i right and w right j right ↔ w (cid:48)
right l right on 2nd input.
In S the possible transformed pairs are
[Equation]
[Table]
Again
[Table]
the crucial observation is that the transformed training set just look like one sampled from a quasimetric space where
– the quasimetric space has one less set of right-pattern elements, – the quasimetric space has one more set of left-pattern elements, and – transformed training set is bad on that extra left pattern (given by the extra set of left-pattern elements),
which can be easily veriﬁed by comparing the transformed training set with the requirements in Eqs.
(84) and (85).
Therefore
[Table]
when S is bad on both a left pattern and a right pattern (necessarily on disjoint sets of pairs)
[Table]
we consider the following orthogonal transform composed of:
(a) both transforms speciﬁed above (which only transforms 2nd inputs), (so that after this we obtain another possible training set of same size from the quasimetric space that is only different up to some permutation of X ) (b) a permutation of X (on both inputs) so that the bad left-pattern nodes and the bad
This transforms gives another possible training set of same size from the same quasimetric space, also is bad on a left pattern and a right pattern .
Moreover, with a particular way of select bad patterns (e.g., by the order of the subscripts), this process is reversible .
Therefore, we have deﬁned a way to pair up all such bad training sets.
Consider the predictors ˆ d before and ˆ d after trained on these two training sets (before and after transform) with an learning algorithm equivariant to orthogonal transforms.
Assuming that they satisfy non-negativity and Identity of Indiscernibles, we have,
• The predictors have the same distortion over respective training sets.
Therefore
we denote
this
distortion
as dis S ( ˆ d )
without specifying the predictor
ˆ
d
or training set S .
• the predictors must predict the same on heldout pairs in the sense that
[Equation]
Focusing on the ﬁrst
we denote
[Equation]
without specifying the predictor ˆ d or the speciﬁc y and z
However, the quasimetric constraints on heldout pairs ( y left i left , z h left ) and ( y right j right ,
z h right )
are completely different (see the left vs. right part of Fig. 6).
Therefore, as shown in Fig.
6, assuming non-negativity , one of the two predictors must have total violation at least
[Equation]
Fixing a large enough c , two terms in the max of Eq.
(91) can equal for some ˆ d ( y, z ) , and are respectively decreasing and increasing in ˆ d ( y, z ) .
In that case, we have
[Equation]
for δ > 0 such that
[Equation]
Solving the above quadratic equation gives
[Equation]
leading to
[Equation]
Therefore
[Equation]
choosing c ≥ f 2 n (4 f n + 1) 2 gives
[Equation]
Hence, for training sets that are bad on both a left pattern and a right pattern
[Equation]
we have shown a way to pair them up such that
• each pair of training sets have the same size, and • the algorithm fail on one of each pair by producing a distance predictor that
– has either distortion over training set ≥ f n , or violation ≥ f n , and – has test MSE ≥ f n .
Remark B.1.
Note that all training sets of size m has equal probability of being sampled.
Therefore, to prove the theorem, it sufﬁces to show that with probability 1 − o (1) , we can sample a training set of size m that is bad on both a left pattern and a right pattern .
Consider sampling training set as individually collecting each pair with a certain probability p , and carefully analyze the conditions to sample a training set with the special properties with high probability 1 − o (1) .
In probabilistic methods, it is often much easier to work with independent random variables.
Therefore, instead of considering uniform sampling a training set S of ﬁxed size m , we consider including each pair in S with probability p , chosen independently.
We will ﬁrst show result based on this sampling procedure via a second moment argument, and later extend to the case with a ﬁxed-size training set.
First, let’s deﬁne some notations that ignore constants:
[Equation]
We start with stating a standard result from the second moment method (Alon & Spencer, 2004).
Corollary B.2 (Corollary 4.3.5 of (Alon & Spencer, 2004)).
Consider random variable
X =
X 1 + X 2 + · · · + X n , where X i is the indicator random variable for event
A i .
Write i ∼ j if i (cid:54) =
j and the pair of events ( A i , A j ) are not independent.
Suppose the following quantity does not depend on i :
[Equation]
If E [ X ] → ∞ and ∆ ∗ (cid:28) E
[ X ] , then X ∼ E [ X ] with probability 1 − o (1) .
We will apply this corollary to obtain conditions on p such that S with probability 1 − o
[Equation]
1) is bad on some left pattern, and conditions such that S with probability 1 − o (1) is bad on some right pattern.
A union bound would then give the desired result.
• S is bad on some left pattern.
Recall that a left pattern is speciﬁed by i left , j left , l left , h left all ∈ [ k ] :
[Equation]
Therefore, we consider k 4 = ( n 12 ) 4 events of the form
A i left ,j left ,l
[Equation]
[Equation]
is bad on the left pattern at i left , j left , l left , h left } .
[Equation]
Obviously, these events are symmetrical, and the ∆ ∗ in Eq.
(104) does not depend on i .
By the quasimetric space construction and the requirement for S to be bad on a left pattern in Eqs.
(84) and (85), we can see that ( i left , j left , l left , h left ) ∼ ( i (cid:48) left , j (cid:48) left , l (cid:48) left , h (cid:48) left ) only if i left =
i (cid:48) left or j left =
j (cid:48) left or l left =
l (cid:48) left or h left
= h (cid:48) left .
Therefore, we have
E [
[Equation]
(include 4 pairs & exclude 10 pairs)
[Equation]
(share j left )
+
[Equation] [Equation]
(share i left )
+ n
[Equation]
p
[Equation] [Equation]
p )
[Equation]
(share l left )
+ n
[Equation]
p
[Equation]
(1 − p )(share h left ) + n 2 p 2 (1 − p )(share j left , i left ) + n 2 p 4 (1 − p )(share j left , l left ) + n 2 p 4 (1 − p )(share j left , h left ) + n 2 p 2 (1 − p )(share i left , l left ) + n 2 p (1 − p )(share i left , h left ) + n 2 p 3 (1 − p )(share l left , h left )
+ n (1 − p )(share i left , l left , h left ) + np 3 (1 − p )
(share j left , l left , h left )
+ np (1 − p )(share j left , i left , h left ) + np 2 (1 − p ) (share j left , i left , l left ) ∼ n 3 p 2 (1 − p ) 7 + n
( p 2 (1 − p ) 4 + p (1 − p ) 6 ) (107) + n ((1 − p ) 3 + p 2 (1 − p )) .
Therefore, to apply Corollary B.2, we need to have
[Equation] [Equation]
[Equation] [Equation]
which gives
[Equation]
as
a
sufﬁcient
condition to
for
S to be
bad on
some
left pattern with probability 1 − o
• S is bad on some right pattern.
Recall that a right pattern is speciﬁed by i right , j right , l right , h right all ∈ [ k ] :
[Equation]
Similarly, we consider k 4 = ( n 12 ) 4 events of the form
[Equation]
on the left pattern
[Equation]
i
[Equation]
j
[Equation]
l
[Equation]
Again, these events are symmetrical, and ∆ ∗ in Eq.
(104) does not depend on i .
Similarly, we have
[Equation]
include
[Equation]
& exclude 10 pairs)
[Equation]
(share i
[Equation]
(share
[Equation]
(share h
[Equation] [Equation]
(share l
[Equation] [Equation]
share i
[Equation]
(share i
[Equation]
, h right
[Equation]
(share
[Equation] [Equation]
right
[Equation]
(share j right , h right )
[Equation]
(share j right , l right )
[Equation]
(share
[Equation]
right , l right )
[Equation]
(share j right , h right , l right )
[Equation]
(share i right , h right , l right )
[Equation]
(share i right , j right , l right )
[Equation]
(share i right , j right , h right )
[Equation]
[Equation] [Equation]
Therefore, to apply Corollary B.2, we need to have
[Equation] [Equation]
[Equation] [Equation]
122) which gives p
[Equation] [Equation]
as a sufﬁcient condition to for S to be bad on some right pattern with probability 1 − o (1) .
So, by union bound, as long as
[Equation]
S is bad on some left pattern and some right pattern with probability 1 − o (1) .
Extend to ﬁxed-size training sets and show that, under similar conditions, we sample a training set with the special properties with high probability 1 − o (1) .
To extend to ﬁxed-size training sets, we consider the following alteration procedure:
(a) Sample training set S by independently include each pair with probability p (cid:44) m
+ δ n 2 , for some δ > 0 .
(b) Show that with high probability 1 − o (1) , we end up with [ m, m + 2 δ ] pairs in S .
(c) Make sure that p satisfy Eq.
(125) and Eq.
(126) so that S is bad on some left pattern and some right pattern with high probability 1 − o (1) .
(d) Randomly discard the additional pairs, and show that with high probability 1 − o (1) this won’t affect that S is bad on some left pattern and some right pattern.
We now consider each step in details:
(a) Sample training set S by independently include each pair with probability p (cid:44) m
+ δ n 2 , for some δ > 0 .
For p (cid:44)
m + δ
n 2 , the number of pairs in the training set is distributed as
[Equation]
(b) Show that with high probability 1 − o (1) , we end up with [ m, m + 2 δ ] pairs in S .
Standard Binomial
concentration tells us that,
[Equation]
which
can
be satisﬁed if
[Equation]
(c) Make sure that p satisfy Eq.
(125) and Eq.
(126) so that S is bad on some left pattern and some right pattern with high probability 1 − o (1) .
Therefore, we want
[Equation]
(d) Randomly discard the additional pairs, and show that with high probability 1 − o (1) this won’t affect that S is bad on some left pattern and some right pattern.
Consider
any
speciﬁc bad
left pattern
and
a right pattern in
S .
It is
sufﬁcient that
we
don’t break these two patterns during discarding
Since we only discard pairs, it sufﬁces to only consider the pairs we want to preserve, which are a total of 8 pairs across two patterns.
Each such pair is discarded the probability ≤ 2 δm , since we remove at most 2 δ pairs.
By union bound,
P [ all 8 pairs are preserved ] ≥
[Equation] [Equation]
[Equation]
Hence, it sufﬁces to make sure that
[Equation]
Collecting all requirements, we have
[Equation]
[Equation] [Equation]
[Equation] [Equation]
Assume that
[Equation] [Equation]
[Equation]
It can be easily veriﬁed that using δ (cid:44) n 1 . 1 satisﬁes all conditions.
Hence, for a uniformly randomly sampled training set S with size m , S is bad on some left pattern and some right pattern with high probability 1 − o (1) , as long as
[Equation] [Equation]
This is exactly the condition we need to prove the theorem (see Remark B.1).
This concludes the proof.
B.3. D ISCUSSIONS
Training set
size
dependency
Intuitively, when the training set has almost all pairs, violation can be lowered by simply ﬁtting training set well; when it is small and sparse, the learning algorithm may have an easier job ﬁnding some consistent quasimetric.
Thm.
4.6 shows that, outside these two cases, algorithms equivariant to orthogonal transforms can fail.
Note that for the latter case, Thm. 4.
requires the training fraction to decrease slower than n − 1 / 2 , which rules out training sizes that is linear in n .
We leave improving this result as future work.
Nonetheless, Thm. 4.6 still covers common scenarios such as a ﬁxed fraction of all pairs, and highlights that a training-data-agnostic result (such as the ones for PQEs) is not possible for these algorithms.
Proof techniques.
In embedding theory, it is quite standard to analyze quasimetrics as directed graphs due to their lack of nice metric structure.
In the proof for Thm. 4.6, we used abundant techniques from the probabilistic method, which are commonly used for analyzing graph properties in the asymptotic case, including Corollary B.2 from the second moment technique, and the alteration technique to extend to ﬁxed-size training sets.
While such techniques may be new in learning theory, they are standard for characterizing asymptotic probabilities on graphs, which quasimetrics are often analyzed as (Charikar et al., 2006; Mémoli et al., 2018).
To provide more intuition on why these techniques are useful here, we note that the construction of a training set of pairs is essentially like constructing an Erd ˝ os-Rényi random graph on n 2 vertices.
Erd ˝ os-Rényi (undirected) random graphs come in two kinds:
• Uniformly sampling a ﬁxed number of m edges; • Adding an edge between each pair with probability p , decided independently.
The latter, due to its independent decisions, is often much easy to analyze and preferred by many.
The alteration technique (that we used in the proof) is also a standard way to transfer a result on a random graph of the latter type, to a random graph of the former type (Bollobás & Béla, 2001).
Readers can refer to (Alon & Spencer, 2004; Bollobás & Béla, 2001; Erd˝os & Rényi, 1959) for more in-depth treatment of these topics.
Generalization to other transforms.
The core of this construction only relies on the ability to swap (concatenated) inputs between ( x, y ) ↔ ( x, y (cid:48) ) and between ( y, w ) ↔ ( y, w (cid:48) ) via a transform.
For instance, here the orthogonal transforms satisfy this requirement on one-hot features.
Therefore, the result can also be generalized to other transforms and features with the same property.
Our stated theorem focuses on orthogonal transforms because they correspond to several common learning algorithms (see Lemma 4.5).
If a learning algorithm is equivariant to some other transform family, it would be meaningful to generalize this result to that transform family, and obtain a similar negative result.
We leave such extensions as future work.
B.3. C OROLLARY OF D ISTORTION AND V IOLATION FOR U NCONSTRAINED MLP S
Corollary B.3 (Distortion and Violation of Unconstrained MLPs).
Let ( f n ) n be an arbitrary sequence of desired violation values.
There is an inﬁnite collection of quasimetric spaces (( X n , d n )) n =1 , 2 ,... with |X n | = n , X n ⊂ R n such that MLP trained with squared loss in NTK regime converges to a function ˆ d that either
• fails non-negativity, or • vio ( ˆ d ) ≥ f n ,
with probability 1 / 2 − o (1) over the random training set S of size m , as long as S does not contain almost all pairs 1 − m/n 2 = ω ( n − 1 / 3 ) , and does not only include few pairs m/n 2 = ω ( n − 1 / 2 ) .
Proof of Corollary B.3.
This follows directly from Thm. 4.6 and standard NTK convergence results obtained from the kernel regression optimality and the positive-deﬁniteness of the NTK.
In particular,
Proposition 2 of (Jacot et al., 2018) claims that the NTK is positive-deﬁnite when restricted to a hypersphere.
Since the construction in proof of Thm.
4.6 uses one-hot features, the input (concatenation of two features) lie on the hypersphere with radius √ 2 .
Hence, the NTK is guaranteed positive deﬁnite.
[Figure]
B.3. E MPIRICAL V ERIFICATION OF THE F AILURE C ONSTRUCTION
We train unconstrained MLPs on the toy failure construction discussed in Sec.
4.2 (reproduced as
Fig. 6).
The MLP uses 12-1024-1 architecture with ReLU activations, takes in the concatenated one-hot features, and directly outputs predicted distances.
Varying c ∈ { 1 , 10 , 100 , 1000 } , we train the above MLP 5 times on each of the two patterns in Fig. 6, by regressing towards the training distances via MSE loss.
In Fig.
7, we can see that all training runs conclude with small training error, and indeed the trained
MLPs predict very similarly on the heldout pair, regardless whether it is trained on the left or right pattern of Fig. 6, which restricts the heldout pair distance differently.
This veriﬁes our theory (Thm. 4.6 and Corollary B.3) that algorithms equivariant to orthogonal transforms (including MLPs in NTK regime) cannot distinguish these two cases and thus must fail on one of them.
C P ROOFS AND D ISCUSSIONS FOR S EC . 5: P OISSON Q UASIMETRIC E MBEDDINGS (PQE S ) C. N ON - DIFFERENTIABILITY OF C ONTINUOUS -V ALUED S TOCHASTIC P ROCESSES
In this section we formalize the argument presented in Sec.
5.3 to show why continuous-valued stochastic processes lead to non-differentiability.
Fig.
8 also provides a graphical illustration of the general idea.
Proposition C.1 (Quasimetric Embeddings with Continuous-Valued Stochastic Processes are not Differentiable).
Consider any R k -valued stochastic process { R ( u ) } u ∈ R d such that u (cid:54) = u (cid:48) =
⇒
P [ R ( u ) =
R ( u (cid:48) )] < c for some universal constant c < 1 .
Then P [ R ( u ) ≤ R ( u (cid:48) )] is not differentiable at any u = u (cid:48) .
Proof of Proposition C.1.
Assume that the quantity is differentiable.
Then it must be continuous in u and v .
We will use the ( (cid:15), δ ) -deﬁnition of continuity.
At any u ∈ R d , consider small (cid:15) ∈ (0 , 1 − c 3 ) .
By continuity, since
P [ R ( u
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Table]
[Figure]
we can ﬁnd (
cid:15) ∈ R
d such that
[Equation] [Equation]
[Equation]
However, by assumption, P [ R ( u ) =
R ( u + δ )] < c .
Therefore,
[Equation] [Equation]
which implies
[Equation]
By contradiction, the quantity must not be differentiable at any u = u (cid:48) .
C. PQE-GG: G AUSSIAN - BASED M EASURE AND G AUSSIAN S HAPES
In Sec. 5.1, we presented the following PQE-LH formulation for Lebesgue measures and half-lines:
[Equation]
Here, u i,j and v i,j receive zero gradient when u i,j ≤ v i,j .
Gaussian shapes parametrization.
We therefore consider a set parametrization where no one set is entirely contained in a different set— the regions regions ⊂ R 2 between an axis and a 1D Gaussian density function of ﬁxed variance σ 2 shape = 1 .
That is, for each given u ∈ R , we consider sets
[Equation]
[Equation]
where f N ( b ; µ, σ 2 ) denotes the density of 1D Gaussian N ( µ, σ 2 ) with mean µ and variance σevaluated at b .
Since the Gaussian density function have unbounded support, these sets, which are translated versions of each other, never have one set fully contained in another.
For latent u ∈ R
h × k reshaped as 2D, our set parametrizations are,
[Equation]
i ∈ [
[Equation] [Equation]
[Equation]
A Gaussian-based measure.
These subsets of R 2 always have Lebesgue measure 1 , which would make PQE symmetrical (if used with a (scaled) Lebesgue measure).
Thus, we use an alternative
R 2 measure given by the product of a R Lebesgue measure on the b -dimension (i.e., dimension of the function value of the Gaussian density) and a R Gaussian measure on the a -dimension (i.e., dimension on the input of the Gaussian density) centered at 0 with learnable variances ( σ 2 measure ) i,j .
To avoid being constrained by the bounded total measure of 1 , we also optimize learnable positive scales
c i,j > 0 .
Hence, the each Poisson process has a mean measure as the product of a R Lebesgue measure and a R Gaussian with learnable standard deviation, then scaled with a learnable scale.
Note that the Gaussian measure should not be confused with the Gaussian shape.
Their parameters also are fully independent with one another.
Computing measures of Gaussian shapes and their intersections.
The intersection of two such
Gaussian shapes is formed by two Gaussian tail shapes, reﬂected around the middle point of the two
Gaussian means (since they have the same standard deviation σ shape = 1 ).
Hence, it is sufﬁcient to describe how to integrate a Gaussian density on a Gaussian measure over an interval.
Applying this with different intervals would give the measure of the intersection, and the measures of the two Gaussian shapes.
Omit indices i, j for clarity.
Formally, we integrate the Gaussian density f N ( a ; u, σ 2 shape ) over the centered Gaussian measure with variance σ 2 measure , which has density f N ( a ; 0 , σ 2 measure ) :
[Equation]
which is also another Gaussian integral (e.g., considered as integrating the product measure along the a line of the form y
= x + u ).
After standard algebraic manipulations (omitted here), we obtain (cid:90)
[Equation] [Equation]
for
[Equation]
This can be easily evaluated using statistical computing packages that supports computing the error function and/or Gaussian CDF.
Moreover, this ﬁnal form is also readily differentiable with standard gradient formulas.
To summarize,
• each set A ( u ) has total measure
[Equation]
• the intersection of A ( v ) and A ( u 2 ) , for v ≤ u 2 has measure
[Equation]
Interpretation and representing any total order.
Consider two Gaussian shapes A ( v ) and A ( u 2 ) .
Note that the Gaussian-based measure µ Gaussian is symmetric around and centered at 0 .
Therefore,
[Equation] [Equation]
[Equation]
Moreover, scaling the rates of a Poisson makes it more concentrated (as a Poisson’s mean grows as the square of its standard deviation) so that lim c →∞ P [Pois( cµ 1 ) ≤ Pois( cµ 2 )]
= 1 µ 1 <µ 2 for µ 1 (cid:54) = µ 2 .
Then any total order can be represented as the limit of a Poisson process with Gaussian shapes, with the shapes’ having their means arranged according to the total order, as the scale on the
Gaussian-based measure grows to inﬁnity.
C. T HEORETICAL G UARANTEES FOR PQE S
Theorem 5.2 (Distortion and violation of PQEs).
Under the assumptions of Sec. 4, any quasimetric space with size n and treewidth t admits a PQE-LH and a PQE-GG with distortion O ( t log 2 n ) and violation 1 , with an expressive encoder (e.g., a ReLU network with ≥ 3 layers and polynomial width).
In Sec.
5.4, we presented the above theoretical distortion and violation guarantees for PQE-LH and
PQE-GG.
Furthermore, we commented that the same guarantees apply to more generally to PQEs satisfying a mild condition.
Here, we ﬁrst precisely describe this condition, show that PQE-LH and
PQE-GG do satisfy it, state and prove the general result, and then show the above as a straightforward corollary.
C.3. T HE C ONCENTRATION P ROPERTY
Recall that PQEs are generally deﬁned with measures µ and set parametrizations A as
[Equation]
where
[Equation]
Because the measures µ and set parametrizations A themselves may have parameters (e.g., as in
PQE-GG), we consider them as classes of PQEs.
E.g., PQE-GG is a class of PQEs such that the µ is the speciﬁc Gaussian-based form, and A is the speciﬁc Guassian-shape.
Deﬁnition C.2 (Concetration Property of PQEs).
Consider a PQE class with h mixtures of quasipartition distributions, each from k Poisson processes.
We say that it has concentration property if it satisﬁes the following.
Consider any ﬁnite subset of X (cid:48)
⊂ X , and arbitrary function g : X → R h × k .
There exists a sequence of (( f ( n ) , µ ( n ) , A ( n ) ) n such that
• f ( n ) : X (cid:48)
→ R d , • µ ( n ) , A ( n ) are valid members of this PQE, • E π z ∼ Π PQE z ( µ i ,A i ) (cid:2)
π z ( f ( n ) ( x (cid:48) ) , f ( n ) ( y (cid:48) ))
(cid:3)
uniformly converges to 1 − (cid:81) j 1 g ( x ) i,j ≤ g ( y ) i,j , over all
mixtures
i
and
pairs x
, y ∈ X
(cid:48) .
A sufﬁcient condition.
It sufﬁces to make the probabilities
[Equation]
along some PQE sequence uniformly converge to the indicators
[Equation]
This is sufﬁcient since product of bounded functions is uniformly convergent, if each function is.
Both statements below together form a sufﬁcient condition for Eq.
(159) to uniformly converge to
Eq.
For any g , there exists a speciﬁc PQE of this class satisfying • Measures (of set differences) are consistent with g with some margin (cid:15) > 0 : ∀ i ∈
[ h ] ,
[Equation]
∈ [ k ] , x ∈ X (cid:48)
[Equation]
y ∈ X (cid:48) ,
[Equation]
[Equation] [Equation]
• Either of the following: – One side must be zero: ∀ i ∈
[ h ] , j ∈ [ k ] , x ∈ X , y ∈ X , ( µ i,j
( A i,j ( f ( x ))
\ A i,j ( f ( y
[Equation]
– Max measure is bounded by some constant c > 0 :
[Equation]
For any given speciﬁc PQE of this class, for any positive scale d > 0 , there is another PQE (with same formulation) whose measures (of set differences) equal exactly those of the given PQE scaled by d .
We now show that this is a sufﬁcient condition.
Note that a Poisson distribution has standard deviation equal to square root of its mean.
This means that as we scale the rate of a Poisson, it becomes more concentrated.
Applying to Poisson race probability, we have, for 0 ≤ µ 1 + (cid:15) < µ 2 ,
[Equation]
the
other direction of Poisson race probability
[Equation]
Therefore, applying to scaled versions of the PQE from Item 1 above, we have thus obtained the desired sequence, where Eq.
(159) uniformly converges to Eq.
(160) with rate O (1 /d ) .
Lemma C.3.
PQE-LH and PQE-GG both have the concentration property.
Proof of Lemma C.3.
We show that both classes satisfy the above sufﬁcient condition.
• PQE-LH: Lebesgue measure λ and half-lines.
WLOG, since X is countable, we assume that g satisﬁes
[Equation]
, y ∈ X (cid:48) .
The encoder in Item 1 above f : X → R h × k can simply be g .
We then have µ i,j
( A i,j ( f ( y ))
\ A i,j ( f ( x )))
=
Leb(( −∞ , g ( y )]
\ ( −∞ , g ( x )])
=
( g ( y ) i,j − g ( x ) i,j ) + .
This ensures that one side is always zero.
Furthermore, scaling can be done by simply scaling the encoder f .
Hence, PQE-LH satisﬁes this constraint.
• PQE-GG: Gaussian-based measure and Gaussian shapes (see Appendix C.2).
Because X (cid:48) is ﬁnit, we can have positive constant margin for the PQE requirements in Item 1.
(Inﬁnite X (cid:48) does not work because the total measure is ﬁnite (for a speciﬁc PQE-GG with speciﬁc values of the scaling).)
Concretely, we satisfy both requirements via
– in descending order of g ( · ) i,j we assign Gaussian shapes increasingly further from the origin; – scaling comes from that we allow scaling the Gaussian-based measure.
Hence, PQE-GG satisﬁes this constraint for ﬁnite X .
C.3. A G ENERAL S TATEMENT
We now state the general theorem for PQEs with the above concentration property.
Theorem C.4 (Distortion and violation of PQEs (General)).
Consider any PQE class with the concentration property.
Under the assumptions of Sec. 4, any quasimetric space with size n and treewidth t admits such a PQE with distortion O ( t log 2 n ) and violation 1 , with an expressive encoder (e.g., a ReLU network with ≥ 3 hidden layers, O ( n ) hidden width, and O ( n 2 ) quasipartition distributions, each with O ( n ) Poisson processes.).
Before proving this more general theorem, let us extend a result from Mémoli et al.
Lemma C.5 (Quasimetric Embeddings with Low Distortion; Adapted from Corollary 2 in
Mémoli et al.
Let M =
( X, d ) be a quasipseudometric space with treewidth t , and n = | X | .
Then M admits an embedding into a convex combination (i.e., scaled mixture) of O ( n 2 ) quasipartitions with distortion O ( t log 2 n ) .
Proof of Lemma C.5.
The distortion bound is proved in Corollary 2 in (Mémoli et al., 2018), which states that any quasipseudometric space with n elements and t treewidth admits an embedding into a convex combination of quasipartitions with distortion O ( t log 2 n ) .
To see that n 2 quasipartitions sufﬁce, we scrutinize their construction of quasipartitions in Algorithmof (Mémoli et al., 2018), reproduced below as Algorithm 1.
[Table]
Let H 1 , . . .
, H m be the sub-trees of H ∗ below root ( H ∗ ) that are hierarchical trees of separators of C m respectively
v
G
u
≤
[Table]
Many concepts used in Algorithm 1 are not relevant for our purpose (e.g., r -bounded quasipartition).
Importantly, we observe that for a given quasimetric space, the produced quasipartition is entirely determined by the random choice of z in Step 1, which is only used to compare with distance values between node pairs.
Note that there are n 2 node pairs, whose minimum distance is exactly 0 (i.e., distance from a node to itself).
Since z ≥ 0 , there are at most n 2 choices of z that lead to at most ndifferent quasipartitions, for all possible values of r .
The construction used to prove Corollary 2 of (Mémoli et al., 2018) uses exactly quasipartitions given by this algorithm.
Therefore, the lemma is proved.
Lemma C.5 essentially proves the ﬁrst half of Thm.
C.4.
Before proving the full Thm.
C.4, we restate the following result from (Hiraguchi, 1951), which gives us a bound on how many total orders are needed to represent a general partial order (i.e., quasipartition).
Theorem C.6 (Hiraguchi’s Theorem (Hiraguchi, 1951; Bogart, 1973)).
Let ( X, P ) be a partially ordered set such that | X | ≥ 4 .
Then there exists a mapping f : X → R (cid:98)| X | / 2 (cid:99) such that
[Equation]
Proof of Thm.
C.4.
It immediately follows from Lemma C.5 and Thm.
C.6 that any quasimetric space with n elements and treewidth t admits an embedding with distortion O ( t log 2 n ) into a convex combination of n 2 quasipartitions, each represented with an intersection of O ( n ) total orders.
Because the PQE class has concentration property, for any ﬁnite quasimetric space, we can simply select a PQE that is close enough to the desired convex combination of n 2 quasipartitions, to obtain distortion O ( t log 2 n ) .
Since each Poisson process in PQE takes a constant number of latent dimensions, we can have such a PQE with O ( n 3 ) -dimensional latents and n 2 quasipartition distributions.
It remains only to prove that we can compute such required latents using the described architecture.
Consider any x ∈ X ⊂ R d .
Since X is ﬁnite, we can always ﬁnd direction u x ∈ R d such that ∀ y ∈ X \ { x } , y T u
x (cid:54) =
x T u x .
That is, x has a unique projection onto u x .
Therefore, we can have c, b + , b − ∈ R such that
[Equation]
but for y ∈ X \ { x } , we have, for some a > 0 , either
[Equation]
or
[Equation]
Then, consider computing two of the ﬁrst layer features as, on input z , [
[Equation]
which, if z = x , is [1 , 1] ; if z (cid:54)
= x , is either [0 , 2 + a ] or [2 + a, 0] , for some a > 0 .
Then, one of the second layer features may sum these two features and threshold it properly would single out x , i.e., activate only when input is x .
After doing this for all x ∈ X , we obtain an n -dimensional second layer feature space that is just one-hot features.
The third layer can then just be a simple embedding look up, able to represent any embedding, including the one allowing a PQE to have distortion O ( t log n ) , as described above.
Because quasimetric embeddings naturally have violation 1 , this concludes the proof.
C.3. P ROOF OF T HM . 5.2: D ISTORTION AND VIOLATION OF PQE S
Proof of Thm. 5.2.
Lemma C.3 and Thm.
C.4 imply the result.
To see that polynomial width is sufﬁcient, note that the hidden width are polynomial by Thm.
C.4, and that the embedding dimensions needed to represent each of the O ( n 3 ) Poisson processes is constant 1 in both PQE-LH and PQE-GG.
Hence the latent space is also polynomial.
This concludes the result.
C.3. D ISCUSSIONS
Dependency on log n . log n dependency frequently occurs in distortion results.
Perhaps the most well-known ones are Bourgain’s Embedding Theorem (Bourgain, 1985) and the JohnsonLindenstrauss Lemma (Johnson & Lindenstrauss, 1984), which concern metric embeddings into
Euclidean spaces.
Dependency
on treewidth t
Treewidth t here works as a complexity measure of the quasimetric.
We will use a simple example to illustrate why low-treewidth is easy.
Consider the extreme case where the quasimetric is the shortest-path distance on a tree, whose each edge is converted into two opposing directed ones and assigned arbitrary non-negative weights.
Such a quasimetric space has treewidth 1 (see Defn. 2.2).
On a tree,
the shortest path between two points is ﬁxed, regardless
of the weights
assigned, 2. for
each internal node u and one of its child c , the followings are quasipartitions
[Equation]
Hence it can be exactly represented as a convex combination of quasipartitions.
However, both of observations becomes false when the graph structure becomes more complex (higher treewidth) and the shortest paths can are less well represented as tree paths of the tree composition.
Comparison with unconstrained MLPs.
Thm.
C.4 requires a poly-width encoder to achieve low distortion.
This is comparable with deep unconstrained MLPs trained in NTK regime, which can reach 0 training error (distortion 1 on training set) in the limit but also requires polynomial width (Arora et al., 2019).
Quasipseudometrics
and inﬁnite distances
Thm.
C.4 relies on our assumptions that ( X , d ) is not a quasipseudometric space and has all ﬁnite distances.
In fact, if we allow a PQE to have inﬁnite convex combination weights, it can readily represent quasipseudometric spaces with inﬁnite distances.
Additionally, PQE can still well approximate the quasimetric space with inﬁnities replaced with any sufﬁciently large ﬁnite value (e.g., larger than the maximum ﬁnite distance).
Thus, this limit is generally not important in practice (e.g., learning γ -discounted distances), where a large value and inﬁnity are usually not treated much differently.
Optimizing quasimetric embeddings.
From Thm.
C.4, we know that optimizing PQEs over the training set S w.r.t.
distortion achieves low distortion (and optimal violation by deﬁnition).
While directly optimizing distortion (or error on log distance or distance ratios, equivalently) seems a valid choice, such objectives do not always train stably in practice, with possible inﬁnities and zeros.
Often more stable losses are used, such as MSE over raw distances or γ -discounted distances γ d , for γ ∈ (0 , 1) .
These objectives do not directly relate to distortion, except for some elementary loose bounds.
To better theoretically characterize their behavior, an alternative approach with an average-case analysis might be necessary.
C. I MPLEMENTATION OF P OISSON Q UASIMETRIC E MBEDDINGS (PQE S )
Sec.
5.2 mentioned a couple implementation techniques for PQEs.
In this section, we present them in full details.
C.4. N ORMALIZED M EASURES
Consider a PQE whose each of j expected quasipartitions is deﬁned via k Poisson processes, with set parametrizations u
→ A i,j ( u ) , i ∈ [ h ] , j ∈ [ k ] .
To be robust to the choice of k , we instead use the normalized set parametrizations A (cid:48) i,j :
[Equation]
This does not change the PQE’s concentration property (Defn.
C.2) or its theoretical guarantees (e.g.,
Thms.
5.2 and C.4).
C.4. O UTPUTTING γ -D ISCOUNTED D ISTANCES
Recall the PQE quasimetric formulation in Eq.
(13), for α i ≥ 0 , and encoder f : X → R d : ˆ d ( x, y ) (cid:44) (cid:88)
i
[Equation]
With discount factor γ ∈ (0 , 1)
[Equation]
we can write the γ -discounted PQE distance as
[Equation]
Therefore, instead of learning α i ∈ [
[Equation]
∞ )
[Equation]
we can learn bases β i ∈ (0 , 1) such and deﬁne the γ -discounted PQE distance as
[Equation]
These bases β i ∈ (0 , 1) can be parametrized via a sigmoid transform.
Consider quasimetric learning w.r.t.
errors on γ -discounted distances (e.g., MSE).
Unlike the parametrization with directly learning the convex combination weights α
i ’s, such a parametrization (that learns the bases β i ’s) does not explicitly include γ and thus can potentially be more stable for a wider range of γ choices.
Initialization.
Consider learning bases β
i ’s via a sigmoid transform: learning b i and deﬁning β
i (cid:44) σ ( b i ) .
We must take care in initializing these b
i
’s
so that σ ( b i ) ’s are not too close to 0 or 1 , since we take a product of powers with these bases.
To be robust to different h numbers of quasipartition distributions, we initialize the each b i to be from the uniform distribution
[Equation]
[Equation]
which means that, at initialization, (cid:89)
[Equation]
185) providing a good range of initial outputs, assuming that the exponents (expected outputs of quasipartition distributions) are close to 0 .
Alternatively, b i ’s maybe parametrized by a deep linear network, a similar initialization is employed.
See Appendix C.4.3 below for details.
C.4. L EARNING L INEAR /C ONVEX C OMBINATIONS WITH D EEP L INEAR N ETWORKS
Deep linear networks have the same expressive power as regular linear models, but enjoy many empirical and theoretical beneﬁts in optimization (Saxe et al., 2013; Pennington et al., 2018;
Huh et al., 2021).
Speciﬁcally, instead of directly learning a matrix ∈
R m × n , a deep linear network (with bias) of l layers learns a sequence of matrices
[Equation]
[Equation] [Equation]
[Equation] [Equation]
where the linear matrix can be obtained with
[Equation]
and we require min
[Equation]
In our case, the convex combination weights for the quasipartition distributions often need to be large, in order to represent large quasimetric distances; in Poisson process mean measures with learnable scales (e.g., the Gaussian-based measure described in Appendix C.2), the scales may also need to be large to approximate particular quasipartitions (see Appendix C.3.1).
Therefore, we choose to use deep linear networks
to optimize these parameters.
In particular, • For the convex combination weights for h quasipartition distributions, – When learning the convex combination weights { α i } i ∈ [ h ] , we use a deep linear network to parametrize a matrix ∈ R 1 × h (i.e., a linear map from R h to R ), which is then viewed as a vector ∈ R h and applied an element-wise square transform a → a 2 to obtain non-negative weights α ∈ [0 , ∞ ) h ; – When learning the bases for discounted quasimetric distances β
i ’s (see Appendix C.4.2), we use a deep linear network to parametrize a matrix ∈ R h × 1 , which is then viewed as a vector ∈ R h and applied an element-wise sigmoid transform a → σ ( a ) to obtain bases β ∈ (0 , 1) h .
Note that here we parametrize
a matrix ∈
R h × 1 rather than R 1 × h as above for α i ’s. The reason for this choice is entirely speciﬁc to the initialization scheme we use (i.e., (fullyconnected layer weight matrix initialization, as discussed below).
Here the interpretation of a linear map is no longer true.
If we use R 1 × h , the initialization method would lead to the entries distributed with variance roughly 1 /n , which only makes sense if they are then added together.
Therefore, we use R h × 1 , which would lead to constant variance.
• For scales of the Poisson process mean measure, such as PQE-GG, we consider a slightly different strategy.
Consider a PQE formulation with h × k independent Poisson processes, from which we form h quasipartition distributions, each from k total orders parametrized by k Poisson processes.
The Poisson processes are deﬁned on sets
[Equation]
[Equation]
use mean measures
[Equation] [Equation]
and set parametrizations
[Equation]
to compute quantities
[Equation]
u
[Equation]
Scaling each mean measure independently.
Essentially, adding learnable scales (of mean measures) w ∈ [0 , ∞ )
h × k (or, equivalently, { w i,j ∈ [0 , ∞ ) } i,j ) gives a scaled set of measures { w i
[Equation]
j · µ i,
[Equation]
} i ∈
[Equation] [Equation]
This means that the quantities in Eq.
(197) becomes respectively scaled as
[Equation]
for u ∈ R
[Equation]
Convex combinations of all measures.
However, we can be more ﬂexible here, and allow not just scaling each measure independently, but also convex combinations of all measures .
Instead of having w as a collection of h × k scalar numbers ∈ [0 , ∞ ) , we have a collection of ( h × k ) vectors each having length ( h × k ) (or h × k -shape tensors)
[Equation]
and have the quantities in Eq.
(197) respectively scaled and combined as (cid:88)
[Equation]
for u ∈ R d
[Equation]
v ∈ R d
[Equation]
i ∈
[ h
[Equation]
∈
[Equation]
Note that these still are valid Poisson processes for a PQE.
Speciﬁcally, the new Poisson processes now all use the same set parametrization (as the collection of original ones), with different measures (as different weighted combinations of the original measures).
This generalizes the case where each mean measure is scaled independently (as w can be diagonal).
Therefore, we will apply this more general strategy using convex combinations of all measures .
Similarly to learning the convex combination weights of quasipartition distributions, we collapse a deep linear network into a tensor ∈ R
h × k × h × k , and apply an element-wise square a → a 2 , result of which is used as the convex combination weights w to ensure non-negativity.
Initialization.
For initializing the matrices ( M 1 , M 2 , . . .
, M l ) of a deep linear network (Eq.
(190)), we use the standard weight matrix initialization of fully-connected layers in PyTorch (Paszke et al., 2019).
The bias matrix B (Eq.
(191)) is initialized to all zeros.
When used for learning the bases for discounted quasimetric distances
β i ’s (as described in Appendix C.4.2), we have a deep linear network parametrizing a matrix ∈ R h × 1 , initialized in the same way as above (including initializing B as all zeros).
Consider the matrix up to before the last one:
[Equation]
M ∗ is essentially a projection to be applied on each row of the last matrix M l ∈ R h × m l − 1 , to obtain b i (which is then used to obtain bases β
i (cid:44) σ ( b i ) ).
Therefore, we simply rescale the M ∗ subspace for each row of M l and keep the orthogonal space intact, such that the projections would be distributed according to the distribution speciﬁed in Eq.
[Equation]
which has good initial value properties, as shown in Appendix C.4.2.
C.4. C HOOSING h THE N UMBER OF Q UASIPARTITION D ISTRIBUTIONS AND k THE N UMBER OF P OISSON P ROCESSES FOR E ACH Q UASIPARTITION D ISTRIBUTION
A PQE (class) is deﬁned with h × k independent Poisson processes with means { µ i,j } i ∈ [ h ] ,j ∈ [ k ] along with h × k set parametrizations { A i,j } i ∈ [ h ] ,j ∈ [ k ] .
For k pairs of means and set parametrizations, we obtain a random quasipartition.
A mixture (convex combination) of the resulting h random quasipartitions gives the quasimetric.
The choices of µ and A are ﬂexible.
In this work we explore PQE-LH and PQE-GG as two options, both using essentially the same measure and parametrization across all i, j (up to individual learnable scales).
These two instantiations both perform well empirically.
In this section we aim to provide some intuition on choosing these two hyperparameters h and k .
h the Number of Quasipartition Distributions
Theoretical result Thm.
C.4 suggest thats, for a quasimetric space with n elements, n 2 quasipartition distributions sufﬁce to learn a low distortion
[Equation]
c : Train : Test
[Equation]
embedding.
Since this is a worst-case result, the practical scenario may require much fewer quasipartitions.
For instance, Appendix C.3.4 shows that O ( n ) quasipartitions is sufﬁcient for any quasimetric space with a tree structure.
In our experiments, h ∈ [8 , 128] quasipartition distributions are used.
k the Number of Poisson Processes for Each Quasipartition Distribution (Random Partial
Order)
It is well-known that such intersection of sufﬁciently many total orders can represent any partial order (Trotter, 1995; Hiraguchi, 1951).
This idea is equivalent with the dominance drawing dimension of directed graphs (Ortali & Tollis, 2019), which concerns an order embedding of the vertices to preserve the poset speciﬁed by the reachability relation.
In this graph theoretical view, several results are known.
(Felsner et al., 2010) prove that planar graphs have at most 8 dimension.
(Ortali & Tollis, 2019) show that the dimension of any graph with n vertices is at most min( w P , n 2 ) , where w P the maximum size of a set of incomparable vertices.
A simpler and more fundamental result can be traced to Hiraguchi from 1951:
Theorem C.6 (Hiraguchi’s Theorem (Hiraguchi, 1951; Bogart, 1973)).
Let ( X, P ) be a partially ordered set such that | X | ≥ 4 .
Then there exists a mapping f : X → R (cid:98)| X | / 2 (cid:99) such that ∀
[Equation]
x )
[Equation]
f ( y ) coordinate-wise .
Thm.
C.6 states that n 2 dimensions generally sufﬁce for any poset of size n ≥ 4 .
In our formulation, this means that using k = n 2 Poisson processes (giving n 2 random total orders) will be maximally expressive.
In practice, this is likely unnecessary and sometimes impractical.
In our experiments, we choose a small ﬁxed number k = 4 .
D E XPERIMENT S ETTINGS AND A DDITIONAL R ESULTS
Computation power.
All our experiments run on a single GPU and ﬁnish within 3 hours.
GPUs we used include NVIDIA 1080, NVIDIA
Ti, NVIDIA
Ti, NVIDIA Titan Xp, NVIDIA
Titan RTX, and NVIDIA Titan V.
D. E XPERIMENTS FROM S EC . 3.2: A T OY E XAMPLE
In Sec. 3.2 and Fig. 2, we show experiment results on a simple 3 -element quasimetric space
Quasimetric space.
The quasimetric space has 3 elements with one-hot features ∈ R 3 .
Thequasimetric and training pairs are shown in Fig.
Unconstrained network.
The unconstrained
network
has architecture
, with ReLU
activations
[Figure]
Metric embedding.
The embedding space is 32 -dimensional, upon which corresponding metric is applied.
The encoder network has architecture 6-128-128-32, with ReLU activations.
Asymmetric dot products.
The embedding space is 32 -dimensional.
The two inputs are encoded with a different encoder of architecture 6-128-128-32, with ReLU activations.
Then the dot product of the two 32 -dimensional vector is taken, which parametrizes a distance estimate
Poisson Quasimetric Embeddings.
The embedding space is 32 -dimensional, which parametrizes 8 quasimetric distributions, each from 4 independent Poisson processes using (scaled) Lebesgue measure and half-lines.
We use deep linear networks, as described in Appendix C.4.3.
A deep linear network (without bias) of architecture 8-32-32-1 parametrizes the convex combination weights { α i } i ∈ [8] .
Another deep linear network (without bias) of architecture 32-64-64-32 parametrizes convex combination weights of the mean measures d ∈ [0 , ∞ ) 32 × 32 .
Note that these do not give many more effective parameters to PQEs as they are equivalent with simple linear transforms.
Optimization.
All models are trained w.r.t.
MSE on distances with the Adam optimizer (Kingma & Ba, 2014) with learning rate 0 .
0003 for 1000 iterations (without mini-batching since the training set has size 8 ).
Additional results.
Results with additional formulations (together with the ones presented in Fig. 2) are shown in Fig.
D. E XPERIMENTS FROM S EC . 6: E XPERIMENTS
Triangle inequality regularizer.
For methods that do not inherently respect triangle inequalities (e.g., unconstrained networks and asymmetrical dot products), we explore training with a regularizer that encourages following these inequalities.
By sampling random triplets uniformly over the training set, the regularizer is formulated as,
[Equation]
where the γ -discounted terms and the squared form allows easier balancing with the training loss, which, across all experiments, are MSEs on some γ -discounted distances.
PQE settings.
Across all experiments of this section, when given an encoder architecture mapping input to an R d latent space, we construct PQEs according to the following general recipe, to obtain the two PQEs settings used across all experiments: PQE-LH (PQE with Lebesgue measure and half-lines) and PQE-GG (PQE with Gaussian-based measure and Gaussian shapes, see see Appendix C.2):
• (Assuming d is a multiple of 4 ,)
We use h (cid:44) =
d/ 4 quasipartition distributions, each given by k (cid:44) 4 Poisson processes; • A deep linear network (see Appendix C.4.3), is used for parametrizing the convex combination weights α ∈ R d/ 4 or the bases β ∈ R d/ 4 (see Appendix C.4.2), we follow the initialization and parametrization described in Appendix C.4.3, with hidden sizes [ n hidden , n hidden , n hidden ] (i
e.,matrices/layers), where n hidden (cid:44) max(64 , 2 1+ (cid:100) log 2 ( d/ 4) (cid:101) ) .
• For PQE-GG, – The learnable σ 2 measure ∈ (0 , ∞ ) d (one for each Poisson Process) is achieved by optimizing the log variance, which is initialized as all zeros.
– The Gaussian-based measures need learnable scales.
We use a deep linear network to parametrize the [0 , ∞ )
d × d weights for the convex combinations of measures, as described in Appendix C.4.3.
Similarly, it has hidden sizes [ n hidden , n hidden , n hidden ] (i
e., 4 matrices/layers), where n hidden (cid:44) max(64 , 2 1+ (cid:100) log 2 d (cid:101) ) .
Note that the PQEs add only a few extra effective parameters on top of the encoder ( d for PQE-LH, and d + d 2 for PQE-GG), as the deep linear networks do not add extra effective parameters.
Mixed space metric embedding settings.
Across all experiments of this section, when given an encoder architecture mapping input to an R d latent space, we construct the metric embedding into mixed space as follows:
• (Assuming d is a multiple of 4 ,)
We use (1) a ( d/ 2) -dimensional
Euclidean space (2) a ( d/ 4) - dimensional (cid:96) 1 space, and a ( d/ 4) -dimensional spherical distance space (without scale).
• Additionally, we optimize three scalar values representing the log weights of the convex combination to mix these spaces.
DeepNorm and WideNorm method overview and parameter count comparison with PQEs.
Both DeepNorm and
WideNorm
parametrize asymmetrical norms.
When
used
to
approximate quasimetrics
they
are applied as ˆ d ( x, y ) (cid:44) f AsymNorm ( f Enc ( x ) − f Enc ( y )) , where f Enc is the encoder mapping from data space to an R d
latent
space and f
AsymNorm
is
either
the
DeepNorm
or the WideNorm predictor on that
latent
space (Pitis et
al
• DeepNorm is a modiﬁcation from Input Convex Neural Network (ICNN; Amos et al.
(2017)), with
restricted weight matrices and activation functions for positive homogeneity (a requirement of asymmetrical norms), and additional concave function for expressivity.
For an input latent space of R d , consider an n -layer DeepNorm with width w (i.e., ICNN output size) and the suggested intermediate MaxReLU activation and MaxMean ﬁnal aggregation (see (Pitis et al., 2020) for details of these functions).
This DeepNorm predictor f DeepNorm (on latent
space)
has
#parmaters of f DeepNorm =
[Equation]
matrices
[Equation]
input
[Equation]
each layer
[Equation]
matrices
[Equation]
neighboring
[Equation]
[Table]
[Equation]
which is on the order of O ( nw max( d, w )) .
In the common case where the hidden size w is chosen to be on the same magnitude as d , this becomes O ( nd 2 ) .
• WideNorm is based on the observation that
[Equation]
is an asymmetric norm when W is non-negative, where :: denotes vector concatenation.
WideNorm then learns many such norms each with a different W matrix parameter, before (again) feeding the norm values into a concave function and aggregating them together with MaxMean.
For an input latent space of R d
consider a WideNorm with c such learned norms with W matrices of shape R (2 d ) × w ≥ >.
This WideNorm predictor f WideNorm (on latent space), has
[Equation]
[Table]
[Equation]
[Table]
which is on the order of O ( cdw ) .
In the common case where both the number of components c and the output size of each component (before applying the l 2 -norm) are chosen to be on the same magnitude as d , this becomes O ( d 3 ) .
For both DeepNorm and WideNorm, their parameter counts are much larger than the number of effective parameters of PQEs ( d for PQE-LH and d + d 2 for PQE-GG).
For a 256 -dimensional latent space, this difference can be on the order of 10 6 ∼ 10 7 .
DeepNorm and WideNorm settings.
Across all experiments of this section, we evaluate 2 DeepNorm settings and 3 WideNorm settings, all derived from the experiment setting of the original paper (Pitis et al., 2020).
For both DeepNorm and WideNorm, we use MaxReLU activations, MaxMean aggregation, and concave function of 5 components.
For DeepNorm, we use 3 -layer networks with 2 different hidden sizes: 48 and 128 for the 48 -dimensional latent space in random directed graphs experiments, 512 and 128 for the 512 -dimensional latent space in the large-scale social graph experiments, 128 and 64 for the 128 -dimensional latent space in ofﬂine Q-learning experiments.
For
WideNorm, we components of size 32 and experiment with 3 different numbers of components: 32 , 48 , and 128 .
Error range.
Results are gathered across 5 random seeds, showing both averages and population standard deviations.
D.2. R ANDOM D IRECTED G RAPHS Q UASIMETRIC L EARNING
Graph generation.
The random graph generation is controlled by three parameters d , ρ un and ρ di .
d is the dimension of the vertex features.
ρ un speciﬁes the fraction of pairs that should have at least one (directed) edge between them.
ρ di speciﬁes the fraction of such pairs that should only have one (directed) edge between them.
Therefore, if ρ un = 1 , ρ di = 0 , we have a fully connected graph; if ρ un = 0 .
5 , ρ di = 1 , we have a graph where half of the vertex pairs have exactly one (directed) edge between them, and the other half are not connected.
For completeness, the exact generation procedure for a graph of n vertices is the following:
randomly add ρ un · n 2 undirected edges, each represented as two opposite directed edges; 2.
optimize R n × d vertex feature matrix using Adam (Kingma & Ba, 2014) w.r.t.
L align ( α = 2) + 0 .
· L uniform ( t = 3) from (Wang & Isola, 2020), where each two node is considered a positive pair if they are connected; 3.
randomly initialize a network f of architecture d -4096-4096-4096-4096-1 with tanh activations; 4.
for each connected vertex pair ( u, v ) , obtain d u → v (cid:44)
f
( feature ( u ))
− f ( feature ( v )) and d v → u = − d u → v ; 5.
for each ( u, v ) such that d u → v is among the top 1 − ρ di / 2 of such values (which is guaranteed to not include both directions of the same pair due to symmetry of d u → v ), make v → u the only directed edge between u and v .
We experiment with three graphs of 300 vertices and 64 -dimensional vertex features:
• Fig.
A graph generated with ρ un = 0 .
15 , ρ di = 0 .
85 ; • Fig.
12: A sparser graph generated with ρ un = 0 .
05 , ρ di = 0 .
85 ; • Fig.
A sparse graph with block structure by 1. generating 10 small dense graphs of 30 vertices and 32 -dimensional vertex features, using ρ un = 0 .
18 , ρ di = 0 .
generating a sparse 10 -vertex “supergraph” with 32 -dimensional vertex features, using ρ un = 0 .
22 , ρ di = 0 .
for each supergraph vertex (a) associating it with a different small graph, (b) for all vertices of the small graph, concatenate the supergraph vertex’s feature to the existing feature, forming 64 -dimensional vertex features for the small graph vertices, (c) picking a random representative vertex from the small graph, 4.
connecting all 10 representative vertices in the same way as their respective supergraph vertices are connected in the supergraph.
Architecture.
All encoder based methods (PQEs, metric embeddings, dot products) use 64-128128-128-48 network with ReLU activations, mapping 64 -dimensional inputs to a 48 -dimensional latent space.
Unconstrained networks use a similar 128-128-128-128-48-1 network, mapping concatenated the 128 -dimensional input to a scalar output.
Data.
For each graph, we solve the groundtruth distance matrix and obtain 300 2 pairs, from which we randomly sample the training set, and use the rest as the test set.
We run on 5 training fractions evenly spaced on the logarithm scale, from 0 .
01 to 0 .
Training.
We use 2048 batch size with the Adam optimizer (Kingma & Ba, 2014), with learning rate decaying according to the cosine schedule without restarting (Loshchilov & Hutter, 2016) starting from 10 − 4 to 0 over 3000 epochs.
All models are optimized w.r.t.
MSE on the γ -discounted distances, with γ = 0 .
When running with the triangle inequality regularizer, 683 ≈ 2048 / 3 triplets are uniformly sampled at each iteration.
Full results and ablation studies.
Figs.
11 to 13 show full results of all methods running on all three graphs.
In Fig.
14, we perform ablation studies on the implementation techniques for PQEs mentioned in Appendix C.4: outputting discounted distance and deep linear networks.
On the simple directed graphs such as the dense graph, the basic PQE-LH without theses techniques works really well, even surpassing the results with both techniques.
However, on graphs with more complex structures (e.g., the sparse graph and the sparse graph with block structure), basic versions of PQELH and PQE-GG starts to perform badly and show large variance, while the versions with both techniques stably trains to the best results.
Therefore, for robustness, we use both techniques in other experiments.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
D.2. L ARGE -S CALE S OCIAL G RAPHS Q UASIMETRIC L EARNING
Data source.
We choose the Berkeley - StanfordWebGraph (Leskovec & Krevl, 2014) as the largescale directed social graph, which consists of 685 , 230 pages as nodes, and 7 , 600 , 595 hyperlinks as directed edges.
Additionally, we also use the Youtube social network (Leskovec & Krevl, 2014;
Mislove et al., 2007) as a undirected social graph, which consists of 1 , 134 , 890 users as nodes, and 2 , 987 , 624 friendship relations as undirected edges.
Both datasets are available from the SNAP website (Leskovec & Krevl, 2014) under the BSD license.
Data processing.
For each graph, we use node2vec to obtain 128 -dimensional node features (Grover & Leskovec, 2016).
Since the graph is large, we use the landmark method (Rizi et al., 2018) to construct training and test sets.
Speciﬁcally, we randomly choose 150 nodes, called landmarks, and compute the distances between these landmarks and all nodes.
For directed graph, this means computing distances of both directions.
From the obtained pairs and distances, we randomly sample 2 , 500 , 000 pairs to form the training set.
Similarly, we form a test set of 150 , 000 from a disjoint set of 50 landmarks.
For the undirected graph, we double the size of each set by reversing the pairs, since the distance is symmetrical.
Architecture.
All encoder based methods (PQEs, metric embeddings, dot products) use 1282048-2048-2048-512 network with ReLU activations and Batch Normalization (Ioffe & Szegedy, 2015) after each activation, mapping 128 -dimensional inputs to a 512 -dimensional latent space.
Unconstrained networks use a similar 256-2048-2048-2048-512-1 network, mapping concatenated the 256 -dimensional input to a scalar output.
Training.
We use 1024 batch size with the Adam optimizer (Kingma & Ba, 2014), with learning rate decaying according to the cosine schedule without restarting (Loshchilov & Hutter, 2016) starting from 10 − 4 to 0 over 80 epochs.
All models are optimized w.r.t.
MSE on the γ -discounted distances, with γ = 0 .
When running with the triangle inequality regularizer, 342 ≈ 1024 / 3 triplets are uniformly sampled at each iteration.
[Table] [Table]
[Table] [Table]
Full results.
Tables 2 and 3 show full results of distance learning on these two graphs.
On the directed Berkeley - StanfordWebGraph , PQE-LH performs the best (w.r.t.
discounted distance MSE).
While PQE-GG has larger discounted distance MSE than some other baselines, it accurately predicts ﬁnite distances and outputs large values for unreachable pairs.
On the undirected Youtube graph, perhaps as expected, metric embedding methods have an upper hand, with the best performing method being an Euclidean space embedding.
Notably, DeepNorms and WideNorms do much worse than
PQEs on this symmetric graph.
D.2. O FFLINE Q-L EARNING
As shown in Proposition A.4 and Remark A.5, we know that a quasimetric is formed with the optimal goal-reaching plan costs in a MDP M =
( S , A , R , P , γ ) where each action has unit cost (i.e., negated reward).
The quasimetric is deﬁned on X (cid:44) S ∪ ( S × A ) .
Similarly, Tian et al.
(2020) also make this observation and propose to optimize a distance function by Q-learning on a collected set of trajectories.
The optimized distance function (i.e., Q-function) is then used with standard planning algorithms such as the Cross Entropy Method (CEM) (De Boer et al., 2005).
The speciﬁc model they used is an unconstrained network f : ( s, a, s (cid:48) )
→ R , outputting discounted distances (Q-values).
Due to the existing quasimetric structure, we explore using PQEs as the distance function formulation.
We mostly follow the algorithm in Tian et al.
except for the following minor differences:
• Tian et al.
(2020) propose to sample half of the goal from future steps of the same trajectory, and half of the goal from similar states across the entire dataset, deﬁned by a nearest neighbor search.
For simplicity, in the latter case, we instead sample a random state across the entire dataset.
• In Tian et al.
(2020), target goals are deﬁned as single states, and the Q-learning formulation only uses quantities distances from state-action pairs ( s,
a ) ∈ S
× A to states s (cid:48) : ˆ d (( s, a ) , s (cid:48) ) .
However, if we only train on ˆ d (( s, a ) , s (cid:48) ) , quasimetric embeddings might not learn much about the distance to state-action pairs, or from states, because it may simply only assign ﬁnite distances to ˆ d (( s, a ) , s (cid:48) ) , and set everything else to inﬁnite.
To prevent such issues, we choose to use state-action pairs as target goals, by adding a random action.
Then, the embedding methods only need to embed state-action pairs.
In planning when the target is actually a single goal s (cid:48) ∈ S , we use the following distance/Qfunction
[Equation]
Such a modiﬁcation is used for all embedding methods (PQEs, metric embeddings, asymmetrical dot products).
For unconstrained networks, we test both the original formulation (of using single state as goals) and this modiﬁcation.
Environment.
The environment is a grid-world with one-way doors, as shown in of Fig.
15, which is built upon gym-minigrid (Chevalier-Boisvert et al., 2018)
(a project under Apache 2.0 License).
The agent has 4 actions corresponding to moving towards 4 directions.
When it moves toward a direction that is blocked by a wall or an one-way door, it does not move.
States are represented as 18 -dimensional vectors, containing the 2D location of the agent (normalized to be within [ − 1 , 1] 2 ).
The other dimensions are always constant in our enviroment as they refer to information that can not be changed in this particular environment (e.g., the state of the doors).
The agent always starts at a random location in the center room (e.g., the initial position of the red triangle in Fig. 15).
The environment also deﬁnes a goal sampling distribution as a random location in one of the rooms on the left or right side.
Note that this goal distribution is only used for data collection and evaluation.
In training, we train goal-conditional policies using the goal sampling mechanism adapted from Tian et al.
(2020), as described above.
Training trajectories.
To collect the training trajectories, we use an (cid:15) -greedy planner with groundtruth distance toward the environment goal, with a large (cid:15) = 0 . 6 .
Each trajectory is capped to have at most 200 steps.
[Figure]
Architecture.
All encoder based methods (PQEs, metric embeddings, dot products) use 18-20482048-2048-1024 network with ReLU activations and Batch Normalization (Ioffe & Szegedy, 2015) after each activation, mapping a 18 -dimensional state to four 256 -dimensional latent vectors, corresponding to the embeddings for all four state-action pairs.
Unconstrained networks use a similar architecture and take in concatenated 36 -dimensional inputs.
With the original formulation with states as goals, we use a 36-2048-2048-2048-256-4 network to obtain a R |A| output, representing the distance/Q-values from each state-action pair to the goal; with the modiﬁed formulation with state-action pairs as goals, we use a 36-2048-2048-2048-256-16 network to obtain a R |A|×|A| output.
Training.
We use 1024 batch size with the Adam optimizer (Kingma & Ba, 2014), with learning rate decaying according to the cosine schedule without restarting (Loshchilov & Hutter, 2016) starting from 10 − 4 to 0 over 1000 epochs.
Since we are running Q-learning, all models are optimized w.r.t.
MSE on the γ -discounted distances, with γ = 0 .
When running with the triangle inequality regularizer, 341 ≈ 1024 / 3 triplets are uniformly sampled at each iteration.
Planning details.
To use the learned distance/Q-function for planning towards a given goal, we perform greedy 1-step planning, where we always select the best action in A according to the learned model, without any lookahead.
In each of 50 runs, the planner is asked to reach a goal given by the environment within 300 steps.
The set of 50 initial location and goal states is entirely decided by the seed used, regardless of the model.
We run each method 5 times using the same set of 5 seeds.
Full results.
Average results across 5 runs are shown in Fig.
15, with full results (with standard deviations) shown in Fig.
Planning performance across the formulations vary a lot, with PQEs and the Euclidean metric embedding being the best and most data-efﬁcient ones.
Using either formulation (states vs. state-action pairs as goals) does not seem to affect the performance of unconstrained networks.
We note that the the asymmetrical dot product formulation outputting discounted distance is similar to Universal Value Function Approximators (UVFA) formulation (Schaul et al., 2015); the unconstrained network outputting discounted distance with states as goals is the same formulation as the method from Tian et al.
[Figure] [Figure]