L EARNING TO E XTRAPOLATE : A T RANSDUCTIVE A PPROACH
Anonymous authors
Paper under double-blind review
A BSTRACT
Machine learning systems, especially overparameterized deep neural networks, can generalize to novel testing instances drawn from the same distribution as the training data.
However, they fare poorly when evaluated on out-of-support testing points.
In this work, we tackle the problem of developing machine learning systems that retain the power of overparametrized function approximators, while enabling extrapolation to out-of-support testing points when possible.
This is accomplished by noting that under certain conditions, a “transductive” reparameterization can convert an out-of-support extrapolation problem into a problem of within-support combinatorial generalization.
We propose a simple strategy based on bilinear embeddings to enable this type of combinatorial generalization, thereby addressing the out-of-support extrapolation problem.
We instantiate a simple, practical algorithm applicable to various supervised learning problems and imitation learning tasks.
I NTRODUCTION
Generalization is a central problem in machine learning.
Typically, one expects generalization when the test data is sampled from the same distribution as the training set, i.e out-of-sample generalization.
However, in many scenarios of interest, test data is sampled from a different distribution from the training set, i.e out-of-distribution ( OOD ).
In some OOD scenarios, the test-distribution is assumed to be known during training – a common assumption made by meta-learning methods.
Several works have tackled a more general scenario of “reweighted” distribution shift (Koh et al., 2021; Quinonero-Candela et al., 2008) where the test distribution shares support with the training distribution, but has a different and unknown probability density; this setting can be tackled via distributional robustness approaches (Sinha et al., 2018; Rahimian & Mehrotra, 2019).
The goal of our paper is to ﬁnd structural conditions under which we can generalize to test data which has support outside that of the train distribution.
Formally, assume the problem of learning function h :
ˆ y =
h θ ( x ) using data { ( x i , y i ) } N i =1 ∼ D train , where x i ∈ X train , the train domain.
We are interested in making accurate predictions h ( x ) for x / ∈ X train (see examples in Fig 1).
Consider an example task of predicting actions to reach a desired goal (Fig 1b).
During train, goals are provided from the blue cuboid ( x ∈ X train ) , but test time goals are from the orange cuboid ( x / ∈ X train ).
If f is modeled using a deep neural network, its predictions on test goals in the blue area are likely to be accurate, but in the orange area the performance can be arbitrarily poor unless further domain knowledge is incorporated.
This challenge manifests itself in a variety of real-world problems, ranging from supervised learning problems like object classiﬁcation (Barbu et al., 2019) to sequential decision making with reinforcement learning (Kirk et al., 2021) and imitation learning (de Haan et al., 2019).
Reliably deploying learning algorithms in unconstrained environments requires one to account for this type of “out-of-support” distribution shift.
It is widely accepted that if one can identify some structure in the training data that constrains the behavior of optimal predictors on novel data, then extrapolation may become possible.
Several methods can extrapolate if the nature of distribution shift is known apriori: convolution neural networks are appropriate if at test-time a training pattern appears at an out-of-distribution translation.
Similarly, accurate predictions can be made for object point clouds in out-of-support orientations by building in SE(3) equivariance (Deng et al., 2021b; Simeonov et al., 2022).
Another way to extrapolate is if the model class is known apriori: ﬁtting a linear function to a linear problem will
[Figure]
extrapolate.
Similarly, methods like NeRF (Zhang et al., 2022a) use physics of image formation to learn a 3D model of a scene which can synthesize images from novel viewpoints.
In this work, we propose an alternative structural condition under which out-of-support extrapolation is feasible.
Typical machine learning approaches are inductive : decision making rules are inferred from training data and employed for test-time predictions.
An alternative to induction is transduction (Gammerman et al., 1998) or analogy-making where a test example is compared with training examples to make a prediction.
Our main insight is that in the transductive view of machine learning, out-of-support extrapolation can be reparameterized as a combinatorial generalization problem, which, under certain low-rank and coverage conditions (Shah et al., 2020; Agarwal et al., 2021b;
Andreas et al., 2016; Andreas, 2019), admits a solution.
In this work we show how we can (i) re-parameterize out-of-support inputs h ( x test )
→ h (∆ x, x (cid:48) ) ,
where x (cid:48) ∈ X train , when provided a representation of measure of difference ∆ x between x test and x (cid:48) .
We then (ii) Provide conditions under which h (∆ x, x (cid:48) ) makes accurate predictions for unseen combinations of (∆ x, x (cid:48) )
based on a theoretically justiﬁed bilinear modeling approach: h (∆ x, x (cid:48) )
→ f (∆ x ) T g ( x (cid:48) ) , where f and g map their inputs into two vector spaces of the same dimension.
Finally, (iii) empirical results demonstrate generality of extrapolation of our algorithm on a wide variety of tasks: (a) regression for analytical functions and high-dimensional point cloud data; (b) sequential decision making tasks such as goal-reaching for a simulated robot.
S ETUP
Notation.
Given a space of inputs X and targets Y , we aim to learn a predictor h θ : X → P ( Y )which best ﬁts a ground truth function h (cid:63) :
X → Y .
Given some non-negative loss function (cid:96) :
Y × Y → R ≥ 0 on the outputs (e.g., square loss), and a distribution D over X , risk is deﬁned as
[Equation]
Various choices of train ( D train ) and test ( D test ) distributions yield different generalization settings:
In-Distribution Generalization.
This setting assumes D test =
D train .
The challenge is to ensure that with N samples from D train , the expected risk R ( h θ ; D test )
=
R ( h θ ; D train ) is small.
This is a common paradigm in both empirical supervised learning (e.g. Simonyan & Zisserman (2014)) and in standard statistical learning theory (e.g. Vapnik (2006)).
Out-of-Distribution ( OOD ).
This is more challenging and requires accurate predictions when
D train (cid:54) =
D test .
When the ratio between the density function of D test to that of D train is bounded, rigorous OOD extrapolation guarantees exist and are detailed in Appendix B.1.
Such a situation arises when D test shares support with D train but is differently distributed as depicted in Fig 2a.
Out-of-Support ( OOS ).
There are innumerable forms of distribution shift in which density ratios are not bounded.
The most extreme case is when the support of D test is not contained in that of D train .
I.e., when there exists some X (cid:48)
⊂
X such that P x ∼D test [ x ∈ X (cid:48) ] > 0 , but P x ∼D train [ x ∈ X (cid:48) ]
=(see Fig 2b).
We term the problem of achieving low risk on such a D test as OOS extrapolation.
Out-of-Combination ( OOC ).
This is a special case of OOS .
Let X =
X 1 ×
X 2 be the product of two spaces.
Let D train , X 1 , D train , X 2 denote the marginal distributions of x 1 ∈ X 1 , x 2 ∈ X 2 under D train , and D test , X 1 , D test , X 2 under D test .
In OOC learning, D test , X 1 , D test , X 2 are in the support of D train , X 1 , D train , X 2 , but the joint distributions D test need not be in the support of D train .
[Figure]
R EDUCING O UT - OF -S UPPORT TO O UT - OF -C OMBINATION
While general OOS extrapolation can be arbitrarily challenging, we shall argue that there are at least some concrete conditions under which OOC extrapolation is feasible.
Motivated by this insight, we (a) explain how to use transductive predictors to convert OOS prediction to OOC prediction , and (b) describe how bilinear predictors can extrapolate for particular classes of OOC problems (Section 3.2).
Therefore, by converting OOS problems to OOC ones, we are able to use bilinear predictors to enable extrapolation to out-of-support test points.
3. T RANSDUCTIVE P REDICTORS : C ONVERTING OOS TO OOC
To convert OOS to OOC , we require the input space X to have group structure, i.e. there are addition and subtraction operators such that x + x (cid:48) , x − x (cid:48) are well-deﬁned for x, x (cid:48) ∈ X .
Let ∆ X := { x − x (cid:48) :
x, x (cid:48) ∈ X} .
We propose a transductive re-parameterization h θ : X → P ( Y ) with a deterministic function ¯ h θ : ∆ X × X → Y as
[Equation]
where x (cid:48) is referred to as an anchor point for a query point x .
Under this re-parameterization, the training distribution can be rewritten as a joint distribution of train ∆ x = x − x (cid:48) and x (cid:48) as follows
[Equation]
This is just representing the prediction for every point from the training distribution in terms of it’s relationship to other points in the training distribution.
At test time, we are presented with point x ∼ D test that may be from an OOS distribution.
To make a prediction on this OOS x , we make the observation that with a careful selection of an anchor point x (cid:48) , our reparameterization may be able to convert this OOS problem into a more manageable OOC one, since representing the test point x in terms of it’s difference from training points can still be an “in-support” problem.
To do so, we select an anchor point x (cid:48) from D train as follows.
For a radius parameter ρ > 0 , deﬁne the distribution of chosen anchor points D trns ( x ) (referred to as a transducing distribution) as
[Equation]
where X train denotes the set of x in our training set, and we denote ∆ X train :
= { x 1 − x 2 : x 1 , x 2 ∈
X train } .
Intuitively, our choice of D trns ( x ) selects anchor points x (cid:48) to transduce from the training distribution, subject to the resulting differences ( x − x (cid:48) ) being close to a “seen” ∆ x ∈ ∆ X train .
In doing so, both the anchor point x (cid:48) and the difference ∆ x have been seen individually at training time, albeit not in combination.
This allows us to express the prediction for a OOS query point in terms of an in-support anchor point x (cid:48) and an in-support difference ∆ x
(but not jointly in support).
This choice of anchor points induces a joint test distribution of ∆ x = x − x (cid:48) and x (cid:48) :
[Equation]
As seen from Fig 3, the marginals of ∆ x and x (cid:48) under ¯ D test , are individually in the support of those under ¯ D train .
Still, as Fig 3 reveals, since x test is out-of-support, the joint distribution of ¯ D test is not covered by that of ¯ D train (i.e. the combination of x 1 and x test have not been seen together before); precisely the OOC regime.
Moreover, if one tried to transduce all x (cid:48) ∼ D train to x ∼ D test at test time (e.g. transduce point x 3 to x test in the ﬁgure) then we would lose coverage of the ∆ x -marginal.
By doing transduction to keep both the marginal x (cid:48) and ∆ x in support, we are ensuring that we can convert difﬁcult OOS problems into (potentially) more manageable OOC ones.
3. B ILINEAR REPRESENTATIONS FOR OOC LEARNING
Without additional assumptions, OOC extrapolation may be just as challenging as OOS .
However, with certain low-rank structure it can be feasible (Shah et al., 2020; Agarwal et al., 2021b;
Athey et al., 2021).
Following Agarwal et al.
(2021a), we recognize that this low-rank property can be leveraged implicitly for our reparameterized OOC problem even in the continuous case (where x, ∆ x do not explicitly form a ﬁnite dimensional matrix) using a bilinear representation of the transductive predictor in Eq.
(3.1), ¯ h θ (∆ x, x (cid:48) )
= (cid:104)
f θ (∆ x ) , g θ ( x (cid:48) )
(cid:105) .
Here f θ , g θ map their respective inputs into a vector space of the same dimension, R p .
If the output space is K dimensional, then we independently model the prediction for each dimension using a set of K bilinear embeddings:
[Equation]
While ¯ h θ ,k are bilinear in embeddings f θ ,k , g θ ,k , the embeddings themselves may be parameterized by general function approximators.
The effective “rank” of the transductive predictor is controlled by the dimension of the continuous embeddings f θ ,k (∆ x ) , g θ ,k ( x (cid:48) ) .
We now introduce three conditions under which extrapolation is possible.
As a preliminary, we write µ 1 (cid:28) κ µ 2 for two (unnecessarily normalized) measures µ 1 , µ 2 to denote µ 1 ( A ) ≤ κµ 2 ( A ) for all events A .
This notion of coverage has been studied in the reinforcement learning community (termed “concentrability” (Munos & Szepesv´ari, 2008)), and implies that the support of µ 1 is contained in the support of µ 2 .
The ﬁrst condition we need is the following notion of combinatorial coverage.
Assumption 3.1 (Bounded combinatorial density ratio) .
We assume that ¯ D test has κ -bounded combinatorial density ratio with respect to ¯ D train , written as ¯ D test (cid:28) κ, comb ¯ D train .
This means that there exist distributions D ∆ X ,i and D X ,j , i, j ∈ { 1 , 2 } , over ∆ X and X , respectively, such that
D i ⊗ j :=
D ∆ X ,i ⊗ D X ,j satisfy
[Equation]
Assumption 3.1 can be interpreted as a generalization of a certain coverage condition in matrix factorization.
The following distribution is elaborated upon in Appendix B.5 and Fig Figure 9b therein.
Intuitively, the distribution D 1 ⊗ 1 corresponds to the top-left block of a matrix, where all rows and columns are in support.
D 1 ⊗ 2 and D 2 ⊗ 1 correspond to the off-diagonals, and D 2 ⊗ 2 corresponds to the bottom-right block.
The condition requires that ¯ D train covers the entire top-right block and off-diagonals, but need not cover the bottom-right block D 2 ⊗ 2 .
¯ D test , on the other hand, can contain samples from this last block.
The idea is that, for problems with appropriate low-rank structure, samples from the top-left block, and samples from the off-diagonal blocks, uniquely determine the bottom-right block.
Recall that p denotes the dimension of the parametrized embeddings in Eq.
Our next assumptions are that the ground-truth embedding is low-rank (we refer to as “bilinearly transducible”), and that under the “top-block” distribution D 1 ⊗ 1 , its component factors do not lie in a strict subspace of R p , i.e. not degenerate.
Assumption 3.2 (Bilinearly transducible) .
For each component k ∈ [ K ] , there exists f (cid:63),k :
∆ X → R p and g (cid:63),k :
X → R p such that for all x ∈ X , the following holds with probability 1 over x (cid:48) ∼ D trns ( x ) : h (cid:63),k ( x ) =
¯ h (cid:63),k (∆ x, x (cid:48) ) :
= (cid:104) f
(cid:63),k (∆ x ) , g (cid:63),k ( x (cid:48) )
(cid:105) , where ∆ x = x − x (cid:48) .
Further, the ground-truth predictions are bounded: max k sup ∆ x,x (cid:48) | ¯ h (cid:63),k (∆ x, x (cid:48) )
| ≤ M .
Assumption 3.3.
For all components k ∈ [ K ] , the following lower bound holds for some σ 2 > 0 :
[Equation]
The following is proven in Appendix B; due to limited space, we defer further discussion of the bound to Appendix B.4.1.
Theorem 1.
Suppose Assumptions 3.1 to 3.3 all hold, and that the loss (cid:96) ( · , · ) is the square loss.
Then, if the training risk satisﬁes R ( h θ ; D train ) ≤ σ 2 4 κ , then the test risk is bounded as
[Equation]
3. O UR P ROPOSED A LGORITHM : B ILINEAR T RANSDUCTION
Our basic proposal for OOS extrapolation is unweighted transduction , depicted in Algorithm 1: at train time, a predictor ¯ h θ is trained to make predictions for train points x i drawn from the training set D train based on their similarity with other points x j also drawn from D train :
¯ h θ ( x i − x j , x j ) .
The train pairs x i , x j are sampled uniformly from D train .
At test time, for an OOS point x test , we ﬁrst select an anchor point x i from the train set which has similarity with the test point x test − x i that is within some radius ρ of the train similarities distribution ∆ X train .
We then predict the value for x test based on the anchor point x i and the similarity of the test and anchor points: ¯
h θ ( x test − x i , x i ) .
For the supervised regression setting, we compute differences directly between inputs x i , x j ∈ X .
For the goal-conditioned imitation learning setting, we compute difference between states x i , x j ∈
X sampled uniformly over demonstration trajectories.
At test time, we select an anchor trajectory based on the goal, and transduce each anchor state in the anchor trajectory to predict a sequence of actions for a test goal.
As explained in Appendix C.1, there may be scenarios where it is beneﬁcial to more carefully choose which points to transduce.
To this end, we outline a more sophisticated proposal, weighted transduction , that achieves this.
Pseudo code for unweighted variants is depicted in Algorithm 1, and weighted in Algorithm 2 in the Appendix.
Note that we assume access to a representation of X for which the subtraction operator captures differences that occur between train points and between train and test points.
In Appendix D we describe examples of such representations as deﬁned in standard simulators or which we extract with dimensionality reduction techniques (PCA).
Alternatively, one might potentially extract such representations via self-supervision Kingma & Welling (2013) and pre-trained models Upchurch et al.
(2017), or compare data points via another similarity function.
E XPERIMENTS
We answer the following questions through an empirical evaluation: (1) Does reformulating OOS extrapolation as a combinatorial generalization problem allow for extrapolation in a variety of supervised and sequential decision making problems?
(2) How does the particular choice of training distribution and data-generating function affect performance of the proposed technique?
(3) How important is the choice of the low-rank bilinear function class for generalization?
(4) Does our method scale to high dimensional state and action spaces?
We ﬁrst analyze OOS extrapolation on analytical domains and then on more real world problem domains depicted in Fig 8.
4. A NALYZING OOS EXTRAPOLATION ON ANALYTIC PROBLEMS
We compare our method on regression problems generated via 1-D analytical functions (described in Appendix D.1) against standard neural networks with multiple fully-connected layers trained and tested in ranges [20 , 40] and [10 , 20] ∪ [40 , 50] (with the exception of Fig 4c trained in [ − 1 , 1] and tested in [ − 1 . 6 , − 1] ∪ [1 , 1 . 6] ).
We use these domains to gain intuition about the following questions:
What types of problems satisfy the assumptions for extrapolation?
While we outlined a set of conditions under which extrapolation is guaranteed, it is not apparent which problems satisfy these assumptions.
To understand this better, we considered learning functions with different structure: a periodic function with mixed periods (Fig 4a), a sawtooth function (Fig 4b) and a polynomial function (Fig 4c).
Standard deep networks (yellow) ﬁt the training points well (blue), but fail to extrapolate to OOS inputs (orange).
In comparison, our approach (pink) accurately extrapolates on periodic functions but is much less effective on polynomials.
This is because the periodic functions have symmetries which induce low rank structure under the proposed re-parameterization.
[Figure]
Going Beyond Known Inductive Biases Given our method extrapolates to periodic functions in
Fig 4 (which displays shift invariance), one might argue a similar extrapolation can be achieved by building in an inductive bias for periodicity / translation invariance.
In Fig 6, we show that bilinear transduction in fact is able to extrapolate even in cases that the ground truth function is not simply translation invariant, but is translation equivariant, showing that bilinear transduction can capture equivariance.
Moreover, bilinear transduction can in fact go beyond equivariance or invariance.
To demonstrate broader generality of our method we consider a piecewise periodic function that also grows in magnitude (Fig 7).
This function is neither invariant nor equivariant to translation, as the group symmetries in this problem do not commute.
The results in Fig 7 demonstrate that while the baselines fail to do so (including baselines that bake in equivariance (green)), bilinear transduction successfully extrapolates.
The important thing for bilinear transduction to work is when comparing training instances, there is a simple (low-rank) relationship between how their labels are transformed.
While it can capture invariance and equivariance, as Fig 7 shows, it is more general.
How does the relationship between the training distribution and testing distribution affect extrapolation behavior?
We try and understand how the range of test points that can be extrapolated to, depends on the training range.
We show in Fig 5 that for a particular “width” of the training distribution (size of the training set), OOS extrapolation only extends for one “width” beyond the training range since the conditions for ∆ X being in-support are no longer valid beyond this point
4. A NALYZING OOS EXTRAPOLATION ON LARGER SCALE DECISION MAKING PROBLEMS
To establish that our method is useful for complex and real-world problem domains, we varied the complexity (i e., working with high-dimensional observation and action spaces) and the learning setting (regression and sequential decision making).
Baselines: To assess the effectiveness of our proposed scheme for extrapolation via transduction and bilinear matrix completion, we compare with the following non-transductive baselines.
Linear Model : linear function approximator to understand whether the issue is one of overparameterization, and whether linear models would solve the problem.
Neural Networks : typical training of overparameterized neural network function approximator via standard empirical risk minimization.
Alternative Techniques with Neural Networks (DeepSets) : an alternative architecture for combining multiple inputs (DeepSets Zaheer et al.
(2017)), that are meant to be permutation invariant and encourage a degree of generalization between different pairings of states and goals.
Finally, we compare with a
Transductive Method without a mechanism for Structured Extrapolation (Transduction) : transduction with no special structure for understanding the impact of bilinear embeddings and low rank structure.
This baseline uses reparameterization, but just parameterizes h θ as a standard neural network.
We present additional comparison results introducing periodic activations Tancik et al.
(2020) in Table 4 in the Appendix.
OOS Extrapolation in Sequential Decision Making : We tested different types of extrapolation, with results reported in Table 1.
More details about training and evaluation are provided in Appendices D.1 and D.2.
• Extrapolation to OOS Goals: We considered two tasks from the Meta-World benchmark (Yu et al., 2020) where a simulated robotic agent needs to either reach or push a target object to a goal location (column 2 in Fig 8).
Given a set of expert demonstrations reaching/pushing to goals in the blue box ( [0 , 0 . 4] ×
[0 . 7 , 0 . 9] × [0 . 05 , 0 . 3] and [0 , 0 .
3] × [0 . 5 , 0 .
7] × { 0 . 01 } ), we tested generalization to OOS goals in the orange box ( [ − 0 .
4 , 0] × [0 . 7 , 0 . 9] × [0 . 05 , 0 . 3] and [ − 0 .
3 , 0] × [0 . 5 , 0 .
7] ×{ 0 . 01 } ), using a simple extension of our method described in Appendix D.2 to perform transduction over trajectories rather than individual states.
We quantify performance by measuring the distance between the conditioned and reached goal.
Results in Table 1 show that on the easy task of reaching, training a linear or typical neural network based predictor extrapolate as well as our method.
However, for the more challenging task of pushing an object, our extrapolation is better by an order of magnitude than other baselines, showing the ability to generalize to goals in a completely different direction.
• Extrapolation with large state and action space: Next we tested our method on grasping and placing an object to OOS goal-locations in R 3 with an anthropomorphic “Adroit” hand with a much larger action ( R 30 ) and state ( R 39 ) space (column 3 in Fig 8).
Results conﬁrm that bilinear transduction scales up to high dimensional state-action spaces as well and is naturally able to grasp the ball and move it to new target locations ( [ − 0 .
3 , 0] × [ − 0 .
3 , 0] × [0 . 15 , 0 . 35] ) after trained on target locations in ( [0 , 0 . 3] × [0 , 0 . 3] × [0 . 15 , 0 . 35] ).
These results are better appreciated in the video attached with the supplementary material, but show the same trends as above, with bilinear transduction signiﬁcantly outperforming standard inductive methods and non-bilinear architectures.
• Extrapolation to OOS Dynamics: Lastly, we consider problems involving extrapolation not just in location of the goals for goal-reaching problems, but across problems with varying dynamics.
Speciﬁcally, we consider a slider task where the goal is to move a slider on a table to strike a ball such that it rolls to a ﬁxed target position (column 4 in Fig 8).
The mass of the ball varies across episodes and is provided as input to the policy.
We train and test on a range of masses ( [60 , 130] and [5 , 15] ).
We ﬁnd that bilinear transduction adjusts behavior and successfully extrapolates to new masses, showing the ability to extrapolate not just to goals, but also to varying dynamics.
Importantly, bilinear transduction is signiﬁcantly less prone to variance than standard inductive or permutation invariant architectures.
This can be
seen
from a heatmap over various training architectures and training seeds as shown in Fig 12 and Table 5 in Appendix C.4.
While other methods can sometimes show success, only bilinear transduction is able to consistently accomplish extrapolation
OOS Extrapolation in Higher Dimensional Regression Problems .
To scale up the dimension of the input space, we consider the problem of predicting valid grasping points in R 3 from point clouds of various objects (bottles, mugs and teapots) with different orientations, positions and scales (column 1 in Fig 8).
At train and test, objects undergo z -axis orientation ( [0 , 1 . 2 π ] and [1 . 2 π, 2 π ] ), translation ( [0 , 0 . 5] × [0 , 0 . 5] and [0 . 5 , 0 .
7] × [0 . 5 , 0 .
7] ) and scaling ( [0 . 7 , 1 . 3] and [1 . 3 , 1 . 6] ).
In this domain, we represent entire point clouds by a low dimensional representation of the point cloud obtained via PCA.
We consider situations where we train on individual objects where the train set consists of various rotations, translations and scales and standard unweighted bilinear transduction is applied.
We also consider situations where the objects are not individually identiﬁed but instead a single grasp point predictor is trained on the entire set of bottles, mugs and teapots.
We assume access to category labels at training time, but do not require this at test time.
For a more in-depth discussion of this assumption and a discussion of how this can be learned via weighted -transduction, we refer readers to Appendix C.1.
The tests show that while training performance is comparable on all instances, we ﬁnd that extrapolation behavior is signiﬁcantly better for our method.
This is true both for single object cases as well as scenarios where all objects are mixed together (Table 1).
These experiments show that bilinear transduction can work on feature spaces of high dimensional data such as point clouds.
Note that while we only predict a single grasp point here, a single grasp point prediction can be easily generalized to predict multiple keypoints instead (Manuelli et al., 2019), enabling success on more challenging control domains.
For further visualizations and details, please see Appendix C and D.
A BRIDGED R ELATED W ORK
Generalization from train data to test data of the same distribution has been studied extensively, both practically (e.g. Simonyan & Zisserman (2014)) and theoretically Vapnik (2006); Bousquet &
[Table]
Elisseeff (2002); Bartlett & Mendelson (2002).
Our focus in on performance on distributions with examples which may not be covered by the training data, as described formally in Section 2.
Here we provide a discussion of the most directly related approaches, and defer an extended discussion to Appendix A. Even more so than generalization, extrapolation necessitates leveraging structure in both the data and the learning algorithm.
Along these lines, past work has focused on structured neural networks , which hardcode such symmetries as equivariance (Cohen & Welling, 2016;
Simeonov et al., 2022), Euclidean symmetry Smidt (2021) and periodicity (Abu-Dakka et al., 2021;
Parascandolo et al., 2016) into the learning process.
Other directions have focused on general architectures which seem to exhibit combinatorial generalization, such as transformers (Vaswani et al., 2017), graph neural networks (Cappart et al., 2021) and bilinear models (Hong et al., 2021).
In this work, we focus on bilinear architectures with what we term a transductive parametrization.
We demonstrate that this framework can in many cases learn the symmetries (e.g. equivariance and periodicity) that structured neural networks harcode for, and achieve extrapolation in some regimes in which these latter methods cannot.
A bilinear model with low inner dimension is equivalent to enforcing a low-rank constraint on one’s predictions.
Low-rank models have commanded broad popularity for matrix completion (Mnih & Salakhutdinov, 2007; Mackey et al., 2015).
And whereas early analysis focused on the missing-at-random setting (Cand ` es & Recht, 2009; Recht, 2011; Candes & Plan, 2010) equivalent to classical in-distribution statistical learning, we adopt more recent perspectives on missing-not-at-random data, (Shah et al., 2020; Agarwal et al., 2021b;
Athey et al., 2021), which tackle the combinatorial-generalization setting described in Section 3.2.
In the classical transduction setting (Joachims, 2003; Gammerman et al., 1998; Cortes & Mohri, 2006), the goal is to make predictions on a known set of unlabeled test examples for which the features are known; this is a special case of semi-supervised learning.
We instead operate in a standard supervised learning paradigm, where test labels and features are only revealed at test time.
Still, we ﬁnd it useful to adopt a “transductive parametrization” of our predictor, where instead of compressing our predictor into parameters alone, we express predictions for labels of one example as a function of other, labeled examples.
An unabridged version of related work is in Appendix A.
D ISCUSSION
In this work, we consider the problem of out-of-support extrapolation in regression and sequential decision making problems.
We show that under a set of assumptions, extrapolation problems can be reparameterized using transduction to be viewed as combinatorial generalization problems.
This allows us to leverage techniques from low-rank matrix completion in order to solve the combinatorial generalization problem.
While our work serves as an initial study of the circumstances under which problem structure can be both discovered and exploited for understanding extrapolation, there are a number of natural questions for further research.
First, can we classify which set of real-world domains ﬁt our assumptions, beyond the domains we have demonstrated?
Second, can we learn a latent space in which differences ∆ x are meaningful for high dimensional domains?
And lastly, are there more effective schemes for selecting anchor points?
E THICS S TATEMENT
Bias In this work on extrapolation via bilinear transduction, we can only aim to extrapolate to data that shifts from the train set in ways that exist within the train data.
Therefore the performance of this method relies on the diversity within the train data, allowing some forms of extrapolation but not to ones that do not exist within the train distribution.
Dataset release We provide in detail in the Appendix the parameters and code bases and datasets we use to generate our data.
In the future, we plan to release our code base and expert policy weights with which we collected expert data.
R EPRODUCIBILITY S TATEMENT
We describe our algorithms in Section 3.3 and the complete proof of our theoretical results and assumptions in Appendix B. Extensive implementation details regarding our algorithms, data, models, and optimization are provided in Appendix D.
In addition, we plan in the future to release our code and data.
C ONTENTS 1 Introduction
2 Setup
Reducing Out-of-Support to Out-of-Combination 3.
Transductive Predictors: Converting OOS to OOC . . . . . . . . . . . . . . . . . .
Bilinear representations for OOC learning . . . . . . . . . . . . . . . . . . . . . .
Our Proposed Algorithm: Bilinear Transduction . . . . . . . . . . . . . . . . . . .
4 Experiments 4.
Analyzing OOS extrapolation on analytic problems . . . . . . . . . . . . . . . . .
Analyzing OOS extrapolation on larger scale decision making problems . . . . . .
5 Abridged Related Work 6 Discussion 7 Ethics Statement 8 Reproducibility Statement
A Unabridged Related Work
B Theoretical Results
B.1 Generalization under bounded density ratio . .
B.2 Extrapolation for Matrix Completion . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 General Analysis for Combinatioral Extrapolation . . . . . . . . . . . . . . . . . .
B.3.
Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 Extrapolation for Transduction . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.4.
Further Remarks on Theorem 1 . . . . . . . . . . . . . . . . . . . . . . .
B.5 Connections to matrix completion . . . . . . . . . . . . . . . . . . . . . . . . . .
C Additional Results
C.1 Learning which points to transduce: Weighted Transduction
C.2 Additional Results on a Simple 2-D Example
C.3 Additional Results on Supervised Grasp prediction
C.4 Additional Results on Imitation Learning
C.5 Choice of Architecture
D Implementation Details
D.1 Train/Test distributions and Analytic Function Details
D.2 Training details . .
D.2.1 Details of Bilinear Transduction Implementation . . . . . . . . . . . . . . D.2.2 Details of Training Distributions . . . . . . . . . . . . . . . . . . . . . . . D.2.3 Details of Model Architectures . . . . . . . . . . . . . . . . . . . . . . . .
D.2.4 Details of Model Optimization . . . . . . . . . . . . . . . . . . . . . . . .
A U NABRIDGED R ELATED W ORK
Here we discuss various prior approaches to extrapolation in greater depth, focusing on areas not addressed in our abridged discussion in the main text.
Approaches which encode explicit structure.
One popular approach to designing networks which extrapolate to novel data has been the equivariant neural networks , ﬁrst proposed by Cohen & Welling (2016).
The key idea is that, if it is known there is a group G which acts on both the input and target domains of the predictor, and if it is understood that the true predictor must satisfy the equivariance property
[Equation]
for all g ∈ G (here g · () denotes group action), then one can explicitly encode for predictors h θ satisfying the same identity.
Similarly, one can encode for invariances, where h (cid:63) ( g · x ) = h (cid:63) ( x ) for all g ∈ G .
Deng et al.
(2021a) extended the original equivariance framework to accommodate situations where the group G corresponds to rotations ( SO (3) ), and Simeonov et al.
(2022) proposed neural descriptor ﬁelds which handle rigid-body transformations SE (3) ).
For a broader review on how other notions of symmetry can be encoded into machine learning settings, consult Smidt (2021), and AbuDakka et al.
; Parascandolo et al.
for how periodic structure can be explicitly built in.
We remark that in many of these approaches, the group/symmetry must be known ahead of time.
While attempt to learn global (Benton et al., 2020) or local (Dehmamy et al., 2021) equivariances/invariances, there are numerous forms of structure that can be represented as a group symmetry, and for which these methods do not apply.
As we show in our experiments, there are a number of group structures which are not captured by equivariance (such as Fig 7), which bilinear transduction captures.
Architectures favorable to extrapolation.
There are numerous learning architectures which are purported to be favorable to extrapolation to novel-domains.
For example, graph neural networks (GNNs) have been use to facilitate reasoning behavior in combinatorial environments (Battaglia et al., 2018; Cappart et al., 2021), and the implicit biases of GNNs have received much theoretical study in recent years (see Jegelka (2022) and the references therein).
Another popular domain for extrapolations has been sequence-to-sequence modeling, especially in the context of natural language processing.
Here, the now-renowned Transformer architecture due to Vaswani et al.
(2017), as well as its variants based on the same “attention mechanism” (e.g. Kitaev et al.
(2020)) have become incredibly popular.
Supervised by a sufﬁciently diverse set of tasks and staggering amount of data, these have shown broad population in various language understanding tasks (Devlin et al., 2018), text-to-image generation (Ramesh et al., 2021), and even quantitative reasoning (Lewkowycz et al., 2022).
Attention-based models tend to excel best in tasks involving language and when trained on massively large corpora; their ability to extrapolate in manipulation tasks (Zhou et al., 2022) and in more modest data regimes remains an area of ongoing research.
Moreover, while recent research has aimed to study their (in-distribution) generalization properties (Edelman et al., 2022), their capacity for “reasoning” more generally remains mysterious (Zhang et al., 2022b).
Lastly, a line of research has studied various bilinear models (e.g. Hong et al.
; Shah et al.
(2020)), motivated by the extension of literature on matrix factorization discussed in our abridged related work.
However, these methods require additional ﬁne tuning when applied to novel goals and do not display zero shot extrapolation.
Moreover, Shah et al.
(2020) requries a discretization of the state-action space in order to formulate the matrix completion problem, whereas bilinear transduction shows how the results hold in continuous bilinear form, without the need for discretization.
Along similar lines, the Deep-Sets architecture of Zaheer et al.
(2017) aims for combinatorial extrapolation by embedding tokens of interest in a latent vector space on which addition operators can be deﬁned.
Zhou et al.
(2022) compares the Deep Sets and Transformers approaches as policy architectures in reinforcement learning, ﬁnding that neither architecture uniformly outperforms the other.
Distributional Robustness.
One popular approach to out-of-distribution learning has been distributional robustness, which seeks predictors that perform on a family of “nearby” shifted distributions
Sinha et al.
; Rahimian & Mehrotra (2019).
These approaches are well suited to OOD settings where the test distributions have the same support, but have differing (but boundedly different) probability densities.
Meta-Learning and Multi-Task Learning: Meta-learning and multi-task learning methods aim to learn policies/predictors that can either zero-shot or very quickly generalize to new problems Thrun & Pratt (1998);
Finn et al.
; Santoro et al.
; Vinyals et al.
; Kalashnikov et al.
; Crawshaw (2020);
Rosenbaum et al.
; Caruana (1997).
These methods however usually make strict distributional assumptions about the train and test distribution of tasks being the same.
So while the particular data distribution may be different, the meta-level assumption is still an in-distribution one Fallah et al.
Yin et al.
(2019) shows that under distribution shift, meta-learning methods can fail dramatically, and this problem will be exacerbated when supports shift as well.
Multi-task learning methods often make the assumption that training on some set of training tasks can provide good pre-training before ﬁnetuning on new tasks Julian et al.
Su et al.
; Meftah et al.
However, it is poorly understand how these tasks actually relate to each other and the study has been largely empirical.
Generalization in Reinforcement Learning and Imitation Learning:
Reinforcement learning problems are unique in that they are sequential and interactive.
Distribution shift may not just be across different MDPs but within the states of a single MDP as well Filos et al.
; Kirk et al.
; Ghosh et al. (2021).
While we do not explicitly focus on the reinforcement learning setting, our results are largely in the imitation learning setting.
Moreover, we are not looking at the ﬁnetuning setting but rather evaluating zero-shot performance on out of support goals, dynamics and arbitrary contexts.
While certain works Simeonov et al.
; Wang et al.
(2022) do explore ideas of equivariance and invariance in RL and imitation learning, they are largely restricted to the either the SE (3) or SO (2) groups.
Moreover, as we discussed in our experiments, bilinear transduction is able to capture more complex notions than equivariance.
The empirical evaluations we perform are in a large part on domains in goal conditioned reinforcement learning.
Most goal conditioned RL works do not consider extrapolation beyond the range of training goals, but rather standard statistical generalization amongst training goals Kaelbling (1993);
Andrychowicz et al.
; Nair et al.
While certain architectures Hong et al.
(2021) do seem to show some level of extrapolation, they still require signiﬁcant ﬁnetuning.
We show that with low-rank assumptions under reparameterization, zero-shot extrapolation is possible.
Our work can also loosely be connected to analogy making techniques in machine learnig Reed et al.
; Ichien et al.
; Hansen-Estruch et al. (2022); Mitchell (2021).
We formalize this notion in the context of out-of-support extrapolation and show that transduction allows for analogy making under low-rank structure assumptions.
B T HEORETICAL R ESULTS
This appendix is organized into four parts.
• Appendix B.1 introduces the bounded density condition, and reproduces a folklore guarantee for extrapolation when one distribution has bounded density with respect to the other.
• Appendix B.2 provides a basic guarantee for extrapolation in the context of missing-notat-random matrix completion depicted in Fig 9, based on results in Shah et al.
• Appendix B.3 formalizes a notion of bounded combinatorial density ratio (Deﬁnition B.2), in terms of which we can establish out-of-combination extrapolation guarantees by leveraging Appendix B.2.
Finally, Appendix B.4 applies the results of the previous section, stating and proving formal guarantees for our proposed transductive predictors.
B. G ENERALIZATION UNDER BOUNDED DENSITY RATIO
The following gives a robust, quantitative notion of when one distribution is in the support of another.
For generality, we state this condition in terms of general positive measures µ 1 , µ 2 , which need not be normalized and sum to one.
Deﬁnition B.1 ( κ -bounded density ratio) .
Let µ 1 , µ 2 be two measures over a space Ω .
We say µhas κ -bounded density with respect to µ 2 , which we denote µ 1 (cid:28) κ µ 2 , if for all measurable event
A ⊂ Ω , µ 1 [ A ] ≤ κµ 2 [ A ] .
Stating Deﬁnition B.1 for general probability affords us the ﬂexibility to write example, P 1 (cid:28) κ P 2 + P 3 , as P 2 + P 3 is a nonnegative measure with total mass 1 + 1 = 2 .
Remark B.1 (Connection to concentrability) .
The parameter κ is known in the off-policy reinforcement learning literature as the concentratability coefﬁcient (see, e.g. Munos & Szepesv´ari (2008)), and appears in controlling the performance of a rollout policy π 1 trained on a behavior policy π 2 by asserting that P π 1 (cid:28) κ P
π 2 , where P π i is, for example, a state-action visitation probability under π i .
Remark B.2 (Density nomenclature) .
The nomenclature “density” refers to the alternative deﬁnition in terms of Radon-Nikodym derivatives.
To avoid technicalities, this can be best seen when µand µ 2 are continuous densitives over Ω
= R d with densities p 1 ( · ) and p 2 ( · ) ; e.g. for i ∈ { 1 , 2 } .
µ i
A ] = (cid:82) x ∈
A p i ( x )d x .
Then µ 1 (cid:28) κ µ 2 is equivalent to
[Equation]
The following lemma motivates the use of Deﬁnition B.1.
Its proof is standard but included for completeness.
Lemma B.1.
Let µ 1 , µ 2 be measures on the same measurable space Ω , and suppose that µ 2 (cid:28) κ µ 1 .
Then, for any nonnegative function φ , µ 2 [ φ ] ≤ κµ 1 [ φ ] .
4 In particular, if D test (cid:28) κ D train , then as long as our loss function (cid:96) ( · , · ) is nonnegative,
[Equation]
Thus, up to a κ -factor, R ( h θ ; D test ) inherits any in-distribution generalization guarantees for
R ( h θ ; D train ) .
Proof.
As in standard measure theory (c.f. C¸inlar (2011, Chapter 1)), we can approximate any φ ≥by a sequence of simple functions φ n ↑ φ , where φ n ( ω )
=
(cid:80) k n
i =1 c n,i I { ω ∈
A n,i } , with A n,i ⊂ Ω and c n,i ≥ 0 .
For each φ n , we have
[Equation]
c
[Equation]
µ 1 [ φ n ] .
The result now follows from the monotone convergence theorem.
To derive the special case for
D test and D train , apply the general result with nonnegative function φ ( x ) =
E y ∼ h θ ( x ) (cid:96) ( y, h (cid:63) ( x ))
(recall (cid:96) ( · , · ) ≥ 0 by assumption), µ 1 = D train and µ 2 = D test
B. E XTRAPOLATION FOR M ATRIX C OMPLETION
In what follows, we derive a simple extrapolation guarantee for matrix completion.
The following is in the spirit of the Nystr¨om column approximation (see e.g. Gittens & Mahoney (2013)), and our proof follows the analysis due to Shah et al.
Throughout, consider
[Equation]
where we decompose ˆ M , M (cid:63) into blocks ( i, j ) ∈ { 1 , 2 } 2 for dimension n i × m j .
Lemma B.2.
Suppose that ˆ M is rank at most p , M (cid:63) is rank p , and ∀ ( i, j ) (cid:54) =
(2 , 2) , (cid:107) ˆ M i,j − M (cid:63)i,j (cid:107) F ≤ (cid:15), and (cid:107) M
(cid:63)i,j (cid:107) F ≤ M, where (cid:15) ≤ σ p ( M (cid:63) 11 ) / 2 .
Then,
[Equation]
Proof.
The proof mirrors that of Shah et al.
(2020, Proposition 13).
We shall show below that ˆ M is of rank exactly p .
Hence, Shah et al.
(2020, Lemma 12) gives the following exact expression for the bottom-right blocks,
[Equation]
where above ( · ) † denotes the Moore-Penrose pseudoinverse.
Since (cid:107) ˆ M 11 − M (cid:63) 11 (cid:107) op ≤ (cid:107) ˆ M 11 −
M (cid:63) 11 (cid:107) F ≤ (cid:15) ≤ σ p ( M (cid:63) 11 ) / 2 , Weyls inequality implies that ˆ M 11 is rank p (as promised), and (cid:107) ˆ M † 11 (cid:107) op ≤ 2 σ p ( M (cid:63) 11 ) − 1 .
Similarly, as (cid:107) ˆ M 12 − M (cid:63) 12 (cid:107) op ≤ σ p ( M (cid:63) 11 ) / 2 ≤ M/ 2 , so (cid:107) ˆ M 12 (cid:107) op ≤ 32 M .
Thus,
[Equation]
Next, using a perturbation bound on the pseudoinverse 5 due to Meng & Zheng (2010, Theorem 2.1),
[Equation]
Therefore, we conclude
[Equation]
B. G ENERAL A NALYSIS FOR C OMBINATIORAL E XTRAPOLATION
We now provide our general analysis for combinatorial extrapolation.
To avoid excessive subscripts, we write X =
W × V rather than X =
X 1 × X 2 as in the main body.
We consider extrapolation under the following deﬁnition of combinatorial support.
Deﬁnition B.2 (Bounded combinatorial density ratio) .
Let D , D (cid:48) be two distributions over a product space W × V .
We say D (cid:48) has κ -bounded combinatorial density ratio with respect to D , written as
D (cid:48) (cid:28) κ, comb D , if there exist distributions D W ,i and D V ,j , i, j ∈ { 1 , 2 } , over W and V , respectively, such that D i ⊗ j := D W ,i ⊗ D V ,j satisfy (cid:80)
[Equation]
i,j )
[Equation]
2 , 2) D i ⊗ j (cid:28) κ D , and D (cid:48) (cid:28) κ (cid:80)
i
[Equation]
j =
[Equation]
2 D i ⊗ j
[Equation]
For simplicity, we consider scalar predictors, as the general result for vector valued estimators can be obtained by stacking the components.
Speciﬁcally, we consider a ground-truth predictor h
[Equation]
and estimator ˆ h of the form
[Equation]
Lastly, we choose the (scalar) square-loss, yielding the following risk
[Equation]
Throughout, we assume that all expectations that arise are ﬁnite.
Our main guarantee is as follows.
Theorem 2.
Let D , D (cid:48) be two distributions on W ×V satisfying D (cid:48)
(cid:28) κ, comb D , with corresponding factor distributions D W ,i and D V ,j , i, j ∈ { 1 , 2 } .
Deﬁne the effective singular value
[Equation]
and suppose that max 1 ≤ i,
[Equation]
≤
[Equation]
B.3. P ROOF OF T HEOREM
First, let us assume the following two conditions hold; we shall derive these conditions from the conditions of Theorem 2 at the end of the proof:
[Equation]
Our strategy is ﬁrst to prove a version of Theorem 2 for when W and V have ﬁnite cardinality by reduction to the analysis of matrix completion in Lemma B.2, and then extend to arbitrary domains via a limiting argument.
Lemma B.3.
Suppose that Eq.
(B.3) hold, and in addition, that W and V have ﬁnite cardinality.
Then,
[Equation]
Proof Lemma B.3.
By adding additional null elements to either W or V , we may assume without loss of generality that |W| = |V| = d , and enumerate their elements { w 1 , . . .
, w d } and { v 1 , . .
, v d } .
Let p i,a =
Pr w ∼D W ,i [ w = w a ] and q j,b
= Pr v ∼D V ,j [ v = v b ] .
Consider matrices ˆ M , M (cid:63)
∈ R 2 d × 2 d , with d × d blocks
[Equation]
a q j,b · h (cid:63) ( w a , v b ) .
We then verify that (cid:107) ˆ M ij − M (cid:63)ij (cid:107) 2F = d (cid:88) a,b = p i,a q j,b (ˆ h ( w a , v b ) − h (cid:63) ( w a , v b ))=
E D i ⊗ j [(ˆ h ( w, v ) − h (cid:63) ( w, v )) 2 ] =
R (ˆ h ; D i ⊗ j ) , (B.4) and thus (cid:107) ˆ M ij − M (cid:63)ij (cid:107) 2F ≤ (cid:15) 2 for ( i, j ) (cid:54) = (2 , 2) .
Furthermore, deﬁne the matrices ˆ A
i , ˆ B j via ( ˆ A i ) a := √ p
i,a ˆ f ( w a ) (cid:62) , ( ˆ B j ) b := √ q j,b ˆ g ( v b ) (cid:62) , and deﬁne A (cid:63) i , B (cid:63) j similarly.
Then, ˆ M = (cid:20)
ˆ A 1 ˆ
A(cid:21) (cid:20) ˆ Bˆ B(cid:21) (cid:62) ,
M (cid:63) =
(cid:20) A (cid:63) 1 A (cid:63) (cid:21) (cid:20) B (cid:63)
B (cid:63)(cid:21) (cid:62) , showing that rank( ˆ M 1 ) , rank( ˆ M 2 ) ≤ p .
Finally, by Eq.
(B.2),
[Equation]
Lastly, we have
[Equation]
Thus, Eq.
(B.4) and Lemma B.2 imply that
[Equation]
Lemma B.4.
Suppose that Eq.
(B.3) hold, but unlike Lemma B.4, W and V need not be ﬁnite spaces.
Then, still, it holds that
[Equation]
Proof of Lemma B.4.
For n ∈ N , deﬁne h (cid:63),n =
(cid:104) f (cid:63),n , g (cid:63),n (cid:105) and ˆ h n = (cid:104)
ˆ f n , ˆ g n (cid:105) , where f (cid:63),n , ˆ f n , ˆ g n , g (cid:63),n are simple functions (i.e. ﬁnite range, see the proof of Lemma B.1) converging to f (cid:63) , ˆ f, g (cid:63) , ˆ g .
Deﬁne
[Equation]
By the dominated convergence theorem 7 ,
[Equation]
In particular, as (cid:15) 2 ≤ σ 2 / 4 , then applying Lemma B.3 for n sufﬁciently large,
[Equation]
Indeed, for any ﬁxed n , all of ˆ f n , ˆ g n , f (cid:63),n , g (cid:63),n are simple functions, so we can partition W and V into sets on which these embeddings are constant, and thus treat W and V as ﬁnite domains; this enables the application of Lemma B.3 applies.
Finally, using the dominated covergence theorem one last time,
[Equation]
We can now conclude the proof of our proposition.
Proof of Theorem 2.
As D (cid:48) (cid:28) κ (cid:80) i,j D i ⊗ j and (cid:80) i,j (cid:54) =
[Equation]
[Equation] [Equation]
D
[Equation]
j
[Equation]
, Lemma B.1 and additivity of the integral implies
[Equation]
Moreover
[Equation]
setting (cid:15) 2 :=
κ R (ˆ h ; D ) [Equation] we have
[Equation]
Thus
[Equation]
for R (ˆ h ; D ) <
σ 2 (cid:63) 4 κ , Eq.
(B.3) holds and thus Lemma B.4 entails
[Equation]
Thus, combining with Eq.
(B.5),
[Equation]
completing the proof.
B. E XTRAPOLATION FOR T RANSDUCTION
Leveraging Theorem 2, this section proves Theorem 1, thereby providing a formal theoretical justiﬁcation for predictors of the form Eq.
Proof of Theorem 1.
We argue by reducing to Theorem 2.
The parameterization of the stochastic predictor h θ in Eq.
(3.1), followed by Assumption 3.2 allows us to write
[Equation]
In the above display, the joint distribution of ( x − x (cid:48) , x (cid:48) ) is precisely given by ¯ D train (see Eq. (3.2)).
Hence,
[Equation]
Further, as (cid:96) ( y, y (cid:48) )
= (cid:107) y − y (cid:48)
(cid:107) 2 is the square loss and decomposes across coordinates,
[Equation]
D
[Equation]
[Equation]
By the same token,
[Equation]
To conclude the proof, we remain to show that for all k ∈ [ K ] , we have
[Equation]
where C prob = κ 2 (cid:18) 1 + 64 M 4 σ 4 (cid:19) .
[Equation]
Indeed, for each k ∈ [ K ] , we have
[Equation]
Hence Eq.
(B.7) holds by invoking Theorem 2 with the correspondences W ← ∆ X , V ← X , σ (cid:63) ← σ , M (cid:63) ← M and κ ← κ .
This concludes the proof.
B.4. F URTHER R EMARKS ON T HEOREM
• The singular value condition, min { σ p ( E D ∆ X , 1 [ f (cid:63),k f (cid:62) (cid:63),k ]) , σ p ( E D X , 1 [ g (cid:63),k g (cid:62) (cid:63),k ]) } ≥ σ 2 > 0 , mirrors the non-degeneracy conditions given in the past work in matrix completion (c.f.
(Shah et al., 2020)).
• The support condition sup ∆ x,x (cid:48) | ¯ h (cid:63),k (∆ x, x (cid:48) )
| ≤ M is a mild boundedness condition, which (in light of Theorem 2) can be weakened further to
[Equation]
where D i ⊗ j are the contituent distributions witnessing ¯ D test (cid:28) κ, comb ¯ D train .
• The ﬁnal condition, R ( h θ ; D train ) ≤ σ 2 4 κ , is mostly for convenience.
Indeed, as M ≥ σ and κ ≥ 1 , then as soon as R ( h θ ; D train ) >
σ 2 4 κ , our upper-bound on R ( h θ ; D test ) is no better than
[Equation]
which is essentially vacuous.
Indeed, if we also inﬂate M and stipulate that sup
[Equation]
M , we can remove the condition R ( h θ ; D train ) ≤ σ
[Equation]
altogether.
[Figure]
C ONNECTIONS TO MATRIX COMPLETION
[Figure]
Building upon Deﬁnition B.1, Assumption 3.1 introduces a notion of bounded density ratio between ¯
D train and ¯ D test in the OOC setting.
Take the discrete case of matrix completion as an example, as illustrated in Fig 9, the training distribution of (∆ x, x (cid:48) ) covers the support of the (1 , 1) , (1 , 2) , (2 , 1) blocks of the matrix, while the testing distribution of (∆ x, x (cid:48) ) might be covered by any product of the marginals of the 2 × 2 blocks.
With this connection in mind, it is possible to establish the OOC guarantees on ¯ D test as in matrix completion tasks, if the bilinear embedding admits some low-rank structure.
In other words, samples from the top-left and off-diagonal blocks uniquely determine the bottom-right block.
This is a standard setting studied extensively in the matrix completion literature, see (Shah et al., 2020; Agarwal et al., 2021b; Athey et al., 2021) for example.
Our results in Theorem 1 can be viewed as a generalization of these results to the continuous space embeddings.
Finally, note that
Assumptions 3.2 and 3.3 are also adopted correspondingly from the matrix completion literature, where Assumption 3.2 concerns the ground-truth bilinear structure of the embeddings, which mirrors the low-rank structure of the underlying matrix in the matrix completion case; Assumption 3.
corresponds to the non-degeneracy assumption on the the top-left block distribution, as referred to as “anchor” columns/rows in the past work when dealing with matrices (Shah et al., 2020).
C A DDITIONAL R ESULTS C. L EARNING WHICH POINTS TO TRANSDUCE : W EIGHTED T RANSDUCTION
There may be scenarios where it is beneﬁcial to more carefully choose which points to transduce rather than comparing every training point to every other training point.
To this end, we outline a more sophisticated proposal, weighted transduction , that achieves this.
Sampling train pairs uniformly may violate the low rank structure needed for extrapolation discussed in Section 3.2.
Consider a dataset of point clouds of various 3D objects in various orientations where the task is to predict their 3D grasping point.
A predictor required to make predictions from similarities between point clouds of any two different objects and orientations may need to be far more complex than simpler predictors trained on only similar objects.
Therefore, it may be useful to identify which train points to transduce, rather than transduce between all pairs of points during training.
Hence, we introduce weighted transduction described in Algorithm 2, which identiﬁes which train points to transduce, rather than transduce between all pairs of points during training.
During train, a transductive predictor ¯ h θ is trained together with a scalar bilinear weighting function ω θ .
¯ h θ is trained to predict values as in the unweighted transduction algorithm, yet weighted by ω θ .
ω θ predicts weights for data points and their potential anchors based on anchor point similarity and the anchor point, thereby deciding which points to transduce.
We can either pre-train ω θ , or jointly optimize ω θ and ¯ h θ .
At test time, we select an anchor point based on the learned weighting function, and make predictions via ¯ h θ .
We show results on learning with weighted transduction on analytic functions in Appendix C.2 and on learning how to grasp from point clouds in Appendix C.3.
=
[Table]
C. A DDITIONAL R ESULTS ON A S IMPLE 2-D E XAMPLE
Weighted transduction with a 2-D example We also considered a 2-D analytic function as shown in Fig 10.
This function has a 2-D input ( x 1 , x 2 ) , and 2-D output ( y 1 , y 2 ) .
The function is constructed by sampling a tiling of random values in the central block (between (1 , 1) and (5 , 5) ).
The remaining tiles are constructed by shifting the central tile by 6 units along the x 1 or x 2 direction.
For tiles shifted in the x 1 -direction ( (6 , 0) or (0 , 6) )
the label has the same ﬁrst dimension and a negated second dimension.
For tiles shifted in the x 2 -direction ( (6 , 0) or (0 , 6) )
the label has the same second dimension and a negated ﬁrst dimension.
This has a particular symmetry in its construction where the problem has a very simple relationship between instances, but not every instance is related to every other instance (only those off by a constant offset).
We show in Fig 10 that weighted transduction is particularly important in this domain as not every point corresponds to every other point.
We perform weighted transduction (seeded by a few demonstration pairs), and this allows for extrapolation to OOS pairs.
While this may seem like an OOC problem, we show through the standard bilinear (without reparameterization) comparison that there is no low rank structure between the dimensions but instead there is low rank structure on reparameterization.
This shows the importance of using weighted transduction and the ability to ﬁnd symmetries in the data with our proposed method.
C. A DDITIONAL R ESULTS ON S UPERVISED G RASP PREDICTION
Qualitative grasping prediction results For qualitative results of all models on the grasping point prediction task, please see Fig 11.
Comparing to an SE(3) equivariant architecture We compare bilinear transduction to a method that builds in SE(3) equivariance domain knowledge into its architecture and operates on point clouds.
Speciﬁcally, we compare to tensor ﬁeld networks Thomas et al.
(2018), and report the results in Table 2. To highlight that SE(3) equivariance architectures can extrapolate to new rotations but lack guarantees on scale, we train on 100 noisy samples of rotated mugs and on a separate set of scaled mugs by in-distribution values described in Appendix D.1. We test extrapolation to 50 mugs with out-of-sample orientations and scales sampled from the distribution described in Appendix D.1. We adapt the Tetris object classiﬁcation tensor ﬁeld network to learn grasp point prediction by training the mean of the last equivariant layer to predict a grasping point and applying the mean squared error (MSE) loss.
Following the implementation in Thomas et al.
(2018), we provide six manually labeled key points of the point cloud that uniquely identify the grasp point: two points on the handle and mug intersection and four points on the mug bottom.
We compare to the linear, neural network and transduction baselines, as well as bilinear transduction.
We ﬁx these architectures to one layer and 32 units.
Under various rotations, the tensor ﬁeld network is able to extrapolate.
For scaling however, the network error increases signiﬁcantly, which is in accordance with the fact that scale transformation guarantees do not exist in SE(3) architectures.
Bilinear transduction can extrapolate in both cases, is comparable to tensor ﬁeld networks on OOS rotations and exceeds its performance on OOS scales.
Weighted grasp point prediction The grasp point prediction results for three objects described in Table 1 are trained with priviledged knowledge of the object types at train time.
I.e., the predictor at train time was provided only with pairs of same object categories but with different positions, orientations and scale.
At test time an anchor was selected based on in-distribution similarities without knowledge of the test sample’s object type.
While this provides extra information, it is still not building in information about the distribution shift we would like to extrapolate to, i.e. the types of transformations the objects undergo: translation, rotation and scale.
We apply the weighted transduction algorithm (Algorithm 2) to this domain and show that with more relaxed assumptions on the required object labels at train time we can learn weights and a predictor that perform as well as with priviledged training (Table 3).
For grasp point prediction, at train, we train a weighting function ω θ ( x j − x i , x i ) to predict whether x i should be transduced to x j based on object-level labels of the data (transduce all mugs to other [Table] mugs and bottles to other bottles).
Then we learn a predictor h θ weighted by ﬁxed ω θ as described in Algorithm 2.
At test time, given point x j we select the train point x i with highest value ω θ ( x j − x i , x i ) and predict y j via h θ ( x j − x i , x i ) , without requiring privileged information about the object category at test time [Table] [Table]
C. A DDITIONAL R ESULTS ON I MITATION L EARNING
Qualitative decision making results For qualitative results of all models on imitation learning domains, please view videos in the supplementary zip ﬁle.
C. C HOICE OF A RCHITECTURE
Periodic activation functions We provide further results on the grasp point prediction and imitation learning domains in Table 4.
This experiment is similar to the one reported in Table 1, but with fourier pre-processing as the initial layer in each model.
For the various domains, adding this pre-processing step is beneﬁcial for some architecture-domain combinations, but does not surpass bilinear transduction on all domains for any method.
Moreover, in most domains bilinear transduction achieves the best results, or close results to the best model.
For implementation details see
Section D.
[Table]
Robustness across seeds and architectures: We show that our results are stable across the choice of architecture and seed.
In Fig 12 we show the performance of our method and baselines over
Meta-World tasks and architectures.
In Table 5 we show the average performance of our method and baselines over the complex domains for ﬁxed architecture parameters over three seeds.
[Table]
Reach
[Table]
Push
[Table]
Slider
[-198 ± 
-][Table]
Adroit [Table] [Figure] push push fourier
Neural Net
DeepSets d e p t h 1 h i d d e n 3d e p t h 1 h i d d e n 5 1d e p t h 1 h i d d e n 1 0 2d
e p t h 2 h i d d e n 3d e p t h 2 h i d
d e n 5 1d e p t h 2 h
i d d e n 1 0 2reach reach fourier push push fourier
Transduction d e p t h 1 h i d d e n 3d e p t h 1 h i d
d e n 5 1d e p t h 1 h i d d e n 1 0 2d
e p t h 2 h i d d e n 3d e p t h 2 h i d
d e n 5 1d e p t h 2 h
i d d
e n 1 0
[Figure]
[Table]
D I MPLEMENTATION D ETAILS D. T RAIN /T EST DISTRIBUTIONS AND A NALYTIC F UNCTION D ETAILS
Analytic functions: We used the following analytic functions for evaluation in Section 4.1.
5(a) Periodic mixture of functions with different periods: We consider the period growing mixture of functions, and do a similar function but with two functions with different periods repeating.
Let us deﬁne x v = x
mod 9 , then the
[Equation]
We can see that this function mixes 2 different periods, to show that our framework can deal with varying functions as well.
5(b)
Sawtooth function: we use a classic sawtooth function
[Equation]
5(c) Randomly chosen polynomial: We sampled a random degree 8 polynomial: f ( x ) =
( x − 0 .
1)( x + 0 . 4)( x + 0 .
7)( x [Equation]
0 . 5)( x
[Equation]
1 . 5)( x
[Equation]
1 . 75)( x + 1)( x
[Equation]
1 . 2) 6 and 8 Periodic growing mixture of functions: Let us deﬁne x v = x mod 3 and x m = (cid:98) x 3 (cid:99)
[Equation]
functions: Let us deﬁne x v = x mod 3 and x m = (cid:98) x 3 (cid:99)
[Equation]
This represents a mixture of different functions - a sinusoid function, a linear function and a quadratic function that repeat and are equivariant on shifting.
These are meant to show that bilinear transduction can capture equivariance.
Grasp point prediction: A train dataset is generated from rotating, translating and scaling a bottle, mug and teapot from the ShapeNet dataset Chang et al.
(2015) with various parameters to predict a grasping point in R 3 from their point clouds.
A test dataset is constructed by applying the same transformations to the objects with OOS parameters.
The point cloud in R N × 3 is represented by a 12-tuple containing the point cloud’s mean point and its three PCA components in R 3 .
In Figwe plot the reduced point cloud state space representing point cloud mean and PCA components for in-distribution and OOS object orientations, positions and scales.
We add Gaussian noise to the point cloud sampled from N (0 , 0 . 005) and to the grasp point label sampled from N (0 , 0 . 002) .
During train, objects are rotated around the z axis θ ∈ [0 , 1 . 2 π ] , translated by ( x, y ) ∈ [0 , 0 . 5] × [0 , 0 . 5] and scaled by a factor of α ∈ [0 . 7 , 1 . 3] .
At test time these transformation parameters are sampled from the following ranges: θ ∈ [1 . 2 π, 2 π ] , ( x, y ) ∈ [0 . 5 , 0 .
7] × [0 . 5 , 0 .
7] and α ∈ [1 . 3 , 1 . 6] .
Meta-World: We evaluate on the reach-v2 and push-v2 tasks.
For reach-v2 , we reduce the observation space to a 6-tuple of the end effector 3D position and target 3D position.
The action space is a 4-tuple of the end effector 3D position and the gripper’s degree of openness.
During train, the 3D target position for reach-v2 is sampled from range [0 , 0 . 4] ×
[0 . 7 , 0 . 9] × [0 . 05 , 0 . 3] and at test time from range [ − 0 .
4 , 0] × [0 . 7 , 0 . 9] × [0 . 05 , 0 . 3] .
The gripper initial position is ﬁxed to [0 , 0 . 6 , 0 . 2] .
For push-v2 the observation space is a 10-tuple of the end effector 3D position, gripper openness, object 3D position and target 3D position.
The action space remains the same as in reach-v2.
The train 3D target position is sampled from [0 , 0 . 3] × [0 . 5 , 0 .
7] × { 0 . 01 } and at test time from [ − 0 .
3 , 0] × [0 . 5 , 0 .
7] × { 0 . 01 } .
The gripper initial position is ﬁxed to [0 , 0 . 4 , 0 . 08] and the object initial position to [0 , 0 . 4 , 0 . 02] .
For both environments, expert and evaluation trajectories are collected for 100 steps.
Adroit hand: We evaluate on the object relocation task in the ADROIT hand manipulation benchmark.
The observation space is a 39-tuple of 30D hand joint positions, 3D object position, 3D palm position and 3D target position.
The action space is a 30-tuple of 3D position and 3D orientation of the wrist, and 24D for ﬁnger positions commanding a PID controller.
During train, the target location is sampled from [0 , 0 . 3] × [0 , 0 . 3] ×
[0 . 15 , 0 . 35] and at test time from [ − 0 .
3 , 0] × [ − 0 .
3 , 0] × [0 . 15 , 0 . 35] .
Expert and evaluation trajectories are collected for 200 steps.
Slider: We introduce an environment for controlling a brush to slide an object to a target.
The observation space is an 18-tuple composed of a 9-tuple representing the brush 2D position, 3D object position and 4D quaternion, an 8-tuple representing the 2D brush velocity, 3D object velocity and 3D rotational object velocity, and the object mass.
The action space is torque applied to the brush.
During train, the object mass is sampled from [60 , 130] and at test time from [5 , 15] .
Expert and evaluation trajectories are collected for 200 steps.
In Figure 14 we plot the actions as a function of time steps for demonstrations, baselines described in Section 4 and Bilinear Transduction for a set of ﬁxed architecture parameters.
D. T RAINING DETAILS D.2. D ETAILS OF B ILINEAR T RANSDUCTION I MPLEMENTATION
We describe the unweighted transduction algorithm implementation in detail.
For grasp point prediction, at train we uniformly select an object o ∈ { bottle , mug , teapot } , and two instances of those objects, o i , o j that were rotated, translated and scaled with different parameters.
We learn to transduce o i to o j , i.e. learn h θ ( o j −
o i , o i ) to predict y j , the grasping point of o j .
At test time, given test point x j with an unknown object category, we select a train point x
i with an “in-distribution” difference x j − x i , as described in Algorithm 1.
We select ρ to be the closest difference x
j − x i to an in-distribution difference over all train points x i .
The differences generated by the optimization process are sampled uniformly to generate a smaller subset for comparison.
For the goal-conditioned imitation learning setting, at train we uniformly sample trajectories τ i = { ( s it , a it ) }
Tt =1 for horizon T speciﬁed in Appendix D.1 and τ j = { ( s jt , a jt ) }
Tt =1 .
We further uniformly sample a time step t from [0 , T ] .
We learn to transduce state s it ∈
τ i to s jt ∈ τ j .
I.e., learn h θ ( s jt − s it , s it ) to predict action a jt .
At test time, we select an anchor trajectory based on the goal (Meta-World, Adroit) or mass (Slider) g j .
This is done by selecting the train trajectory τ i with goal or mass g i that generates an “in-distribution” difference g j − g
i with the test point, as described in
Algorithm 1.
We select ρ to be the 10 -th percentile of the in-distribution differences.
The differences generated by the optimization process are approximated by generating differences from the train data.
We then transduce each state in τ i to predict actions for the test goal or mass.
Given test state s jt , we predict a jt with h θ ( s jt − s it , s it ) , execute action a jt in the environment and observe state s jt +1 .
We complete this process for T steps to obtain τ j .
D.2. D ETAILS OF T RAINING D ISTRIBUTIONS
We train on 1000 samples for all versions of the grasp point prediction domain: single objects, three objects and weighted transduction.
In the sequential decision domains we train on N demonstrations, sequences of state-action pairs, { ( x i , y i ) } Ti =1 where horizon T for each domain is speciﬁed in
D.1.
For Meta-World reach and push, N = 1000 .
For Slider and Adroit N = 100 .
All domains were evaluated on 50 in-distribution out-of-sample points and 50 OOS points.
For each domain, all methods were trained on the same expert dataset and evaluated on the same in-distribution out-of-sample and OOS sets.
We collect expert demonstrations as follows.
For grasp point prediction, we use the Neural
Descriptor Fields Simeonov et al.
(2022) simulator to generate a feasible grasping point by a
Franka Panda arm for three ShapeNet objects (bottle 2d1aa4e124c0dd5bf937e5c6aa3f9597, mug 61c10dccfa8e508e2d66cbf6a91063 and Utah teapot wss.ce04f39420c7c3e82fb82d326efadfe3).
In
Meta-World, we use the expert policies provided in the environment.
For Adroit and Slider we train an expert policy using standard Reinforcement Learning (RL) methods - TRPO Schulman et al.
(2015) for Adroit and SAC Haarnoja et al.
The weighting grasp point prediction network was provided with with 300 labels of positive pairs to be transduced and 300 negative pairs labeled with binary labels.
This is a slightly more relaxed condition than labeling 1000 objects as in the unweighted version.
D.2. D ETAILS OF M ODEL A RCHITECTURES
For analytic domains, we use MLPs (both for NN and bilinear embeddings) with 3 layers ofhidden units each with ReLu activations.
We train all models with periodic activations of fourier features as described in Tancik et al.
Linear model : data points x are processed through a fully connected layer to output predictions y .
Neural Network : data points x are processed through an mlp with k layers, n units each and ReLU activations for hidden layers to output predictions y .
DeepSets : we refer to the elements of the state space excluding target position (Meta-World, Adroit) or mass (Slider) as observations.
Observations and target position or mass are processed separately by two mlps with k layers, n units each and
ReLU activations for hidden layers.
Their embeddings are summed and processed by an mlp with one hidden layer with n units and ReLU activations to produce predictions y .
Transduction : train point x j and x i − x j (for train or test point x i ) are concatenated and processed through an mlp with k layers, n units each and ReLU activations for hidden layers, to predict y i .
Our architecture
Bilinear Transduction : train point x j and ∆ x ij = x
i − x j (for train or test point x i ) are embedded separately to a feature space in R d · m , where d is the dimension of y i , by two mlps with k layers, n units each and ReLU activations for hidden layers.
Each predicted element of y i is the dot product of m -sized segments of these embeddings as described in Eq.
Further, we evaluate for all models with Fourier Embedding as a pre-processing step.
We process inputs x , target or mass or ∆ x through a linear model which outputs × 40 the number of inputs.
We then multiply values by π and process through the sin function.
This layer is optimized as the ﬁrst layer of the model.
In Table 1 we search over number of layers k ∈ { 2 , 3 } , and unit size n ∈ { 32 , 512 , 1024 } .
The bilinear transduction embedding size is m = 32 .
In Tables 3 and 5 we set k = 2 , n = 32 and m = 32 .
The weighting function used for grasp point prediction is the bilinear transduction architecture with k = 2 , n = 128 and m = 32 .
D.2. D ETAILS OF M ODEL O PTIMIZATION
We train the analytic functions for 500 epochs, batch size 32 , and Adam optimizer with learning rate 1 e − 4 .
We optimize the mean squared error (MSE) loss for regression.
We trained all more complex models for 5 k epochs, batch size 32 , with Adam (Kingma & Ba, 2014) optimizer and learning rate 1e-4.
We optimize the L2-norm loss function comparing between ground truth and predicted grasping points or actions for the sequential decision making domains.
Each conﬁguration of hyperparameters was ran and tested on one seed.
We demonstrate the stability of our method across three seeds for a ﬁxed set of hyperparameters in Table 5.
We train the weighted grasp point prediction for 5 k epochs, batch size 16 , and Adam optimizer with learning rate 1 e − 4 .
We optimize the MSE loss between the output and ground truth binary label indicating if a train point should be transduced to another train point.
The weighting function did not require further ﬁnetuning jointly with the bilinear predictor.