Learning to Extrapolate: A Transductive Approach
Authors: Aviv Netanyahu
Abhishek Gupta
Max Simchowitz
Kaiqing Zhang
Pulkit Agrawal
Venue: ICLR 2023 poster (2023)

Abstract: Machine learning systems, especially with overparameterized deep neural networks, can generalize to novel test instances drawn from the same distribution as the training data. However, they fare poorly when evaluated on out-of-support test points. In this work, we tackle the problem of developing machine learning systems that retain the power of overparameterized function approximators while enabling extrapolation to out-of-support test 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 under certain conditions. We instantiate a simple, practical algorithm applicable to various supervised learning and imitation learning tasks.

Published as a conference paper at ICLR 2023
LEARNING TO EXTRAPOLATE: A TRANSDUCTIVE AP- PROACH
Aviv Netanyahu 1.2, Abhishek Gupta 12,3+ Max Simchowitz', Kaiqing Zhang & Pulkit Agrawal 1.2 Improbable AI Lab' MIT
12.3, University of Washington"
University of Maryland, College Park
ABSTRACT
Machine learning systems, especially with overparameterized deep neural net- works, can generalize to novel test instances drawn from the same distribution as the training data. However, they fare poorly when evaluated on out-of-support test points. In this work, we tackle the problem of developing machine learn- ing systems that retain the power of overparameterized function approximators while enabling extrapolation to out-of-suppor test points when possible. This is accomplished by noting that under certain conditions, "transductive" repa- rameterization can convert an out-of-support extrapolation problem into a prob- lem of within-support combinatorial generalization. We propose a simple strategy based on bilinear embeddings enable this type of combinatorial generalization, thereby addressing the out-of-support extrapolation problem under certain condi- tions. We instantiate simple, practical algorithm applicable to various supervised learning and imitation learning tasks.
INTRODUCTION
Generalization is 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-xample general- ization. However, in many scenarios, test data is sampled from a different distribution from the training set, i.e., out-of-ddistribution (OOD). In some OOD scenarios, the test distribution is assumed be known during training a common assumption made by meta-learning methods (Finn al., 2017b). Several works have tackled a more general scenario of "reweighted" distribution shift (Koh al., 2021: Quinonero- Caarro--aanelaaaa et al., 2008) where the test distribution shares support with the train- ing distribution, but has different and unknown probability density. This setting can be tackled via distributional robustness approaches (Sinha et al., 2018; Rahimian & Mehrotra, 2019).
Our paper aims to find structural conditions under which generalization to test data with support outside of the training distribution is possible. Formally, assume the problem of learning function h: h(x) using data { vi)l.N Dirain. where 2; E Xtrains the training domain. We are interested in making accurate predictions h(x)) for x & X (see examples in Fig 1). Consider an example task of predicting actions to reach desired goal (Fig 1b). During training, goals are pro- vided from the blue cuboid (x E Xerain ), but test time goals are from the orange cuboid (x & Xerain). If h is modeled using deep neural network, its predictions on test goals in the blue area likely to be accurate, but for the goals in the orange area the performance can be arbitrarily poor unless further domain knowledge incorporated. This challenge manifests itself variety of real-world problems, ranging from supervised learning problems like object classification (Barbu et al., 2019), sequential decision making with reinforcement learning (Kirk et al., 2021), transferring reinforce- ment learning policies from simulation to real-world (Zhao et al., 2020), imitation learning (de Haan al., 2019), etc. Reliably deploying learning algorithms in unconstrained environments requires ac- counting for such "out-of-support" distribution shift, which we refer to as extrapolation.
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. Sev- cral methods can extrapolate the nature of distribution shift is known apriori: convolution neural
'denotes equal contribution. Correspondence to Aviv Netanyahu <avivn@mit.cdu. Abhishek Gupta wabhgupta@es..wwshigee

---

Published as a conference paper at ICLR 2023
VMAA
(a) Grasp point prediction
(b) Action prediction
Function value prediction
Figure 1: In the real-world the distribution (orange) often has different support than the training distri- bution (blue). Some illustrative tasks: (a) grasp point prediction for object instances with out-of-support scale, (b) action prediction for reaching out-of-support goals, (c) function value prediction for an out-of-support input range. The black crosses show predictions for conventionally trained deep neural network that makes accu- rate predictions for in-support inputs, but fails out-of-support inputs. We propose algorithm that makes accurate out-of-support predictions under set assumptions.
networks are appropriate if training pattern appears at an nut-of-distributio translation in test example (Kayhan & Gemert, 2020). Similarly, accurate predictions can made for object point clouds in out-of-support orientations by building in SE(3) equivariance (Deng et al., 2021b). An- other way to extrapolate is if the model class known apriori: fitting linear function to linear problem will extrapolate. Similarly, methods like NeRF (Zhang et al., 2022a) use physics of image formation to learn a 3D model of scene which can synthesize images from novel viewpoints.
We propose alternative structural condition under which extrapolation is feasible. Typical ma- chine learning approaches are inductive: decision-maie makiing rules are inferred from training data and employed for test-time predictions. An alternative induction is transduction (Gammerman al., 1998) where a test example is compared with the training examples to make a prediction. Our main insight is that in the transductive view of machine learning, extrapolation can be reparameterized as a combinatorial generalization problem, which, under certain low-rank and coverage conditions (Shah et al., 2020; Agarwal et al., 202 Andreas et al., 2016; Andreas, 2019), admits solution.
First we show how we can (i) reparameterize out-of-support inputs hu(ttest) - h Aar, x'), where x' &raain and Az representation of the difference between Ftest and z'. We then (ii) provide conditions under which h(Ax, x²) makes accurate predictions for unseen combinations of (Az, x') based on a theoretically justified bilinear modeling approach: h(Ax,x') x²) -> f(Ar) gxx), where and g map their inputs into vector spaces of the same dimension. Finally, (iii) empirical results demonstrate the 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.
SETUP
Notation. Given a space of inputs X' and targets V, we aim to learn predictor he X + P())) parameterized by 0, which best fits ground truth function h4 X' y. Given some non-negative loss function & y 3 x y T2zo on the outputs (e.g.. squared loss), and a distribution over X, the risk is defined
R(ho: D) == Ee-DEyy-hotr [(y, (x)).
(2.1)
Various training (Dtrain) and test (Drest) distributions yield different generalization settings:
In-Distribution Generalization. This setting assumes Peest Dirain- The challenge to ensure that with N samples from Derain. the expected risk R(hæ; Drest) = R(ho; D is small. This a common paradigm in both empirical supervised learning (e.g., Simonyan Zisserman (2014)) and in standard statistical learning theory (e.g., Vapnik (2006)).
nut-of-Distribution (OOD). This is more challenging and requires accurate predictions when Dtrain # D When the ratio between the density function of Drest to that of D is bounded,
Thhoughout, let P())) denote the of distributions supported y
2

---

Published as a conference paper at ICLR 2023
rigorous OOD extrapolation guarantees exist and detailed in Appendix B.I. Such situation arises when Deest shares support with Duain but is differently distributed as depicted in Fig 2a.
Out-of-Support (00S). There innumerable forms of distribution shift in which density ratios not bounded. The most extreme case is when the support of Deest is not contained in that of D Drrain- Le., when there exists some X' A such that P P..eD [2: X') > 0, but IP P--Dai [a: X¹] (see Fig 2b). We term the problem of achieving low risk on such a Drest as OOS extrapolation.
cat-of-Commination (OOC). This is special case of OOS. Let X' Xi A2 be the product of TWO spaces. Let X1 Derain. denote the marginal distributions of æ1 E X,, 32 E X2 under Draain, and D1 Drest,X;, under Dtest. In OOC learning, Drest,,. Dustt, in the support of X,, X2. but the joint distributions Drest need not be in the support of Drrain
Transduction. In classical transduction (Gammerman et al., 1998), given data points {ri) and their labels { the objective is making predictions for points {x, {x;) In this paper we refer to predictors of that are functions of the labeled data points {a; 1 transductive predictors.
D rest
D Dtrain
Drest
Perain support
D trrain
Pist support
X1
Derain support
Desst support
same support
Pirain support
Drest support
X2
(a) OOD
(b) General 0OS
(c) Structured oOs: OOC
Figure 2: Illustration of different learning settings. (a) in-support out-of distrribution (OOD) learning; (b) general out-of-support (OOS) learning in 1-D (on the left) and 2-D (on the right); (c) out-of- combination (OOC) learning in 2-D.
CONDITIONS AND ALGORITHM FOR OUT-OF-SUPPORT GENERALIZATION
While general OOS extrapolation can arbitrarily challenging, we show that there exist some con- crete conditions under which OOC extrapolation is feasible. Under these conditions, OOS predic- tion problem can be converted into an OOC prediction problem.
Xtrain
23
AXtest A23
(Agg.
X3
X2
X1
Xtest
AX
(A22. (dx xi)
in-support
out-of-support
ofoomb--coinnnni
Figure 3: Illustration of converting OOS to OOC. (Left) Consider training points 21, 22, W' E Arran and 0OS test point &test- During training, we predict he by transducing Ia to ho (Azz, 2s), where Aus = 22 Similarly, at test time, predict he (#re)) by transducing training point æı, via ho (Arret, 21), where Arter: Frest Zi. In this example note that A23 Artest (Right) This conversion yields 00C generalization problem in space A. Xeraini marginal distributions A.Y and Yeaino are covered by training distribution, but their combination not.

---

Published as a conference paper at ICLR 2023
3.1 TRANSDUCTIVE PREDICTORS: CONVERTING OOS TO OOC
We require that input space X' has group structure, i.e., addition and subtraction operators x + x', x' are well-defined for x; x' E X. Let AX == {ar x' x., x' E _]. We propose transductive reparameterization he : X + P()) with deterministic function he A x X - y as
hz(x) == ho(x - 2²),
(3.1)
where s' is referred to as an anchor point for a query point x.
Under this reparameterization, the training distribution can be rewritten as a joint distribution of Ax x - x' and x' as follows:
P (Ac,r') E 1 Pr((Axz,')e E | Derain L'~ Drrsin: Ar -
(3.2)
This is just representing the prediction for querying every point from the training distribution in terms of its relationship to other anchor points in the training distribution.
At test time, we are presented with query point Piest that may be from an OOS distribution. To make a prediction on this OOS 2, we observe that with a careful selection of an anchor point x' from Derain our reparameterization may be able convert this OOS problem into more manageable OOC one, since representing the test point in terms of its difference from training points can still be an "in-support" problem. For a radius parameter p 0, define the distribution of chosen anchor points (x:) (referred to as a transducing distribution) as
P PD)mm Ia' E Pr[æ'
x'
Derain:
inf ArEAN
[[ ((x-xx) --)=''AArr S pl.
(3.3)
where Agrain denotes the set of in our training set, and we denote A.train == {ar 32 : 21; 22 Mtaii J. Intuitively, our choice of Drens (z) selects anchor points x' to transduce from the training distribution, subject the resulting differences (x - x') being close to a "seen" Ar Araaain In doing so, both the anchor point x' and the difference Ar have been seen individually at training time, albeit not combination. This allows us to express the prediction for OOS query point in terms of an in-support anchor point x' and an in supporr difference Ar (but not jointly in supporr..). This choice of anchor points induces joint test distribution of Ax
x' and x':
P Docn ((Aa, x') = Pr((Az, x') E x Drest; x' Derns(r), Ar x xx-xx x').
x'}.
(3.4)
As seen from Fig 3, the marginals of Ar and x' under Drest, are individually in the support those under Diraiu- Still, Fig 3 reveals, since Tiea is out-of-support, the joint distribution of Drest is not covered by that of Dirain (i.e., the combination of x1 and Ftest have not been seen together before); precisely the OOC regime. Moreover, if one tried transduce all x' Durain to Dresst at test time (e.g., transduce point 2i to ftest in the figure) then we would lose coverage of the Au-marginal. By applying transduction keep both the marginal x' and Ar in-support, can convert difficult OOS problems into (potentially) more manageable OOC ones.
ensuring that
3.2 BILINEAR REPRESENTATION 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 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, Ax do not explicitly form finite dimensional matrix), using a bilinear representation of the transductive (oo(s)).o(aa)))
predictor in Eq. (3.1), (Arr x²) = (fo(Az), go(aa)) into a vector space of the same dimension, IR" If the output space is K -idimensional, then we independently model the prediction for each dimension using a set of K bilinear embeddings:
Here fo: ge map their respective inputs
he (Ax, x') (he (Az, 2'), he.k (Ax, 2')); (Az, x') (fe. ga. (x')). (3.5)
While he are bilinear in embeddings fo. Io,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 fo. (Ar).
In implementation, we take (0y;0g), with separate parameters for each embedding.
4

---

Published as a conference paper at ICLR 2023
We now introduce three conditions under which extrapolation is possible. As a preliminary, write Hi <x M2 for two (unnecessarily normalized) measures μ1 H2 to denote μn (1( I k4p2(A) ) for all events A. This notion of coverage has been studied in the reinforcement learning community (termed as "concentrability"" (Munos & Szepesvári, 2008)), and implies that the support of μr1 is contained in the support of j12- The first condition is the following notion of combinatorial coverage.
Assumption& 3.1 (Bounded combinatorial density ratio). We assume that Drest has k-bounded com- binatorial density ratio with respect to Draain. written as Prest <x.comb Durain- This means that there exist distributions Dax and Dx.j 2,j E {1,2}, over Ad and X, respectively, such that Dioj= = DAX, & Dx, satisfy
Σ
Dios <x Derains and D Pesst <x
Σ Diss. 1,11.2
(1,)7
Assumption 3.1 can be interpreted as a generalization of a certain coverage condition matrix factorization. The following distribution is elaborated upon in Appendix B.5 and Fig 9b therein. Intuitively, Dig1 corresponds to the top-left block of a matrix, where all rows and columns are in- support. Di02 and 12æ; correspond the off-diagonals, and D282 corresponds the bottom-right block. The condition requires that Draii covers the entire top-right block and off-diagonals, but need not cover the bottom-right block D2s2. Drest, on the other hand, can contain samples from this last block. 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. (3.5). Our next assump- tions are that the ground-truth embedding is low-rank (we refer to as "bilinearly transducible"), and that under D1@i, 1, its component factors do not lie in strict subspace of R", i.e., not degenerate.
Assumption 3.2 (Bilinearly transducible). For each component E [K]. there exists fa, AX RRP and J+.k X - R2 such that for all x E X. the following holds with probability 1 over (J+, (xx")())) where Ar = x - x'. Further, the ground-truth predictions are bounded: max; suPar,æ' |h,, (Ar, x')l S M.
x' ~ "Drrus (1)) (xx) =
Assumption 3.3. For all components & [K], the following lower bound holds for some o² > 0: min(op(EDi.r [.,)) op (EDios (...DT)))))) 2 0².
(3.6)
The following is proven in Appendix B; due to limited space, we defer further discussion of the bound Appendix B.4.1.
Theorem 1. Suppose Assumptions 3.1 to 3.3 all hold, and that the loss E(.,-) )) is the square loss.
Then, the training risk satisfies R(ho; Derain) S TR. then the test risk bounded
o²
R(he; Drest) S R(ho; Dirain) - x² ( (1+ 64 M²
R(he; Derain) poly(K, M
σt
3.3 OUR PROPOSED ALGORITHM: BILINEAR TRANSDUCTION
Our basic proposal for OOS extrapolation is bilinear transduction, depicted in Algorithm 1: at training time, a predictor he is trained make predictions for training points x; drawn from the training set Xtaain based on their similarity with other points I, E Xtrain ho(#; Izz Kz). The training pairs Z,, are sampled uniformly from Xraain At time, for an OOS point Test we first select an anchor point z; from the training set which has similarity with the test point Twst 2, that within some radius p of the training similarities AXtrain- We then predict the value for Fesst based on the anchor point and the similarity of the test and anchor points: ha (aiee 2;).
For supervised learning, i.e., the regression setting, we compute differences directly between inputs Ja. X. For the goal-conditioned imitation learning setting, we compute difference between states Ij, xy E X sampled uniformly over demonstration trajectories. At test time, we select an anchor trajectory based 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.I. there may be scenarios where it is beneficial more carefully choose which points to transduce. To this end, more sophisticated proposal, weighted transduction, that achieves this. Pseudo code for bilinear transduction is depicted in Algorithm 1, and for the weighted variant in Algorithm 2 in the Appendix.
outline
5

---

Published as a conference paper at ICLR 2023
Note that assume access to representation of X' for which the subtraction operator captures differences that occur between training points and between training and test points. In Appendix D describe examples of such representations defined in standard simulators 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.
Algorithm Bilinear Transduction
Input: Distance parameter p. training set (æ1, y1)...., (I Un)
2: Train: Train 0 on loss L(0) = Σi= Ejzs ((ha(I; i Iy, Tz):
3: Test: For each new Itees, I [[Seett) = (t inf Aerrin || 24 Aal S p). and predict
ho (xreat 2;);;), where iv Uniform(I((raa))
4 EXPERIMENTS
We answer the following questions through an empirical evaluation: (1) Does reformulating oos extrapolation combinatorial generalization problem allow for extrapolation in variety of pervised and sequential decision- maaiinn 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 first analyze oOs extrapolation analytical domains and then on more real world problem domains depicted in Fig 8.
4.1 ANALYZING OOS EXTRAPOLATION ON ANALYTICAL 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 [-1, and tested in 1.6, - [1, 1.6).. We use these domains to gain intuition about the following questions:
What types problems satisfy the assumptions for extrapolation? While we outlined set of conditions under which extrapolation is guaranteed, it is not apparent which problems satisfy these assumptions. understand this better, we considered learning functions with different structure: periodic function with mixed periods (Fig 4a), sawtooth function (Fig 4b) and polynomial function (Fig 4c). Standard deep networks (yellow) the training points well (bluc), but fail to extrapolate to OOS inputs (orange). In comparison, our approach (pink) accurately extrapolates periodic functions but much less effective on polynomials. This is because the periodic functions have symmetries which induce low rank structure under the proposed reparameterization.
Ground Truth (Train)
Ground Truth (Test)
Neural
Bilinear Transduction (Ours)
(a) Mixed periodic function
Sawtooth function
(c) Degree-8 polynomial
Figure 4: Bilinear transduction behavior on I-D regression problems. Bilinear transduction performs well on functions with repeated structure, whereas they struggle arbitrary polynomials. Standard neural nets fail to extrapolate in most settings, even when provided periodic activations (Tancik al., 2020).
Going beyond known inductive biases. Given our method extrapolates to periodic functions in Fig 4 (which displays shift invariance), one might argue that similar extrapolation can be achieved

---

Published as a conference paper at ICLR 2023
by building in an inductive bias for periodicity translation invariance. In Fig show that bilinear transduction fact able to extrapolate even cases that the ground truth function not simply translation invariant, but is translation equivariant, showing that bilinear transduction equivariance. Moreover, bilinear transduction can fact go beyond equivariance or invariance. To demonstrate broader generality of method, consider piecewise periodic function that also grows in magnitude (Fig 7). This function is neither invariant nor equivariant to translation, the group symmetries in this problem do 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 work when comparing training instances, there is simple (low-rank) relationship between how their labels are transformed. While can capture invariance and equivariance, as Fig shows, is more general.
capture
How does the relationship between the training distribution and test distribution affect ex- trapolation 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 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 A. being supput
longer valid beyond this point.
Ground Truth (Train)
Ground Truth (Test)
Neural Net
Equivariant
Bilinear Transduction (Ours)
MM Mm
Figure 5: Performance of trans- ductive predictors as test points go and oOS. Predic- tions are accurate for one "data- width" outside training data.
Figure 6: Prediction on function that displays affine equivariance. Bilinear trandsuction is able capture equivariance without being explicitly encoded.
Figure 7: Prediction func- tion that is neither invariant equivariant. Bilinear transduc- tion able to extrapolate while an equivariant predictor fails.
4.2 ANALYZING OOS EXTRAPOLATION ON LARGER SCALE DECISION- PROBLEMS
To establish that our method is useful for complex and real-world problem domains, we also vary the complexity (i.e., working with nigh-dimensional observation and action spaces) and the learning setting (regression and sequential decision making).
Figure 8: Evaluation domains training (bluc) and oos (orange). (Left to Right:) grasp prediction for var- ious object orientations scales, table-top robotic manipulation for reaching and pushing objects to various targets, dexterous manipulation for relocating objects to various targets, slider control for striking a ball of various mass.
Baselines. To assess the effectiveness of our proposed scheme for extrapolation via trans- duction and bilinear embeddings, we compare with the following non-transductive baselines. Linear Model: linear function approximator to understand whether the issue one of over- parameterization, and whether linear models would solve the problem. Neural Networks: typical training of overparameterized neural network function approximator via standard em- pirical risk minimization. Alternative Techniques with Neural Nenworks (DeepSets): ternative architecture for combining multiple imputs (DeepSets (Zaheer et al., 2017)), that are meant to be permutation invariant and encourage degree of generaliza-
al-

---

Published as a conference paper at ICLR 2023
tion between different pairings of states and goals. Finally, we compare with Transductive Method without a Mechanism for Structured Extraoolation (Transduction): transduc- tion with no special structure for understanding the impact of bilinear embeddings and the low- rank structure. This baseline uses reparameterivation, but just parameterizes he as a standard neural network. We present additional comparison results introducing periodic activations (Tancik al., 2020) in Table 4 in the Appendix.
oos extrapolation in sequential decision making. Table contains results for different types of extrapolation settings. More training and evaluation details 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 set of expert demonstrations
baaching/pushing to goals in the blue box ((0, 0.4] x [0.7, 0.9) x [0.05, 0.3] and [0, 0.3) x (0.5, 0.7] x {0.01)) )) we tested generalization to OOS goals in the orange box (I-0.4, 0] x [0.7, 0.9) x 10.05, 0.3] and 0.3., 0J [0.5, {0.01)), using a simple extension of our method described in Appendix D.2 to perform transduction over trajecto- ries rather than individual states. We quantify performance by measuring the distance between the conditioned and reached goal. Results in Table show that on the easy task of reaching, training linear 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 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 RR3 with an anthropomorphic "Adroit" hand with a much larger action (IR30) and state (I239) space (column 3 in Fig 8). Results confirm that bilinear transduc- tion 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) x [-0.3, 0J x 10.15, 0.35)) after trained on target loca- tions in ((O, 0.3] x [0, 0.3] x [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 significantly 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. Specifically, 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 fixed 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 find 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 significantly less prone to variance than standard inductive or permutation cermutation-ivvarian invariant architectures. This can be seen from a heatmap over various training architec- tures 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 from point clouds of various objects (bottles, mugs, and teapots) with different orientations, positions, and scales (col- umn 1 in Fig 8). At training and test, objects undergo z-axis orientation (00, 1.2x) and [1.2T, 2m) ), translation ((0, 0.5) x [0, 0.5) and [0.5, 0.7) x [0.5, 0.7]) and scaling (00.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 ob- tained via PCA. We consider situations, where we train on individual objects where the training set consists of various rotations, translations, and scales, and standard bilinear transduction is applied. We also consider situations where the objects are not individually identified 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 more in-depth discussion of this assumption and a discussion of how this can be learned via trighted-transductton we refer readers to Appendix C. While training performance is comparable in all instances, we find that extrapolation behavior is significantly better for our method. This is true both for single object cases as well scenarios where all objects are mixed together (Table 1). These experiments show that bi- linear 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 Appendices C and D.
8

---

Published as a conference paper at ICLR 2023
Table 1: Mean and standard deviation over prediction (regression) or final state (sequential decision making) error for OOS samples and over a hyperparameter search.
Task
Expert
Linear
Neural Net
DeepSets
Transduction
Ours
Mug
0.068 1 0.013
0.075 + 0.046
0.055 + 0.043
0.013 + 0.007
Bottle
0.026
0.005
0.05 + 0.051
0.027 + 0.016
0.008 1 0.004
Teapot
0.095
0.02
0.101
0.078
0.043 + 0.02
0.022 + 0.014
0.143
0.116
0.118
0.075
0.112 0.08
0.018 1 0.012
Reach
0.006
0.008
0.007
0.006
0.036
1 0.054
0.19 1 0.200
0.036 1 0.048
0.007 1 0.006
Push
0.012
0.001
0.258
0.063
0.258
0.167
0.199
0.114
0.159 = 0.116
0.02 0.017
Slider
0.105
0.066
0.609 0.07
0.469
0.336
0.274
0.262
0.495
0.330
0.149 1 0.113
Adroit
0.035 + 0.015
0.337 0.075 0.331 = 0.203
0.521 = 0.457
0.409 = 0.32
0.147 + 0.117
5 ABRIDGED RELATED WORK
Generalization from training data to test data of the same distribution has been studied extensively, both practically (Simonyan & Zisserman, 2014) and theoretically (Vapnik, 2006; Bousquet & Elisse- eff, 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 structured neural networks, which hardcode such symmetries as equivariance (Cohen & Welling, 2016; Simeonov et al., 2022), Fuclidean 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 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 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 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-ree at-random setting (Candès & Recht, 2009; Recht, 2011; Can- des & Plan, 2010) equivalent to classical in-distribution statistical learning, we adopt more recent perspectives on missing- not-at-randon data, see e.g.. (Shah et al., 2020; Agarwal et al., 2021b; Athey et al., 2021), which tackle the combinatorial-generrliiiie
setting described in Section 3.2.
In the classical transduction setting (Joachims, 2003; Gammerman et al., 1998; Cortes & Mohri, 2006), the goal is make predictions on 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 test time. Still, we find it useful to adopt "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.
DISCUSSION
Our work serves as initial study of the cireumstances under which problem structure can be both discovered and exploited for extrapolation, combining parametric and non-parametric approaches. The main limitations of this work are assumptions regarding access to a representation of X' and similarity measure for obtaining Ar. Furthermore, we are only guaranteed to extrapolate to regions of X' that admit Az within the training distribution. A number of natural questions arise for further research. First, can we classify which set of real-world domains fits our assumptions, beyond the domains we have demonstrated? Second, can we learn a representation of X' in which differences Ar are meaningful for high dimensional domains? And lastly, are there more effective schemes for selecting anchor points? For instance, mmlogy-making for anchor point selection may reduce the complexity of he. guarantecing low-rank structure.

---

Published as a conference paper at ICLR 2023
7 ETHICS STATEMENT
Bias In this work on extrapolation via bilinear transduction, we can only aim to extrapolate to data that shifts from the training of this method relies on the diversity within the training data, allowing some forms of extrapolation but not ones that do not exist within the training distribution.
ways that exist within the training data. Therefore the performance
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.
REPRODUCIBILITY STATEMENT
We describe algorithms in Section 3.3 and the complete proof of our theoretical results and as- sumptions in Appendix B. Extensive implementation details regarding our algorithms, data, models, and optimization are provided in Appendix D. In addition, plan in the future to release code and data.
ACKNOWLEDGMENTS
We thank Anurag Ajay, Tao Chen, Zhang-War Wei Hong, Jacoh Huh, Leslie Kaelbling, Hannah Lawrence, Richard Li, Gabe Margolis, Devavrat Shah and Anthony Simeonov for the helpful dis- cussions and feedback on the paper. We are grateful to MIT Supercloud and the Lincoln Laboratory Supercomputing Center for providing HPC resources. This research was also partly sponsored by the DARPA Machine Common Sense Program, MIT-IBM Watson AI Lab, the United States Air Force Research Laboratory and the United States Air Force Artificial Intelligence Accelerator and was ac- complished under Cooperative Agreement Number FA8750-19 19:00-1--110 2- 1000. The views and conclusions contained in this document are those of the authors and should not interpreted as representing the official policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes, notwithstanding any copyright notation herein.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Fares Abu-Dakka, Matteo Saveriano, and Luka Peternel. Periodic dmp formulation for quaternion trajectories. In 2021 20th International Conference on Advanced Robotics (ICAR), pp. 658-663.
IEEE, 2021.
Anish Agarwal, Abdullah Alomar, Varkey Alumootil, Devavrat Shah, Dennis Shen, Zhi Xu, and Cindy Yang. Persim: Data-efficient offline reinforeement learning with heterogeneous agents via personalized simulators. Advances in Neural Information Processing Systems, 34: 18644-8876, 2021a.
Anish Agarwal, Munther Dahleh, Devavrat Shah, and Dennis Shen. Causal matrix completion. arXiv preprinl arXiv:2 109. 10009//5555 15154, 2021
Jacob Andreas. Measuring compositionality in representation learning. In 7th International Confer- ence on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRe- view.net, 2019. URL https: /wwpnreveeeeeeeeeeooe/.eee /
net fwwumoidddddioeee htt.rrmmmmddoto0o000o00oooo
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 39-48. IEEE Computer Society, 2016. doi: 10.1 N0.100/CVPR.2016.10 URL https: / /ooii/ii..///// org/ 10. 1109 CVPR. 201 6 12.
Marcin Andrychowicz. Dwight Crow, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin. Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In Is- abelle Guyon, Ulrike von Luxburg. Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vish- wanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5048-5058, 2017. URL https: http://pocediiiccccocoooc ngs nwurips.rr.cmmmm Abstract .tmaee.thtmm......
paper/2017. htt.//////////h///h 453 /
Susan Athey, Mohsen Bayati, Nikolay Doudchenko, Guido Imbens, and Khashayar Khosravi. Ma- trix completion methods for causal panel data models. Journal of the American Statistical Asso- ciation, 176653)) 1730, 2021.
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. Advances in neural information processing systems, 32, 2019.
Peter I. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463 482, 2002.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchezz Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 2018.
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew Wilson. Learning invariances in neural networks from training data. Advances in neural information processing systems, 13:17605- 17616, 2020.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learn ing Research, 2:499-51 526, 2002.
Emmanuel J Candes and Yaniv Plan. Matrix completion with noise. Proceedings of the IEEE, 98 (6):925-936) 2010.
Emmanuel Candès and Benjamin Recht. Exact matrix completion via convex optimization. Foun- dations of Computational mathematics, 96))717-772, 2009.
Quentin Cappart, Didier Chételat, Elias Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković, Combinatorial optimization and reasoning with graph neural networks. arXiv preprint 10Xiv:2102.09544 2021.
11

---

Published as a conference paper at ICLR 2023
Rich Caruana. Multitask learning. Mach. Learn. 28(1):41-75, 1997. doi: 10.1023/A:
1007379606734. URL httpt:t////..... http:://oooooooo /ttp://....//111.111111111111110
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang. Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint ar2iv:1512.03012. 2015.
Erhan Çinlar. Probability and stochastics, volume 261. Springer, 2011.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer- ence on machine learning, pp. 9990-2999. PMLR, 2016.
Corinna Cortes and Mehryar Mohri. On transductive regression. Advances in neural information processing systems, 19, 2006.
Michael Crawshaw. Multi-task learning with deep neural networks: A survey. CoRR, abs/2009.09796, 2020. URL
hwwpg//bs//0......//// 09796.
Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imita- tion learning. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.). Advances in Neural Informa- tion Processing Systems 32: Anmual Conference on Neural Information Processing Sys- tems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11693- 11704, 2019. URL
Abstract html.
Nima Dehmamy, Robin Walters, Yanchen Liu, Dashun Wang, and Rose Yu. Automatic symmetry discovery with lie algebra convolutional network. Advances Neural Information Processing Systems, 34:2503-25155 2021.
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector neurons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 12200-12209, 2021a.
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J. Guibas. Vector neurons: A general framework for so(())equivariiat networks. CoRR, abs/2104.12229, 2021b. URL
httrg//bs//100 12229.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1810.04410044 2018.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pp- PMLR, 2022.
Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. On the convergence theory of gradient- based model-agnostic meta-learning algorithms. In Silvia Chiappa and Roberto Calandra (eds.), The 23rd International Conference on Artificial Intelligence and Statistics, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy], volume 108 Proceedings of Machine Learning Research, pp. 1082-1092. PMLR, 2020. URL BoLLhttp://pooooee 2allah20a.. html.
ngs mlm. www..prres/e...... wwwppp/ssssle...
Angelos Filos, Panagiotis Tigas, Rowan McAllister, Nicholas Rhinchart, Sergey Levine, and Yarin Gal. Can autonomous vehicles identify, recover from, and adapt to distribution shifts? In Pro- ceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learming Research, pp. 3145-3153. PMLR, 2020. URL http:
pwwppedinceeegnneeoeeee / mlr. www.:rrresssiv...eee
html.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leariig for fast adaptation of deep networks. CoRR, hts//703.034400 2017a. URL http: arxiv. orgg abs 1703. 03400.
12

---

Published as a conference paper at ICLR 2023
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135, PMLR, 2017b.
Alex Gammerman, Volodya Vovk, and Vladimir Vapnik. Learning by transduction. In Proceedings of the Fourteenth Conference on Uncertainty Artificial Intelligence, pp. 148-153, 1998.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan Adams, and Sergey Levine. Why generalization in RL is difficult: Epistemic pomdps and implicit partial observability. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann Dauphin, Percy Liang, and Jennifer Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Confer- ence on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 25502 25515, 2021. URL https: htttps:/pproceediiiee httpp//pcceeednggeeeeeee newr ips cc/paper/ de 2377c955c 4dd9 62 Abstract ctttiraactttll.......
d5 (.55ff15337110
Alex Gittens and Michael Mahoney. Revisiting the nystrom method for improved large-scale ma- chine learning. In International Conference Machine Learning, pp. 567-575, PMLR, 2013,
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 Proceedings of Machine Learning Research, pp. 1856-1865, PMLR, 2018. URL http: httpp//rcceenn...... aalr.press. / v80 haarno j aa Bb. html.
Philippe Hansen- Estruch, Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimulation makes analogies in goal-conditioned reinforcement learning. In Kamalika Chaudhuri. Stefanic Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 8407-8426. PMLR, 2022. URL https: html.
Esansen-Estruch,
/ //proceedings.. mll press / vv1
Zhang-War Wei Hong, Ge Yang, and Pulkit Agrawal. Bi-linear value networks for multi-goal reinforce- ment learning. In International Conference on Learning Representations, 2021.
Nicholas Ichien, Qing Liu, Shuhao Fu, Keith Holyoak, Alan L. Yuille, and Hongjing Lu, Vi- sual analogy: Deep learning versus compositional models. CoRR, abs/2105.07065,, 2021. URL https: / arxiv. org/ abs//bb//// 2205. 07065.
Stefanie Jegelka. Theory of graph neural networks: Representation and learning. arXiv preprint arXiv:2204.07697, 2022.
Thorsten Joachims. Transductive learning via spectral graph partitioning. In Proceedings of the 20th international conference on machine learning (ICML-03), pp. 290-297, 2003.
Ryan Julian, Benjamin Swanson, Gaurav S. Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. CoRR, abs/2004.10194 2020. URL https: / wereeiiiorrr//ge htgg/bbs/2004.20
Leslie Pack Kaelbling. Learning to achieve goals. In Ruzena Bajesy (ed.), Proceedings of the 13th International Joint Conference on Artificial Intelligence. Chambéry, France, August 28 September 3, 1993, pp. 1094-1099. Morgan Kaufmann, 1993.
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn. Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic re- inforcement learning at scale. CoRR, abs/2104.08212, 2021. URL https: /ttp:s////xxvvvorrr/ee abss21044 08212.
Osman Semih Kayhan and Jan van Gemert. On translation invariance in cnns: Convolutional layers can exploit absolute spatial location. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14274-14285, 2020.
Diederík P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1e.. 69800 2014.
13

---

Published as a conference paper at ICLR 2023
Diederik P Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational
bayes.
arXiv preprint
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in deep reinforcement learning. CoRR, abs/2111.09794, 2021. URL https: /atpaaaiivvv.///.//eee/eaaaee.ea wwtxxiiorrg org/
abs/ 2211. 09794.
Nikita Kitaev, tukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv 2020.
preprint arXiv:2001
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsub- ramani, Weihua Hu. Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran S. Haque, Sara M. Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang (eds.), Proceed- ings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 5637-5664. PMLR, 2021. URL http: / pww./eiinceeenseeeeeeeee mlr. www.mmrrr.s.......mm.. pwpesssvv333 /v19/v1119// kohhka. html.
Aitor Lewkowyzz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman- et al. Solving quantitative reasoning problems with language models. arXiv preprint arXiv:2206. 2022.
Lester W Mackey, Ameet Talwalkar, and Michael I Jordan. Distributed matrix completion and robust factorization. J. Mach. Learn. Res., 16(1):913-966) 2015.
Lucas Manuelli, Wei Gao, Peter R. Florence, and Russ Tedrake. KPAM: keypoint affordances for category-level robotic manipulation. In Tamim Asfour, Eiichi Yoshida, Jaeheung Park, Henrik Christensen, and Oussama Khatib (eds.), Robotics Research The 19th International Symposium ISRR 2019, Hanoi, Vietnam, October 6-10, 2019, volume of Springer Pruceedings in Advanced Robotics, pp. 132-157. Springer, 2019. doi: (007/978-33000--555500 org 10. 1007/9788877 3-030-0-----000000 030- 90000-955598-800
URL https:
Sara Meftah, Nasredine Semmar, Mohamed Ayoub Tahiri, Youssef Tamaazousti, Hassane Essafi, and Fatiha Sadat. Multi-task supervised pretraining for neural domain adaptation. In Proceedings of the Eighth International Workshop on Natural Language Processing for Social Media, pp. 61-71, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v 1/2020.. socialnlp-1.3 1.8. URL https: / //aaaaaaanttoooooooo
Lingsheng Meng and Bing Zheng. The optimal perturbation bounds the mmoore-penron inverse under the frobenius norm. Linear algebra and its applications, 432(4))956-963, 2010..
Melanie Mitchell. Abstraction and making in artificial intelligence. CoRR, ab//1102.117777 2021. URL https: / wxw...o..gg...o........ www.//rrg/aaaaaaeoeae....eeaee/e abs 10717.
Andriy Mnih and Russ R Salakhutdinov. Probabilistic matrix factorization. Advances in neural information processing systems, 20, 2007.
Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5). 2008
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Vi- sual reinforcement learning with imagined goals. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Gamett (eds.), Ad- vances in Neural Information Processing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 2209-92200 2018. URL https::/ppocceeeeeee /ttcc////ccccceee//////e neurips. httccc///aa// pttppppppppp/aaasseeeee 7ec69dd4 (ic.5ccfA44L54455555544444770000 Abstraact.. html.
Giambattista Parascandolo, Heikki Huttunen, and Tuomas Virtanen. Taming the waves: sine as activation function in deep neural networks. 2016.
Joaquin Quinonero Quinonero-Candela. Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008.
14

---

Published as a conference paper at ICLR 2023
Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review, 2019. URL atttps://axiic org/ /tss//990 05659.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Benjamin Recht. A simpler approach to matrix completion. Journal of Machine Learning Research, 12(12), 2011.
Scott E. Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual malogy-makingg In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 1252-1260, 2015, URL https: ptttp:://pzoceeddngge new.iiipiignneeissssp.ee newrips...p.... wwtpppppe/c/////aaeeeeeeeeeeeeee e0741335 01 a - html.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL https: httpss:///ppeeeeee...oo net/ hetruoorrmicco By8dvM-R-
Adam Santoro, Sergey Bartunov, Matthew M. Botvinick, Daan Wierstra, and Timothy Lillicrap. One-shot learning with memory-augmenteed neural networks. CoRR, abss/605.06065, 2016. URL
http:/ /ttp::///xxivv org/ abs//ab///a /tbbs/lee 665.. 06065.
John Schulman, Sergey Levine, Picter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust gion policy optimization. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, vol- ume 37 of JMLR Workshop and Conference PrcnnceePProceee Proceedings, pp. 1889-1897. 1889-1897.JMLR.000 2015. URL http:/ ppoceedingg mml. press /v3 wwtssschlllannnccccccc 5 html.
Devavrat Shah, Dogyoon Song, Zhi Xu, and Yuzhe Yang. Sample efficient reinforcement learning via low-rank matrix estimation. Advances in Neural Information Processing Systems, 33:12092- 12103, 2020.
Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B Tenenbaum, Alberto Rodriguez, Pulkit Agrawal, and Vincent Sitzmann. Neural descriptor fields: (3)-equivariant object rep- resentations for manipulation. In 2022 International Conference on Robotics and Automation (ICRA), pp. 6394 6400. IEEE, 2022.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1409. 1556, 2014.
Aman Sinha, Hongseok Namkoong, and John C. Duchi. Certifying some distributional robustness with principled adversarial training. In 6th International Conference on Learning Representa- tions, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceed- ings. OpenReview.net, 2018. URL https: /eattts::oppeeeee //tp///eeeeeeeeeee///eeeeee net/ forum H
Tess E Smidt. Euclidean symmetry and equivariance in machine learning. Trends in Chemistry, (2):82-85, 2021.
Gilbert W Stewart. On the perturbation of aseudo-inverses, projections and linear least squares problems. SIAM review, 19(4):634 662, 1977.
Yixuan Su, Lei Shu, Elman Mansimov, Arshit Gupta, Deng Cai, Yi-An Lai, and Yi Zhang. Multi- task pre-training for plug-and-play task-oriented dialogue system. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.). Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). ACL 2022, Dublin. Ireland, May 22-27, 2022, pp. 4661-4676, Association for Computational Linguisties, 2022. doi: 10.18653/v 06653/v1/2022 acI-long.319. URL https httpppp:////...///
/11.66331l//2 aci long.311111 919.
15

---

Published as a conference paper at ICLR 2023
Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil. Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let net- works learn high frequency functions in low dimensional domains. CoRR, abs/2006.117939, 2020.
URL https::///....
abs /200 10000661733000 10739.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and conslation-equiarrinn neural networks for 3d point clouds. arXiv preprint arXiv:/882.088219, 2018.
Sebastian Thrun and Lorien Y. Pratt (eds.). Learning to Learn. Springer, 1998. ISBN 978-
1-4613-7527-2. doi: 10. 11107/978-1-4615-55299999 1-4615-5529-2. URL https: /ttps://dcciooo wtt.:ii.rg/1..0.0000000
Paul Upchurch, Jacob Gardner, Geoff Pleiss, Robert Pless, Noah Snavely, Kavita Bala, and Kilian Weinberger. Deep feature interpolation for image content changes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7064-7073, 2017.
Vladimir Vapnik. Estimation of dependences based on empirical data. Springer Science & Business Media, 2006.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Oriol Vinyals, Charles Blundell, Timothy P. Lillicrap. Koray Kavukcuoglu, and Daan Wierstra. Matching networks for one shot learning. CoRR, abs/1600.04080, 2016. URL org/
abs erg/abs///
6 04080.
Dian Wang, Robin Walters, and Robert Platt. $ (ma)rrmmSSO()SS-equvvvaaaa reinforcement learn- ing. In The Tenth International Conference on Learning Representationss ICLR 2022, Vir- tual Event, April 25-29, 2022. OpenReview.net, 2022. URL Gorum?id=7F9c0lie
Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Meta-learning without memorization. CoRR, abs/1912 2,.08200 2019. URL
1912. 03820.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-word world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, pp. 1094-1100. PMLR, 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Jian Zhang, Yuanqing Zhang, Huan Fu, Xiaowei Zhou, Bowen Cai, Jinchi Huang, Rongfei Jia, Binqiang Zhao, and Xing Tang. Ray priors through reprojection: Improving neural radiance fields for novel view extrapolation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18376-18386 2022a.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint arXiv:2206. 04301, 2022b.
Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund. Sim-to-real transfer in deep rein- forcement learning for robotics: a survey. In 2020 IEEE symposium series on computational intelligence (SSCI), pp. 737-744. IEEE, 2020.
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for com- positional generalization in control. arXiv preprint arXiv:2203.05960, 2022.
16

---

Published as a conference paper at ICLR 2023
CONTENTS
Introduction
Setup
3
Conditions and Algorithm for Out-of-Support Generalization
3.1
Transductive Predictors: Converting oos to ooc
3.2
Bilinear representations for OOC learning
3.3
Our Proposed Algorithm: Bilinear Transduction
Experiments
4.1
Analyzing OOS extrapolation on analytical problems
4.2
Analyzing OOS extrapolation larger scale decision-making problems
5
Abridged Related Work
Discussion
7
Ethics Statement
10
8
Reproducibility Statement
10
9 Acknowledgments
10
Unabridged Related Work
19
Theoretical Results
21
B.1
Generalization under bounded density ratio
B.2 Extrapolation for Matrix Completion
22
B.3
General Analysis for Combinatioral Extrapolation
22
B.3.1 Proof of Theorem 2
23
B.4
Extrapolation for Transduction
B.4.1 Further Remarks on Theorem
26
B.5
Connections to matrix completion
26
Additional Results
27
C.1 Learning which points to transduce: Weighted Transduction
27
C.2 Additional Results Simple 2-D Example
27
C33 Additional Results on Supervised Grasp prediction
28
C.4 Additional Results on Imitation Learning
29
C.5 Choice of Architecture
30
D Implementation Details
31
17

---

Published as a conference paper at ICLR 2023
D.1
Train/Test distributions and Analytic Function Details
31
D.2
Training details
32
D.2.1
Details of Bilinear Transduction Implementation
................................
32
D.2.2
Details of Training Distributions
33
D.2.3 Details of Model Architectures Architettuees...................
34
D.2.4 Details of Model Optimization
34
18

---

Published as a conference paper at ICLR 2023
A UNABRIDGED RELATED WORK
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 that extrapolate to novel data has been the equivariant neural networks, first proposed by Cohen & Welling (2016). The key idea is that, if it is known there is 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
(g 1) gg. (2)
for all g G (here g () denotes group action), then one can explicitly encode for predictors he satisfying the same identity. Similarly, one can encode for invariances, where h, (g x) = h1 (x) for all g E G.
Deng et al. (2021a) extended the original equivariance framework to accommodate situations where the group G corresponds rotations (SO(3)), and Simeonov al. (2022) proposed neural descrip- tor fields which handle agiidbody transformations SE(3)). For a broader review on how other notions of symmetry can be encoded into machine learning settings, consult Smidt (2021), and Abu- Dakka et al. (2021); Parascandolo et al. (2016) 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) equivari- ances/invariances, there are numerous forms of structure that can be represented as a group symme- try, and for which these methods do not apply. As we show in experiments, there are 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 ex- trapolations has been sequence-to-sequeee processing. Here, the now-renowned Transformer architecture due to Vaswani et al. (2017), (), well as its variants based on the same "attention mechanism" " (e.g. Kitaev et al. (2020)) have become incredibly popular. Supervised by a sufficiently 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 al., 2021), and even quantitative reasoning (Lewkowyez al., 2022). Attention-baat 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).
modeling, especially in the context of natural language
Lastly, a line of research has studied various bilinear models (e.g. Hong et al. (2021); Shah et al. (2020)), motivated by the extension of literature matrix factorization discussed in our abridged related work. However, these methods require additional fine tuning when applicd 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 trans- duction 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 ex- trapolation by embedding tokens of interest in latent vector space on which addition operators can be defined. Zhou et al. (2022) compares the Deep Sets and Transformers approaches as policy architectures in reinforcement learning, finding that neither architecture uniformly outperforms the other.
Distributional Robustness. One popular approach to out-of distrribbuttin learning has been distri- butional robustness, which seeks predictors that perform a family of "nearby" shifted distributions
19

---

Published as a conference paper at ICLR 2023
Sinha et al. (2018); Rahimian & Mehrotra (2019). These approaches are well suited OOD set- tings where the test distributions have the same support, but have differing (but boundedly different) probability densities.
Meta-Learning and Multi- Learning: Meta-learning and multi-task learning methods aim learn policies/predictors that can either zero-shot or very quickly generalize new prob- lems Thrun & Pratt (1998); Finn et al. (2017a); Santoro et al. (2016): Vinyals et al. (2016); Kalash- nikov et al. (2021); Crawshaw (2020); Rosenbaum et al. (2018); Caruana (1997). These methods however usually make strict distributional assumptions about the training 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. (2020). Yin et al. (2019) shows that under distribution shift, meta-learning methods can fail dramatically, and this problem will be exacerbated when supports shift well. Multi-task learning methods often make the assumption that training on some set of training tasks can provide good pre- tre-irrining before finetuning on new tasks Julian al. (2020); Su et al. (2022); Meftah et al. (2020). 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. (2020); Kirk et al. (2021); 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 finetuning setting but rather evaluating zero-shot performance arbitrary contexts. While certain works Simeonov et al. (2022); Wang et al. (2022) do explore ideas of equivariance and invariance in RL and imitation learning, they largely restricted to the either the (3) or SO(2) groups. Moreover, as we discussed in our experiments, bilinear transduction is able to capture more complex notions than equivariance.
out of support goals, dynamics and
The empirical evaluations we perform are in large part on domains in goal conditioned reinforce- ment 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 al. (2017); Nair et al. (2018). While certain architectures Hong al. (2021) do seem to show some level of extrapolation, they still require significant finetuning. 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 al. (2015); Ichien et al. (2021); Lansen-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.
20

---

Published as a conference paper at ICLR 2023
B THEORETICAL RESULTS
This appendix organized into four parts.
Appendix B.1 introduces the bounded density condition, and reproduces a folklore guaran- tee for extrapolation when one distribution has bounded density with respect the other. Appendix B.2 provides a basic guarantee for extrapolation in the context of missing-not- at-random matrix completion depicted in Fig 9, based results in Shah et al. (2020).
Appendix B.3 formalizes a notion of bounded combinatorial density ratio (Definition B.2), in terms of which we can establish out-of-combination extrapolation guarantees by lever- aging 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.1 GENERALIZATION 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 μl, P2, which need not be normalized and sum to one.
Definition B.1 (K-bounded density ratio). Let (41 #2 be two measures over space Q. We say has xs-buunded density with respect (12, which we denote f41 <x, 12, if for all measurable event* A C. Q. /1 [A] S kP][A].
Stating Definition B.1 for general probability affords us the flexibility to write example, IP; <<. F2 as P2 + +3 is a nonnegative measure with total mass + 1 2.
Remark B.1 (Connection to concentrability). The parameter is known in the off-policy reinforce- ment learning literature as the concentratability coefficient (see, e.g. Munos & Szepesvári (2008)), and appears in controlling the performance of rollout policy πj trained on a behavior policy Te by asserting that Px,, <,, IP, T22 where P is, for example, state-action visitation probability under Ti- Remark B.2 (Density nomenclature). The nomenclature "density" refers to the alternative defini- tion in terms of Radon-Nikody derivatives. To avoid technicalities, this can be best when μy and μ12 are continuous densitives over Ω = Id with densities pi (-) and p2(-); e.g. for E (1,2). pa [a[AA] - pe(r)dæ. Then μn Ex u12 equivalent
P1 (I sup (2
S K.
The following lemma motivates the use Definition B.1. Its proof is standard but included for completeness.
Lemma B.1. Let μl, U2 be measures on the same measurable space and suppose that #2 Ex P1. Then, for any nonnegative function &. Hzlø] < (p1(0] 4 In particular, if Drest Ex Diraim then as long as our loss function E(.,,) is nonnegative, R(ha; Drest) S KRe Durai).
Thus, up to R-factor, R(he; Dreat) inherits any in-distribution generalization guarantees for Rho Derain).
Proof. As in standard measure theory (c.f. Çinlar (2011, Chapter 1)), we can approximate any by sequence of simple functions on Φ, where (w) Ein Ca,TT(w E An,), with An,s C 5 and Ca, 2 0. For each &n. we have
kn
12 [2(0))]
Σ Cn.all2(Ann)) SK Cn.llh (An,i) M1 [)()) i=1
The result now follows from the monotone convergence theorem. To derive the special case for Drest and Derain, apply the general result with nonnegative function 6(x) = ( C(y, h (x)) (recall ((.,) -) 2 0 by assumption), #1 Dussnn and μ2 Drest-
'For simplicity, we omit concrete discussion of mcasurability concerss throughout. *Here, v(0) ((w)d)(w)) denotes the integration with respect to μ.
21

---

Published as a conference paper at ICLR 2023
B.2 EXTRAPOLATION FOR MATRIX COMPLETION
In what follows, we derive a simple extrapolation guarantee for matrix completion. The following in the spirit of the Nyström column approximation (see e.g. Gittens & Mahoney (2013)), and proof follows the analysis due to Shah et al. (2020). Throughout, consider
M1
M12]
M
M
M
M
M21
M22]
M
M
where we decompose M, M into blocks (i, E 2)2 for dimension X mÿj
Lemma B.2. Suppose that M rank at most p, M rank p. and
v(i,j) + (,,2), [|M,J M4,J-M;,11 ||e S, and ||M; S M,
where oy(M1,122. Then,
M op op(Mi
M 2 IF s 8c
Proof. The proof mirrors that of Shah et al. (2020, Proposition 13). We shall show below that M is of rank exactly Hence, Shah et al. (2020, Lemma 12) gives the following exact expression for the bottom-right blocks,
M22 = M/ M = ME (M,,) M
where above (.) denotes the Moore-Penrose pseudoinverse. Since || - Milllop S [M, Mi ||r 6 (M/ Weyls inequality implies that Mn is rank p (as promised), and [M]] llop S (M; Similarly, as | Mizllop S (M,(M) & M/2, so ||M S
M. Thus,
|M M52lFF < M21 - M |F M 11 lop |M + MM MM - M12lle
Miellop||MM lop [M]] (M1
5cM
20, (M, + M² M (M)))*
Next, using perturbation bound on the pseudoinverse due to Meng & Zheng (2010, Theorem 2.1),
[M]] - IlI S |M Mille (|x l II(M;)) 11p
S.. (o,
Therefore, we conclude
5
4M2
M
M2 - M S
+
S 8c
(Mx o,(Mi
o (M
B.3 GENERAL ANALYSIS FOR COMBINATIORAL EXTRAPOLATION
We now provide our general analysis for combinatorial extrapolation. To avoid excessive subscripts, we write X' = W x V rather than = X1 x X₂ as in the main body. We consider extrapolation under the following definition of combinatorial support.
Definition B.2 (Bounded combinatorial density ratio). Let D, D' be two distributions over a product space W X V. We say D' has x-bounded combinatorial density ratio with respect D, written as 'D' E combobb D. if there exist distributions Dw and Dv.j. i,j E {1, over W and V. respectively, such that Dioj : Dw, & Dv.j satisfy
E(s)) Dies <x Cx D, and D' Ex Eij=1,2 Disj-
Unlike Shah et al. (2020), interested in the Frobenius norm error, we elect for the slightly sharper hound of Meng & Zheng (2010) above than the classical operator bound of Stewart (1977).
22

---

Published as a conference paper at ICLR 2023
For simplicity, consider scalar predictors, as the general result for vector valued estimators be obtained by stacking the components. Specifically, we consider a ground-truth predictor h, and estimator h of the form
h. (fe:ge), h (f.0) f., - R2, 9u.9 R..
(B.1)
Lastly, we choose the (scalar) square-loss, yielding the following risk
[h; D) == ((w.))-D(hh. (w, v) h(w, v))
Throughout, we assume that all expectations that arise are finite. Our main guarantee as follows. Theorem 2. Let D, D' be distributions on W xV satisfying D' <x,comb D, with corresponding factor distribulions Dw,, and Dv.j. i, E 2). Define the effective singular value
of = (ED [f. (w) f4 (w) D)po (EDv,, (v)
(B.2)
and suppose that maxisii.s2 EDigy 14, (w, 2)) = M2. R(h D) S
a T
M R(h; D) poly ni
M.
R(); S D): 2²2 + of
o,
B.3.1 PROOF OF THEOREM
First, let us assume the following two conditions hold; we shall derive these conditions from the conditions of Theorem at the end of the proof.6
v(i.j) * (2.
R(s: se.
ED [h, u) S M2 E<0./2.
(B.3)
Our strategy is first to prove a version of Theorem 2 for when W and V have finite cardinality by reduction the analysis of matrix completion in Lemma B.2, and then extend arbitrary domains via a limiting argument.
Lemma B.3. Suppose that Eq. (B.3) hold, and in addition, that W and V have finite cardinality. Then,
R(h; ((02)))) = |M2 - MM < 64c2 M
Proof Lemma B.3. By adding additional null elements either W or V, we may assume without loss of generality that [W] | = d, and enumerate their elements (wi,..., wa} and (vi,..., va)- Let Pi.a Phw--D.. [w w.] and 9ý.b ProwDy., vz]- Consider matrices M. M
gadx2d with d x d blocks
(Mi)aa = 9i, h(war vo). (Mij)as = VP..a9g. h, h. vb).
We then verify that
[|
M 1
Σ Ps.a P.a9n,(h((aa vb) - h, (wa, un))2
(B.4)
ED.o. [((w v) hw r)) = R(h; Disi):
and thus |M M < for (i.j) # (2. Furthermore, define the matrices A,, B, via
(A:)a == (Pral(wa)) (B,)s (V,,
and define A;, Bj similarly. Then,
A Az
B B2
M
M
B
'Notice that here take 2 as upper bound of Exi,, [b,,,,, v) rather than pointwise upper bound in Theorem 2. This for convenience limiting argument below.
23

---

Published as a conference paper at ICLR 2023
showing that rank (M1), rank( & p. Finally, Eq. (B.2).
op(Min) T) 2 (Ai))
on ((A1) A1) (B (B))
op Σ Plaf (((aa))(()))))) Op Σ [aLLs(vo)(((o)
= o (EDw., (f D)o, (EDv, (((o)(()) = o².
Lastly, we have
Σ
Piohh v6)² ED,,, (), S M².
a.b
Thus, Eq. (B.4) and Lemma B.2 imply that
Rh; D282) = M2 M22 18 I 64c M
ot
Lemma B.4. Suppose that Eq. (B.3) hold, but unlike Lemma B.4, W and V need not be finite spaces. Then, still, it holds that
R(h; = |Ma - Mip S 64e 2 M of
Proof of Lemma B.4. For N, define he, (Srmn Ir and hn fx Jn Ins 9x. are simple functions (i.e. finite range, see the proof of Lemma B.I) converging to f.y f...8 9. Define
(n,sNn); where
2²,n = (f,
D), (r.))...()))...(
M max ED (w,
+n (.)7(2.2
2 max R(hn; Dioj). (2,2(,
By the dominated convergence theorem?.
lim inf 2 +,n 2 oz
lumsup lim sup M2,. S M2
lim sup 2 s e. a21
n21
n21
In partieular, as ² < 02/4 then applying Lemma B.3 for sufficiently large,
M 2
R(hm; D202) S 64
kmn
Indeed, for any fixed all of f.,, In, fa.ns Ir.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 finite domains; this enables the application of Lemma B.3 applies. Finally, using the dominated covergence theorem one last time,
M1 R(h; D282) = lim R(hn; :(o2) I lim 64
M σf
cn < 64
n-
n21
04
We can now conclude the proof of our proposition.
Via standard arguments, one can construct the limiting embeddings fr.n, In Gn. Ir.n in such way that their norms dominated by integrable functions.
24

---

Published as a conference paper at ICLR 2023
Proof of Theorem 2. As D' <x En, Dies and E154(2.2) Dies <n D, Lemma B.1 and additivity of the integral implies
R(h;D') S kR;h; 2202) +K E R(n;)in))
(1.)+(222)
kk D202) + K R(;hD)
(B.5)
Moreover, setting 2= =R=hh D), have
R(h; Diso)) S
E
Diss) S <R(h; =
max (2,)+(2.2)
(13))
Thus, for D): <
0² 1x
Eq. (B.3) holds and thus Lemma B.4 entails
R(h; D202) S 64c² M of
M 64h D)
o
Thus, combining with Eq. (B.5).
R(h:D') s <<:n(D (1 64 Mr
of
completing the proof.
B.4 EXTRAPOLATION FOR TRANSDUCTION
Leveraging Theorem 2, this section proves Theorem 1, thereby providing formal theoretical justi- fication for predictors of the form Eq. (3.5).
Proof of Theorem I. We argue by reducing to Theorem 2. The reparameterizztion of the stochastic predictor he in Eq. (3.1), followed by Assumption 3.2 allows to write
R(ea; Dirain) ExwDDuain Eywhhl=>) E(y, h. (x:)
EewDeram Egrn D [(ho x' x2), hx),))
EnnDDuain Eg'n D D (ho x' x²), Tu xz))))).
In the above display, the joint distribution of (z x x')xx) is precisely given by Derain (see Eq. (3.2)).
Hence,
R(ho; Derain) = E (A, x'), h. (Ar, x'))
Further, E(y, lly 11112 the square loss and decomposes across coordinates.
K
R(ho; Derain)
E E
x') - x'))
(B.6)
By the same token,
R(hg: Drest)
Σ Ep.. (ho x') (Az, x²²)
To conclude the proof, we remain to show that for all & E [K], we have
M Ep..m (he (Ax, 2²) h (A,,,''))) < Cprob ED (ho (AI, x') h h,, (Ax,
(B.7)
where Oproh = 2 + of
Indeed, for each & E [K], we have
(Eq (B.6) S R(ho; Derain)
(by assumption)
o²
ED,,ain (ho, T 2')x
4k
Hence Eq. (B.7) holds by invoking Theorem 2 with the correspondences W + AX, V + X, n
0, M2 M and This concludes the proof.
25

---

Published as a conference paper at ICLR 2023
B.4.1 FURTHER REMARKS ON THEOREM
The singular value condition, min(o,(EDa.)) (fn.). op (IE 19. xg. 2 o² > 0, mirrors the non-degeneracy conditions given in the past work in matrix completion (c.f. (Shah et al., 2020)).
The support condition sup, Ae.." [ x')l S M is a mild boundedness condition, which (in light of Theorem 2) can be weakened further to
max Ep Ih, (Ax, z') S M 1SWJ2
where Dioj are the contituent distributions witnessing Duest <n,comb D
The final condition, Rha; Dtrain) S a is mostly for convenience. Indeed, as M 2 a Th
and 2 1, then as soon as Rho;; Derain) > An better than
o²
our upper-bound on R(ho; Drest) is no
64 M RM
o2 2 6M².
4
which is essentially vacuous. Indeed, also inflate M and stipulate that supan.a' (ho(Ar,2)) I V6M, we can remove the condition R(ho: Derain) S o² ix
together.
B.5 CONNECTIONS MATRIX COMPLETION
hox
m4
(ax(
[xx(x²
11
M11
AX
sampled
missing
hox
faxx
90,k
(a) Matrix completion
(b) Connecting 00C to matrix completion
Figure 9: Illustration of bilincar representations for Icarning, and connection matrix completion. (a) An example low-rank matrix completion, where both M and Mд have rank-p. Blue: support where entries can be accessed, green: entries are missing. (b) An example that low-rank structure facilitates certain forms of ooc, ie. for each & E [K]. the predictor can be represented by bilinear embeddings as ho (Az, x') (Jo,,(Ar) x0.x.... (x))))
Building upon Definition B.I, Assumption 3.1 introduces a notion of bounded density ratio between Derain and D Deess in the OOC setting. Take the discrete case of matrix completion as an example, as illustrated in Fig 9, the training distribution of (Az, x²) covers the support of the (1, 1), (1, 2). (2, 1) blocks of the matrix, while the test distribution of (Az,, x²) might be covered by any product of the marginals of the 2 x 2 blocks. With this connection in mind, it is possible to establish the ooc guarantees on Duest as matrix completion tasks, 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 standard setting studied extensively in the matrix completion literature, (Shah al., 2020; Agarwal et al., 2021b; Athey et al., 2021) for example. Our results in Theorem can be viewed as 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.3 corresponds to the nondeegenyracy 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).
26

---

Published as a conference paper at ICLR 2023
C ADDITIONAL RESULTS
C.I LEARNING WHICH POINTS TO TRANSDUCE: WEIGHTED TRANSDUCTION
There may be scenarios where is beneficial to more carefully choose which points to transduce rather than comparing every training point to every other training point. To this end, we outline more sophisticated proposal, weighted transduction, that achieves this.
Sampling training pairs uniformly may violate the low rank structure needed for extrapolation dis- cussed in Section 3.2. Consider a dataset of point clouds of various 3D objects in various orientations where the task is 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, may be useful to identify which training points to transduce, rather than transduce between all pairs of points during training.
Hence, we introduce weighted transduction described in Algorithm 2, which identifies which train- ing points to transduce, rather than transduce between all pairs of points during training. During training. a transductive predictor hg is trained together with scalar bilinear weighting function wa- he is trained predict values in the bilinear transduction algorithm, yet weighted by wa- we 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 we. or jointly optimize wis and hg. test time, we select anchor point based on the learned weighting function, and make predictions via hg. We show results on learning with weighted transduction on analytic functions in Appendix C.2 and learning how to grasp from point clouds in Appendix C.3.
Algorithm Weighted Transduction
1: Input: distance parameter p. training set (I1,81), (In, No.. regularizer Reg(.)
2: Train: Train 0 loss
(())= Σ wo(a x,, ((o(m: -,,,,,), Reg(O)
3: Test: each new Wtest, predict
he (Frest Tz, Di), where i i~ [P]i = i] x we (eeus x1)
C.2 ADDITIONAL RESULTS ON SIMPLE 2-D EXAMPLE
Weighted transduction with a 2-D example We also considered 2-D analytic function as shown in Fig 10. This function has 2-D input (21,22) and 2-D output (y1, The function is con- structed by sampling a tiling of random values in the central block (between 1) and (5, 5)). The remaining tiles are constructed by shifting the central tile by units along the 21 12 direction. For tiles shifted in the vi -drrectio ((6, 0) or (0, 6)) the label has the same first dimension and a negated second dimension. For tiles shifted in the -irdirection (66 0) or (0, 6)) the label has the same second dimension and a negated first dimension. This has 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 transduc- tion 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 extrapola- tion 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 find symmetries in the data with proposed method.
27

---

Published as a conference paper at ICLR 2023
Unweighted Bilinear Trensductoon
Weighted Bilinerr Transduction
Ground Truth (TRAIN)
Ground Truth (DOS)
Inductive NN
Blinear
ys
Figure 10: Predictions for the 2-D analytic function from (x1, (y1, y2)- (Top) y1 output values for inputs (2), (Bottom) y/ output values. (Left to Right:) In-distribution ground truth values, 0OS ground truth values, neural network predictions, bilinear predictions, bilinear transduction predictions and weighted transduction predictions. Transduction weighting is important in this domain and it is able discover problem problem, since bilinear prediction directly doesn't work, the OOC
symmetry. While this may seem like view on this problem does have low rank structure, while OOS view does.
C.3 ADDITIONAL RESULTS ON SUPERVISED GRASP PREDICTION
Qualitative grasping prediction results For qualitative results of all models on the grasping point prediction task, please see Fig 11.
Bilinear transduction is comparable to an SE(3) equivariant architecture on OOS orientations and exceeds its performance on OOS scales We compare bilinear transduction method that builds in SE(3) equivariance domain knowledge into its architecture and operates on point clouds. Specifically, we compare to tensor field networks Thomas 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, bottles or teapots and on sep- arate set of scaled mugs, bottles teapots by in-distribution values described in Appendix D.1. We test extrapolation to mugs, bottles or teapots with out-of-sample orientations and scales sampled from the distribution described in Appendix D.I. We adapt the Tetris object classification tensor field network to learn grasp point prediction by training the mean of the last equivariant layer to predict grasping point and applying the mean squared error (MSE) loss. Following the implemen- tation in Thomas et al. (2018), we provide six manually labeled key points of the point cloud that uniquely identify the grasp point: two points 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 fix these architectures one layer and units. Under various rotations, the tensor field network is able to extrapolate. For scaling however, the network eгror increases significantly, which is in accordance with the fact that scale transformation guarantees do exist in SE(3) architectures. Bilinear transduction can extrapolate in both cases, is comparable to tensor field networks on 0OS orientations and exceeds its performance on OOS scales.
Weighted grasp point prediction The grasp point prediction results for three objects described in Table I are trained with priviledged knowledge of the object types at training time. Le., the predictor training 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 training time we can learn weights and a predictor that perform as well as with priviledged training (Table
28

---

Published as a conference paper at ICLR 2023
Table Comparing bilinear transduction on grasp point prediction with tensor field networks Thomas et al. (2018), an SE(3) equivaraint network for point clouds. We report mean and standard deviation for for 50 oos samples.
Task
Tensor Field
Linear
Neural Net
Transduction
Ours
Mug rotation
0.009 1 0.003
0.005
0.001
0.066
0.035
0.005
0.001
0.01
0.003
Mug scale
0.027 0.012
0.006
0.001
0.004
0.001
0.004 0.001
0.004 0.001
Bottle rolation
0.008 1 0.004
0.004 1. 0.004
0.085
0.046
0.011
0.026
0.004 0.002
Bottle scale
0.019 0.008
0.003 0.001
0.003
0.001
0.01
0.006
0.003 0.001
Teapot rotation
0,011 + 0.006
0,003 = 0,001
0.055 = 0.03
0.017 = 0.012
0.006 0.002
Teapot scale
0.093 + 0.039
0.003 = 0.001
0.003 0.001
0.003 0.001
0.003 0.001
For grasp point prediction, during training, we train a weighting function wa(x, - T,, -4, 24) to predict whether H; should be transduced to 2, based on object-level labels of the data (transduce all mugs to other mugs and bottles to other bottles). Then we learn predictor hg weighted by fixed wy as described in Algorithm
At test time, given point xy we select the training point 2; with highest value Wr( Fj; .ii) and predict yj via ho (2r about the object category at test time.
Ti), without requiring privileged information
Figure 11: oos grasping point prediction for bottles, mugs and teapots. Each row represents one test datapoint. Columns (Left Right): linear, neural network, transductive and bilinear model predictions (red) and ground truth (green) grasp points. We show that baslines may predict infeasible grasping points (e.g. in the middle an object) whereas bilinear transduction makes near accurate predictions.
C.4 ADDITIONAL RESULTS ON IMITATION LEARNING
Qualitative decision making results For qualitative results of all models on imitation learning domains, please view videos in the supplementary zip file.
29

---

Published as a conference paper at ICLR 2023
Grasping point prediction on three objects with various rotations, translations scales with ground truth weighted training. standard training and learned weighted training for bilinear transduction with a fixed architecture. By training weighting function with fewer labels, learned weighted bilinear transduction achieves similar performance to ground truth weighted training.
Table 3:
Ground Truth
Bilincar Transduction
Learned Weighted
Weighted
(Train)
0.007 0.004
0.01 0.005
0,006 0.003 0.024 = 0.014
(OOS)
0.027 = 0.014
0.068 = 0.041
C.5
CHOICE OF ARCHITECTURE
Periodic activation functions We provide further results on the grasp point prediction and imi- tation learning domains Table 4. This experiment similar to the one reported Table I, but with fourier pre-processing as the initial layer in each model. For the various domains, adding this pre pre-procsssing step beneficial for some architecture-domain combinations, but does not surpass bilinear transduction on all domains for any method. Moreover, in most domains bilinear trans- duction achieves the best results, or close results to the best model. For implementation details see Section D.2.
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 fixed architecture parameters over three seeds.
Table Mean and standard deviation over prediction (regression) or state (sequential decision making) error for OOS samples and over hyperparameter search with fourier pre-processing. While fourier activations are useful for some combinations of models and domains, it is not beneficial across all.
Task
Expert
Linear
Neural Net
DeepSets
Transduction
Ours
Mugs
0.119 + 0.05
0.219 0.114
0.3 = 0.148
0.086 + 0.081
Bottle
0.099 0.041
0.139 0.075
0.158 0.096
0.019 0.011
Тeapot
0.112 + 0.045
0.216 0.107
0.301 0.14
0.041 + 0.026
All
0.142 + 0.064
0.252 0.119
0.33 0.137
0.047 + 0.033
Reach
0.006
0.008
0.007 +: 0.006
0.075 0.082
0.087
0.092
0.117 0.176
0.008 0.008
Push
0.012
0.001
0.182 + 0.13
0.113 0.124
0.121 0.091
0.17 0.12
0.014 + 0.006
Slider
0.105 1 0.066
0.38 1: 0.032
0.151 0.081
0.241 0.077
0.134 0.114
0.256 1: 0.214
Adroit
0.035 + 0.015
0.551 + 0.508
0.398 + 0.435
0.373 + 0.195
0.341 + 0.218
0.365 + 0.31
Neural Net
DeepSets
reach fourter
push
push fourier
0.5
Transduction
Ours
0.4
reach
0.3
reach fourler
0.2
push
0.1
push fourier
0.0
depth hilden 32
depth hidden 512
depth hidden 1024
depth 2 heldan32 32
depth hidden 512
depth hidden 1924
depth hinden 32
depth hidden 512
depth hilden 1024
depth hitson 32
depth hidden 512
depth hidden 1024
Figure 12: Heatmap complementing Table showing mean ooS emror on Meta- World. While some architec- tures can be suitable for some domains, our method is more robust to hyperparameter search and almost always achieves low error.
30

---

Published as a conference paper at ICLR 2023
Table 5: We report mean and standard deviation error over evaluation samples and three seeds. For each domain, the first row is evaluated on in-distribution samples, the second row is on OOS samples. We show (I) bilinear transduction performs well in n-distribution well OOS (2) our algorithm stable across several seeds.
Task
Expert
Linear
Neural Net
DeepSets
Transduction
Ours
Grasping (Train)
0.06 + 0.059
0.006 0.004
0.006 0.003
0.007 0.003
Grasping 00S
0.141 0.114
0.124 0.074
0.099 0.07
0.024 0.014
Reach (Train)
0.006 + 0.008
0.004 + 0.005
0.005 + 0.005
0.005 + 0.005
0.004 + 0.005
0.005 + 0.005
Reach (0OS)
0.007 + 0.006
0,013 + 0.015
0,088 + 0.075
0,007 = 0.007
0.007 = 0.006
Push (Train)
0.012 0.001
0.22 + 1.103
0.046 0.049
0.01 = 0.001
0.022 0.011
0.016 = 0.007
Push O0S
0.26 + 0.063
0.355 + 0.138
0.188 + 0.124
0.337 + 0.151
0.022 + 0.016
Slider (Train)
0.105 + 0.066
0,425 + 0.077
0.191
0.079 + 0.037
0.32 0.072
0.17 = 0.092
Slider O0S
0.512 + 0.06
0.438 0.1191
0.385 + 0.288
0.55 = 0.209
0.163 + 0.126
Adroit (Train)
0.035 + 0.015
0.66 + 0.596
0.139 0.231
0.103 0.18
0.129 + 0.215
0.047 = 0.024
Adroit oos
0.337 + 0.075
0.703 + 0.561
0.422 + 0.241
0.773 = 0.796
0.198 = 0.255
D IMPLEMENTATION DETAILS
D.1 TRAIN/TEST DISTRIBUTIONS AND ANALYTIC FUNCTION DETAILS
Analytic functions:
We used the following analytic functions for evaluation in Section 4.1.
Figure 4a. Periodic mixture of functions with different periods: We consider the period growing mixture of functions, and do similar function but with two functions with different periods repeat- ing. Let us define x, mod 9, then the
3m
sin(10 * x.)
I, < 1
Zm (sin(10) + (Ep 1))
if < <2 if <3
f(x)=
Zm
(sin(10) + (Ee 2)
(D.1)
2m
sin(10 * (x. -3)
2
if < I, <5
Zm (sin(10) + Eu - 1))
if < 2,5<x,5 <
Zm (sin(10) + 1 2-- 2)
if <
We can see that this function mixes different periods, to show that our framework can deal with varying functions as well.
Figure 4b. Sawtooth function: we use classic sawtooth function
f(x) (xx mod Period) Amplitude
if (L Period |) mod 2 = 0
(mplltude - (z mod Period)
Period
Amplitude Period
if Period
mod 2 == (D.2)
Figure 4c. Randomly chosen polynomial: We sampled random degree 8 polynomial: f(x) (x ( 0.1)(x + 0.4)( + 0z7) 0z5) - 1z + 1.75) 1)(x - 1.2)
Figures & 7. Periodic growing mixture of functions: Let define x. = x mod 3 and xm = L
xm
Zm sin 10 x.))
x, if < Xy <
f(x)
Zm (sin(10) + (z,, - 1))
(D.3)
Zm (sin(10) + + (xo - 2))) if 2 < Fy < 3
Figure 6. Shift equivariant mixture functions: Let us define Ve = x mod 3 and 2m =
11
Fmm 3 sin(10 x.))
if F. <
f(x) = Xm * (sin(10) + (x., - 1)) (sin(10) + 1 (x,, 2)2)
if < JT < 2
(D.4)
if < < 3
This represents mixture of different functions sinusoid function, linear function and quadratic function that repeat and are equivariant on shifting. These are meant to show that bilinear transduction capture equivariance.
31

---

Published as a conference paper at ICLR 2023
Grasp point prediction: A training dataset is generated from rotating, translating and scaling bottle, mug and teapot from the ShapeNet dataset Chang et al. (2015) with various parameters to predict a grasping point in RS from their point clouds. A test dataset is constructed by applying the same transformations to the objects with OOS parameters. The point cloud in IR N x3 is represented by a 12-tuple containing the point cloud's mean point and its three PCA components in TR3 In Fig 13 we 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.005) and to the grasp point label sampled from (00, 0.002). During training. objects are rotated around the 2 axis E [0, 1.27], translated by (z. y) E [O, 0.5] [0. 0.5 and scaled by a factor of a E [0.7, 1.3). At test time these transformation parameters are sampled from the following ranges: [1.2x, 2x). (xr, y) E [0.5, 0.7] x [0.5, 0.7] and a E [1.3, 1.6].
Meta- World: We evaluate on the reach-v2 and push-v2 tasks. For reach-v2. we reduce the ob- servation 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- ing, the 3D target position for reach-v2 is sampled from range [0, 0.4] [0.7, 0.9] x [0.05, 0.3] and at test time from range [-0.4,0] 0J x [0.7, 0.9] x [0.05, 0.3). The gripper iniiial position is fixed to [0, 0.6, 0.2). For push-v2 the observation space is a 10-tuple of the end effector 3D position, grip- per openness, object 3D position and target 3D position. The action space remains the same as in reach-v2. The training 3D target position is sampled from [0, 0.3) x [0.5, 0.7] x {0.01} and at test time from [-0.3, 0] x [0.5, 0.7] x {0.01}. The gripper initial position is fixed to [0, 0.4, 0.08) and the object initial position to [0, 0.4, 0.02]. For both environments, expert and evaluation trajectories collected for 100 steps.
We evaluate on the object relocation task in the ADROIT hand manipulation bench- mark. 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 ori- entation of the wrist, and 24D for finger positions commanding PID controller. During train- ing, the target location is sampled from [0, 0.3] x [0, 0.3] x 00115, 0.35] and at test time from [-0.3, 0) x 0.3, 0] x [0.15, 0.35]. Expert and evaluation trajectories are collected for 200 steps.
Adroit hand:
Slider: We introduce an environment for controlling brush to slide an object to 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 training, the object mass is sampled from [60, 130) and test time from (5, 15). Expert and evaluation trajectories are collected for 200 steps. In Figure 14 we plot the actions as function of time steps for demonstrations, baselines described in Section and Bilinear Transduction for a set of fixed architecture parameters.
D.2 TRAINING DETAILS
D.2.1 DETAILS OF BILINEAR TRANSDUCTION IMPLEMENTATION
We describe the bilinear transduction algorithm implementation in detail. For grasp point predic- tion, during training we uniformly select an object E (bottle, mug, teapot), and two instances of those objects, O;. 0y that were rotated, translated and scaled with different parameters. We learn to transduce o; to 0y. L.e learn ho o,, o;) to predict yj, the grasping point of o,. At test time, given test point zr with an unknown object category, we select a training point x; with an "in-distribution" difference æy- - Is, as described in Algorithm 1. We select p to be the closest difference x; -z, to an in-distribution difference over all training points x,, The differences generated by the optimization process are sampled uniformly to generate a smaller subset for comparison.
For the coal-conditioned imitation learning setting, during training we uniformly sample trajectories T, {s} aT) for horizon T specified in Appendix D.1 and T, {(s/ af IT We further uniformly sample a time step from [0, T). We learn to transduce state s) E T. to sf E Tj. Le., leam (sg s/. sf) to predict action a.. At test time, we select an anchor trajectory based on the goal (Meta-World, Adroit) mass (Slider) gj. This is done by selecting the training trajectory т,
32

---

Published as a conference paper at ICLR 2023
40
10
46
11
dh
Figure 13: Reduced point cloud feature space in R 12 From left right: mean and three PCA components. Top: in distribution, bottom: OOS. Blue: bottles, red: mugs, green: teapots.
Figure 14: Slider torque actions function of time steps. Each color represents different trajec- tory and mass. From left to right, top to bottom: in-support demonstrations and Linear, Neural Net, DeepSets, Transduction and Bilinear Transduction policies for oos masses.
with goal or mass that generates an "in-distribution" difference gj - Di with the test point, as described in Algorithm 1, We select p 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 training data. We then transduce each state in T; to predict actions for the test goal or mass. Given test state s/ we predict af with ha(s) st st), execute action af in the environment and observe state - We complete this process for I' steps to obtain Tj.
D.2.2 DETAILS TRAINING DISTRIBUTIONS
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 demonstra- tions, sequences of state-action pairs, { 31) 11)) where horizon 7' for each domain is specified in D. For Meta- World reach and push, N = 1000. For Slider and Adroit = 100. All domains were evaluated on 50 in- in-distribution out-of-sample points and 50 oos points. For each domain, all meth- ods were trained on the same expert dataset and evaluated on the same nn-distribution out-of sampae and OOS sets.
33

---

Published as a conference paper at ICLR 2023
We collect expert demonstrations follows. For grasp point prediction, we 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 2d 2dlaa444222200055555775663333333 and Utah teapot
mug
In 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., 2018)
Meta-Worid,
The weighting grasp point prediction network was provided with with 300 labels of positive pairs be transduced and 300 negative pairs labeled with binary labels. This a slightly more relaxed condition than labeling 1000 objects as in the standard version.
D.2.3 DETAILS OF MODEL ARCHITECTRRES
For analytic domains, we use MLPs (both for NN and bilinear embeddings) with 3 layers of 1000 hidden units each with ReLu activations. We train all models with periodic activations of fourier features as described in Tancik et al. (2020).
Linear model: data points x are processed through fully connected layer output predictions y. Neural Network:
data points r are processed through an mlp with & layers, units each and ReLU activations for hidden layers to output predictions y. DeepSets: we refer 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 mips with & layers, 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- ing point zj and H; -y (for training or test point x.)) are concatenated and processed through an mlp with & layers, n units each and ReLU activations for hidden layers, to predict yı- Our architecture Bilinear Transduction: training point 2, and Ary 2; I. (for training or test point 2;) are embedded separately
feature space in Rd where d is the dimension of yc, by two mips with k layers, n units each and ReLU activations for hidden layers. Each predicted element of y, is the dot product of m-sized segments of these embeddings described in Eq. (3.5).
Further, we evaluate for all models with Fourier Embedding as a pre-preee pre-processing step. We process inputs æ, target or mass Ar through linear model which outputs x40 the number of then multiply values by æ and process through the sin function. This layer is optimized as the first layer of the model.
inputs. We
In Table I we search over number of layers k E {2, (2, and unit size n E {32, 1024]. The bilinear transduction embedding size is 32. In Tables and 5 we k 2, n 32 and m 32.
The weighting function used for grasp point prediction is the bilinear transduction architecture with k = 2, = 128 and m = 32.
D.2.4 DETAILS OF MODEL OPTIMIZATION
We train the analytic functions for 500 epochs, batch size 32, and Adam optimizer with learning rate
1e 4. We optimize the mean squared error (MSE) loss for regression.
We trained all more complex models for 5k epochs, batch size 32, with Adam (Kingma & Ba, 2014) optimizer and learning rate le-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 configuration of hyperparameters was ran and tested on one seed. We demonstrate the stability of our method across three seeds for fixed set of hyperparameters in Table 5.
We train the weighted grasp point prediction for 5k epochs, batch size 16, and Adam optimizer with learning rate le - 4. We optimize the MSE loss between the output and ground truth binary label indicating if training point should be transduced another training point. The weighting function did not require further finetuning jointly with the bilinear predictor.
34


Reviews:
[{"review_id": "eXuASSzuU5", "reviewer": "Reviewer_dNKc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Paper tackles the out-of-support extrapolation problem where the input at test is out of the support of the training distribution. A bilinear model is trained to predict the output $h(x)$ using $\\bar{h}(x’, \\Delta x)$ where $x’$ is some chosen anchor point and \\Delta x is the distance of $x$ to the anchor point. Given an unseen test point, the model is able to extrapolate (in certain situations) by using an anchor point seen in training and the distance from it. \n", "summary_of_the_review": "While the goal of the work is to explore a “general setting does not assume apriori access to distribution shift or model class”, I believe their assumption of bilinear transduction is still strong and is able to extrapolate only in very specific cases. Most of the experiments essentially require symmetries to solve the out-of-support task (for example translation/rotation/scale invariance/equivariance) but such invariant/equivariant methods are not compared with. ", "strength_and_weaknesses": "Out-of-support extrapolation is an interesting but challenging task and the paper provides a novel way of  looking at the problem via a reformulated out-of-combination task. \n\nMy major concerns:\n\n1. While the goal of the work is to explore a “general setting does not assume apriori access to distribution shift or model class”, I believe that their assumption of bilinear transduction is essentially assuming apriori knowledge about the type of distribution shift or model class. While the assumption is weaker than translation invariance/equivariance, it is still strong and is able to extrapolate only in very specific cases. \n\n2. Many of the experiments (periodic functions, push/reach and other robotic manipulations) only require translation invariance/equivariance for out-of-support extrapolation but such models are not compared with. For point-cloud tasks, out-of-support extrapolation requires rotation/translation/scale symmetries but it is hard to judge if these symmetries are learnt automatically from the method as the input to the model is not the point cloud itself, but low dimensional PCA features and the mean of all the points. Adding these symmetries for point-cloud tasks has been explored a lot (e.g., [1]) and should be compared with. \n\n3. I think the approach will have problems with high-dimensional inputs because modeling of each dimension independently is not accurate for most high-dimensional inputs (like images) unless one uses PCA. If PCA or other such techniques is required for the approach, I think this should be discussed along with their approach. \n\n\n\nOther questions/comments:\n\n1. The method is related to translation symmetry due to the $\\Delta x$ term and is able to learn translation invariance/equivariance. But it is not clear to me how much more general it is without using additional data augmentations or PCA. For example, can one expect extrapolation to other symmetries (rotation, scaling) without data augmentations and PCA? \n2. Most of the experiments consist of symmetry due to which extrapolation is possible. What are other types of out-of-support cases that can be solved with the current approach but do not fit into translation/rotation/scale invariance/equivariance?\n3. Please add the difference between train and test distributions for each experiment in the main text. \n\n4. In Meta-world and Adroit benchmarks, please compare with a translation equivariant model. \n5. For point cloud classification, please compare with methods that use the point clouds as inputs directly and not the low dimensional features. \n6. For Slider experiments, it would be interesting to see the ground truth slider torque dynamics for different object masses to see the difference between train and test, and understand how the method extrapolates. \n\n\n\nReferences:\n\n[1] Thomas, Nathaniel, et al. \"Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.\" *arXiv preprint arXiv:1802.08219* (2018).\n\n**Post-rebuttal**\nThe authors have adequately addressed my concerns and I am increasing my score. As in my original review, I think the paper provides a novel way of looking at the challenging out-of-sample problem, is well-formulated and easy to read. After rebuttal, authors have convinced me that the approach empirically extrapolates in many different settings without just using equivariance. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_J8jduOhap", "reviewer": "Reviewer_gXKu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper first discusses out-of-combination scenario, which is a special case of out-of-support scenario. Then the paper propose to use anchor points to convert OOS case to OOC case for improving generalization.", "summary_of_the_review": "The paper has some novelty in recognizing a special case of o.o.s generalization and propose a method to tackle this case. However it is not clear to me how much percentage of real-world o.o.d problems belong to this special case. The proposed method seems to require strong pre-condition ( stationary data distribution w.r.t anchor points) to work", "strength_and_weaknesses": "Strength:\n1. Detailed analysis of different o.o.d scenarios.\n2. Under certain constraints, the proposed method of using anchor points works in improving generalization performance.\n\nWeaknesses:\n1. It looks like to me the proposed method only works for a very limited scenario, in which o.o.s can be easily converted to o.o.c. For example, in the tasks chosen by authors, the data is almost stationary with respect to the anchor point, meaning the data distribution now is only on delta x. I wonder if any real-world o.o.d problems are like this. For example, in image classification task, the test images can be o.o.d because of unseen lightning conditions, unseen camera angles, unseen background and etc. I wonder if the proposed method can be extended to any more complex datasets where the stationary property does not hold true.\n2. The discussion of o.o.c condition is a bit loose. For example the theorem on page 5 (what is an informal theorem anyway) requires the condition to be as depicted in figure 3. I think assumptions in theorem should be defined more rigorously than referring to a figure.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ZSo-jWeNz-c", "reviewer": "Reviewer_WsKo", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the issue about the mismatch between $\\mathcal{D}\\_{test}$ and $\\mathcal{D}\\_{train}$ (i.e., out-of-distribution (OOD)). Different from a common OOD situation that the ratio between the density function of $\\mathcal{D}\\_{test}$ and $\\mathcal{D}\\_{train}$ is bounded, the authors focus on out-of-support where the density ratios are not bounded. To solve such a problem, the authors propose a transductive re-parameterization by introducing an anchor point $x'$ for a query point $x$. By leveraging the techniques from low-rank matrix completion, the authors provide a guarantee for OOS extrapolation using such transductive re-parameterization under certain conditions.", "summary_of_the_review": "I like this paper in general, due to the novelty of the result and the clearness of representation.", "strength_and_weaknesses": "Strength: The topic is interesting. The result seems novel and convincing. The presentation is relatively clear.\n\nWeaknesses: Some places may need more explanation and justification. I have some detailed questions as follows.\n\n1. The theorem in Section 3.2 is derived by leveraging the techniques in matrix completion. I wonder what are the key differences and similarities between this paper's result and the matrix completion result.\n\n2. If I understood correctly, the proposed transductive re-parameterization needs $(x-x')$ and $x'$, so the training is on every pair of $x'$ and $x$. Will this cause some complexity issues during training when the number of training samples is large (as the number of pairs is almost the square of the training samples)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster