Tailoring: encoding inductive biases by optimizing unsupervised objectives at prediction time
Year: 2022

Abstract:
From CNNs to attention mechanisms, encoding inductive biases into neural networks has been a fruitful source of improvement in machine learning. Adding auxiliary losses to the main objective function is a general way of encoding biases that can help networks learn better representations. However, since auxiliary losses are minimized only on training data, they suffer from the same generalization gap as regular task losses. Moreover, by adding a term to the loss function, the model optimizes a different objective than the one we care about. In this work we address both problems: first, we take inspiration from transductive learning and note that after receiving an input but before making a prediction, we can fine-tune our networks on any unsupervised loss. We call this process tailoring, because we customize the model to each input to ensure our prediction satisfies the inductive bias. Second, we formulate meta-tailoring, a nested optimization similar to that in meta-learning, and train our models to perform well on the task objective after adapting them using an unsupervised loss. The advantages of tailoring and meta-tailoring are discussed theoretically and demonstrated empirically on a diverse set of examples.

heading: 1 Introduction
text: The key to successful generalization in machine learning is the encoding of useful inductive biases. A variety of mechanisms, from parameter tying to data augmentation, have proven useful to improve the performance of models. Among these, auxiliary losses can encode a wide variety of biases, constraints, and objectives; helping networks learn better representations and generalize more broadly. Auxiliary losses add an extra term to the task loss that is minimized over the training data.
However, they have two major problems:
1. Auxiliary losses are only minimized at training time, but not for the query points. This leads to a generalization gap between training and testing, in addition to that of the task loss.
2. By minimizing the sum of the task loss plus the auxiliary loss, we are optimizing a different objective than the one we care about (only the task loss).
In this work we propose a solution to each problem:
1. We use ideas from transductive learning to minimize unsupervised auxiliary losses at each query, thus eliminating their generalization gap. Because these losses are unsupervised, we can optimize them at any time inside the prediction function. We call this process tailoring, since we customize the model to each query.
2. We use ideas from meta-learning to learn a model that performs well on the task loss after being tailored with the unsupervised auxiliary loss; i.e. meta-tailoring. This effectively trains the model to leverage the unsupervised tailoring loss in order to minimize the task loss.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Illustrative example Imagine you want to use a neural network to predict the motion of a planetary system: given the positions and velocities of each planet, the network predicts their future positions and velocities. Additionally, we could encode energy and momentum conservation by adding an auxiliary loss encouraging the neural network to conserve energy and momentum for the training examples. However, this does not guarantee that the network will conserve them for test queries. Alternatively, we can exploit that evaluating these conservations requires comparing only the input with the prediction without needing access to the true target. Therefore, we can enforce these conservations by optimizing an unsupervised objective within the prediction function. In doing so, we tailor the model to each individual query to ensure it satisfies energy and momentum conservation. Taking into account this prediction-time adaptation during training leads to a two-layer optimization, where we train to make accurate predictions after encouraging the physical conservations.
Tailoring a predictor Traditionally, supervised learning is approached within the inductive learning framework, shown in the second row of Figure 1. There, an algorithm consumes a training dataset of input-output pairs, ((xi, yi))ni=1, and produces a set of parameters θ̂ by minimizing a supervised loss �n i=1 Lsup(fθ(xi), yi) and, optionally, an unsupervised auxiliary loss �n i=1 Lunsup(θ, xi). These parameters specify a hypothesis fθ̂(·) that, given a new input x, generates an output ŷ = fθ̂(x). This problem setting misses a substantial opportunity: before the learning algorithm sees the query point x, it has distilled the data down to the parameters θ̂, which are frozen during inference, and so it cannot use new information about the particular x that it will be asked to make a prediction for.
Vapnik recognized an opportunity to make more accurate predictions when the query point is known, in a framework that is now known as transductive learning [50, 11], illustrated in the top row of Figure 1. In transductive learning, a single algorithm consumes both labeled data, ((xi, yi))ni=1, and a set of input queries for which predictions are desired, (x(j))j , and produces predictions (ŷ(j))j for each query. In general, however, we do not know queries a priori, and instead, we want an inductive function that makes predictions online, as queries arrive. To obtain such an online prediction function from a transductive system, we would need to take the training data and the single unlabeled query and encapsulate the entire transductive learning procedure inside the prediction function itself. This strategy would achieve our objective of taking x into account at prediction time but would be computationally much too slow [12].
This approach for combining induction and transduction would reuse the same training data and objective for each prediction, only changing the single unlabeled query. Consequently, it would perform extremely similar computations for each prediction. Therefore, we propose to effectively reuse the shared computations and find a “meta-hypothesis” that can then be efficiently adapted to each query. As shown in the third row of Figure 1, we propose to first run regular supervised learning to obtain parameters θ̂. Then, given a query input x, we fine-tune θ̂ on an unsupervised loss Ltailor to obtain cus-
Algorithm 1 MAMmoTh: Model-Agnostic Meta-Tailoring Subroutine Training(f , Lsup, λsup, Ltailor, λtailor, Dtrain ,b)
randomly initialize θ while not done do
Sample batch of samples (xi, yi) ∼ Dtrain forall (xi, yi) do
θxi = θ − λtailor∇θLtailor(θ, xi) // Inner step with tailor loss θ = θ − λsup∇θ � (xi,yi) Lsup � fθxi (xi), yi � // Outer step with supervised loss
return θ
tomized parameters θx and use them to make the final prediction: fθx(x). We call this process tailoring, because we adapt the model to each particular input for a customized fit. Notice that tailoring optimizes the loss at the query input, eliminating the generalization gap on the unsupervised auxiliary loss.
Meta-tailoring Since we will be applying tailoring at prediction time, it is natural to incorporate this adaptation during training, resulting in a two-layer optimization similar to those used in metalearning. Because of this similarity, we call this process meta-tailoring, illustrated in the bottom row of Figure 1. Now, rather than letting θ̂ be the direct minimizer of the supervised loss, we set it to
θ̂ ∈ argmin θ
n�
i=1
Lsup(fτ(θ,Ltailor,xi)(xi), yi).
Here, the inner loop optimizes the unsupervised tailoring loss Ltailor and the outer loop optimizes the supervised task loss Lsup. Notice that now the outer process optimizes the only objective we care, Lsup, instead of a proxy combination of Lsup and Lunsup. At the same time, we learn to leverage Ltailor in the inner loop to affect the model before making the final prediction, both during training and evaluation. Adaptation is especially clear in the case of a single gradient step, as in MAML [19]. We show its translation, MAMmoTh (Model-Agnostic Meta-Tailoring), in algorithm 1.
In many settings, we want to make predictions for a large number of queries in a (mini-)batch. While MAMmoTh adapts to every input separately, it can only be run efficiently in parallel in some deep learning frameworks, such as JAX [10]. Inspired by conditional normalization (CN) [18] we propose CNGRAD, which adds element-wise affine transformations to our model and only adapts the added parameters in the inner loop. This allows us to independently tailor the model for multiple inputs in parallel. We prove theoretically, in Sec. 4, and provide experimental evidence, in Sec. 5.1, that optimizing these parameters alone has enough capacity to minimize a large class of tailoring losses.
Relation between (meta-)tailoring, fine-tuning transfer, and meta-learning Fine-tuning pretrained networks is a fruitful method of transferring knowledge from large corpora to smaller related datasets [17]. This allows us to reuse features on related tasks or for different distributions of the same task. When the data we want to adapt to is unlabeled, we must use unsupervised losses. This can be useful to adapt to changes of task [16], from simulated to real data [52], or to new distributions [46].
Tailoring performs unsupervised fine-tuning and is, in this sense, similar to test-time training(TTT) [46] for a single sample, which adapts to distribution shifts. However, tailoring is applied to a single query; not to a data set that captures distribution shift, where batched TTT sees most of its benefits. Thus, whereas regular fine-tuning benefits from more adaptation data, tailoring would be hindered by adapting simultaneously to more data. This is because tailoring aims at building a custom model for each query to ensure the network satisfies a particular inductive bias. Customizing the model to multiple samples makes it harder, not easier. We show this in Figure 2, where TTT with 6400 samples performs worse than tailoring with a single sample. Furthermore, tailoring adapts to each query one by one, not globally from training data to test data. Therefore, it also makes sense to do tailoring on training queries (i.e., meta-tailoring).
Meta-tailoring has the same two-layer optimization structure as meta-learning. More concretely, it can be understood as the extreme case of meta-learning where each single-query prediction is its own task. However, whereas meta-learning tasks use one loss and different examples for the inner and outer loop, meta-tailoring tasks use one example and different losses for each loop (Ltailor,Lsup). We emphasize that meta-tailoring does not operate in the typical multi-task meta-learning setting. Instead, we are leveraging techniques from meta-learning for the classical single-task setting.
Contributions In summary, our contributions are: 1. Introducing tailoring, a new framework for encoding inductive biases by minimizing unsuper-
vised losses at prediction time, with theoretical guarantees and broad potential applications.
2. Formulating meta-tailoring, which adjusts the outer objective to optimize only the task loss, and developing a new algorithm, CNGRAD, for efficient meta-tailoring.
3. Demonstrating meta-tailoring in 3 domains: encoding hard and soft conservation laws in physics prediction problems (Sec. 5.1 and Sec. 5.2), enhancing resistance to adversarial examples by increasing local smoothness at prediction time (Sec. 5.4), and improving prediction quality both theoretically (Sec. 3.1) and empirically (Sec. 5.3) by tailoring with a contrastive loss.

heading: 2 Related work
text: Tailoring is inspired by transductive learning. However, transductive methods, because they operate on a batch of unlabeled queries, are allowed to make use of the underlying distributional properties of those queries, as in semi-supervised learning [12]. In contrast, tailoring does the bulk of the computations before receiving any query; vastly increasing efficiency. Similar to tailoring, local learning [9] also has input-dependent parameters. However, it uses similarity in raw input space to select a few labeled data points and builds a local model instead of reusing the global prior learned across the whole data. Finally, some methods [21, 33] in meta-learning propagate predictions along the test samples in a semi-supervised transductive fashion.
Similar to tailoring, there are other learning frameworks that perform optimization at prediction time for very different purposes. Among those, energy-based models do generative modeling [2, 27, 32] by optimizing the hidden activations of neural networks, and other models [4, 49] learn to solve optimization problems by embedding optimization layers in neural networks. In contrast, tailoring optimizes the parameters of the model, not the hidden activations or the output.
As discussed in the introduction, unsupervised fine-tuning methods have been proposed to adapt to different types of variations between training and testing. Sun et al. [46] propose to adapt to a change of distribution with few samples by unsupervised fine-tuning at test-time, applying it with a loss of predicting whether the input has been rotated. Zhang et al. [54] build on it to adapt to group distribution shifts with a learned loss. Other methods in the few-shot meta-learning setting exploit test samples of a new task by minimizing either entropy [16] or a learned loss [5] in the inner optimization. Finally, Wang et al. [51] use entropy in the inner optimization to adapt to large-scale variations in image segmentation. In contrast, we propose (meta-)tailoring as a general effective way to impose inductive biases in the classic machine learning setting. Whereas in the aforementioned methods, adaptation happens from training to testing, we independently adapt to every single query.
Meta-learning [44, 7, 48, 28] has the same two-level optimization structure as meta-tailoring but focuses on multiple prediction tasks. As shown in Alg. 1 for MAML [19], most optimization-based meta-learning algorithms can be converted to meta-tailoring. Similar to CNGRAD, there are other meta-learning methods whose adaptations can be batched [40, 3]. Among these, [55, 41] train FiLM networks [39] to predict custom conditional normalization (CN) layers for each task. By optimizing the CN layers directly, CNGRAD is simpler, while remaining provably expressive (section 4). CNGrad can also start from a trained model by initializing the CN layers to the identity function.

heading: 3 Theoretical motivations of meta-tailoring
text: In this section, we study the potential advantages of meta-tailoring from the theoretical viewpoint, formalizing the intuitions conveyed in the introduction. By acting symmetrically during training and prediction time, meta-tailoring allows us to closely relate its training and expected losses, whereas tailoring alone does not have the same guarantees. First, we analyze the particular case of a contrastive tailoring loss. Then, we will generalize the guarantees to other types of tailoring losses.

heading: 3.1 Meta-tailoring with a contrastive tailoring loss
text: Contrastive learning [24] has seen significant successes in problems of semi-supervised learning [37, 26, 13]. The main idea is to create multiple versions of each training image and learn a representation in which variations of the same image are close while variations of different images are far apart. Typical augmentations involve cropping, color distortions, and rotation. We show theoretically that, under reasonable conditions, meta-tailoring using a particular contrastive loss Lcont as Ltailor = Lcont helps us improve generalization errors in expectation compared with performing classical inductive learning.
When using meta-tailoring, we define θx,S to be the θx obtained with a training dataset S = ((xi, yi)) n i=1 and tailored with the contrastive loss at the prediction point x. Theorem 1 provides an upper bound on the expected supervised loss Ex,y[Lsup(fθx,S (x), y)] in terms of the expected contrastive loss Ex[Lcont(x, θx,S)] (analyzed in App. B), the empirical supervised loss 1 n �n i=1 Lsup(fθxi,S (xi), yi) of meta-tailoring, and its uniform stability ζ. Theorem 6 (App. C) provides a similar bound with the Rademacher complexity [6] Rn(Lsup ◦ F) of the set Lsup ◦ F , instead of using the uniform stability ζ. Proofs of all results in this paper are deferred to App. C.
Definition 1. Let S = ((xi, yi))ni=1 and S� = ((x�i, y�i))ni=1 be any two training datasets that differ by a single point. Then, a meta-tailoring algorithm S �→ fθx,S (x) is uniformly ζ-stable if ∀(x, y) ∈ X × Y, |Lsup(fθx,S (x), y)− Lsup(fθx,S� (x), y)| ≤ ζ n .
Theorem 1. Let S �→ fθx,S (x) be a uniformly ζ-stable meta-tailoring algorithm. Then, for any δ > 0, with probability at least 1 − δ over an i.i.d. draw of n i.i.d. samples S = ((xi, yi))ni=1, the following holds: for any κ ∈ [0, 1], Ex,y[Lsup(fθx,S (x), y)] ≤ κEx � Lcont(x, θx,S) � + (1− κ)J , where J = 1n �n i=1 Lsup(fθxi,S (xi), yi) + ζ n + (2ζ + c) � (ln(1/δ))/(2n), and c is the upper bound on the per-sample loss as Lsup(fθ(x), y) ≤ c. In the case of regular inductive learning, we get a bound of the exact same form, except that we have a single θ instead of a θx tailored to each input x. This theorem illustrates the effect of meta-tailoring on contrastive learning, with its potential reduction of the expected contrastive loss Ex[Lcont(x, θx,S)]. In classic induction, we may aim to minimize the empirical contrastive loss 1n̄ �n̄ i=1 Lcont(xi, θ) with n̄ potentially unlabeled training samples, which incurs the additional
generalization error of Ex[Lcont(x, θx,S)]− 1n̄ �n̄
i=1 Lcont(xi, θ). In contrast, meta-tailoring can avoid this extra generalization error by directly minimizing a custom θx on each x: Ex[Lcont(x, θx,S)]. In the case where Ex[Lcont(x, θx,S)] is left large (e.g., due to large computational cost), Theorem 1 still illustrates competitive generalization bounds of meta-tailoring with small κ. For example, with κ = 0, it provides generalization bounds with the uniform stability for meta-tailoring algorithms. Even then, the bounds are not equivalent to those of classic induction, and there are potential benefits of meta-tailoring, which are discussed in the following section with a more general setting.

heading: 3.2 Meta-tailoring with general tailoring losses
text: The benefits of meta-tailoring go beyond contrastive learning: below we provide guarantees for meta-tailoring with arbitrary pairs of tailoring loss Ltailor(x, θ) and supervised loss Lsup(fθ(x), y). Remark 1. For any function ϕ such that Ex,y[Lsup(fθ(x), y)] ≤ Ex[ϕ(Ltailor(x, θ))], Theorems 1 and 6 hold with the map Lcont being replaced by the function ϕ ◦ Ltailor. This remark shows the benefits of meta-tailoring through its effects on three factors: the expected unlabeled loss Ex[ϕ(Ltailor(x, θx,S))], uniform stability ζ , and the Rademacher complexity Rn(Lsup ◦ F). It is important to note that meta-tailoring can directly minimize the expected unlabeled loss Ex[ϕ(Ltailor(x, θx,S))], whereas classic induction can only minimize its empirical version, which results in the additional generalization error on the difference between the expected unlabeled loss and its empirical version. For example, if ϕ is monotonically increasing and Ltailor(x, θ) represents the physical constraints at each input x (as in the application in section 5.1), then classic induction requires a neural network trained to conserve energy at the training points to generalize to also conserve it at unseen (e.g., testing) points. Meta-tailoring avoids this requirement by directly minimizing violations of energy conservation at each point at prediction time.
Meta-tailoring can also improve the parameter stability ζθ defined such that ∀(x, y) ∈ X×Y, �θx,S− θx,S�� ≤ ζθn , for all S, S� differing by a single point. When θx,S = θ̂S − λ∇Ltailor(x, θ̂S), we obtain an improvement on the parameter stability ζθ if ∇Ltailor(x, θ̂S) can pull θ̂S and θ̂S� closer so that �θx,S − θx,S�� < �θ̂S − θ̂S��, which is ensured, for example, if � · � = � · �2 and cos_dist(v1, v2) �v1� �v2� > 1 2 where cos_dist(v1, v2) is the cosine similarity of v1 and v2, with v1 = θ̂S − θ̂S� , v2 = λ(∇Ltailor(x, θ̂S) −∇Ltailor(x, θ̂S�)) and v2 �= 0. Here, the uniform stability ζ and the parameter stability ζθ are closely related as ζ ≤ Cζθ, where C is the upper bound on the Lipschitz constants of the maps θ �→ Lsup(fθ(x), y) over all (x, y) ∈ X × Y under the norm � · �, since |Lsup(fθx,S (x), y)− Lsup(fθx,S� (x), y)| ≤ C�θx,S − θx,S�� ≤ Cζθ n .
Algorithm 2 CNGRAD for meta-tailoring Subroutine Training(f , Lsup, λsup, Ltailor, λtailor, steps,Dtrain ,b) // Only in meta-tailoring
randomly initialize w // All parameters except γ,β; trained in outer loop while not done do
X,Y ∼b Dtrain ; gradw = 0 // Sample batch; initialize outer gradient γ0 = 1b, � l ml ;β0 = 0b, � l ml // Initialize CN layers to the identity for 1 ≤ s ≤ steps do γs = γs−1 − λtailor∇γLtailor(w, γs−1,βs−1, X) // Inner step w.r.t. γ βs = βs−1 − λtailor∇βLtailor(w, γs−1,βs−1, X) // Inner step w.r.t. β γs,βs = γs.detach(), βs.detach() // Only in 1st order CNGrad gradw = gradw +∇wLsup (fw,γs,βs(X), Y ) // Outer gradient w.r.t. w
w = w − λsupgradw // Apply outer step after all inner steps return w Subroutine Prediction(f , w, Ltailor, λ, steps, X) // Both in meta-tailoring & tailoring γ0 = 1X.shape[0], � l ml ;β0 = 0X.shape[0], � l ml
for 1 ≤ s ≤ steps do γs = γs−1 − λ∇γLtailor(w, γs−1,βs−1, X) βs = βs−1 − λ∇βLtailor(w, γs−1,βs−1, X) return fw,γsteps,βsteps(X)
4 CNGRAD: a simple algorithm for expressive, efficient (meta-)tailoring In this section, we address the issue of using (meta-)tailoring for efficient GPU computations. Although possible in JAX [10], efficiently parallelizing MAMmoTh across inputs is not possible in other frameworks. To overcome this issue, building on CAVIA [55] and WarpGrad [20], we propose CNGRAD which adapts only conditional normalization parameters and enables efficient GPU computations for (meta-)tailoring. CNGRAD can also be used in meta-learning, providing a parallelizable alternative to MAML (see App. D).
As done in batch-norm [30] after element-wise normalization, we can implement an element-wise affine transformation with parameters (γ,β), scaling and shifting the output h(l)k (x) of each k-th neuron at the l-th hidden layer independently: γ(l)k h (l) k (x)+β (l) k . In conditional normalization, Dumoulin et al. [18] train a collection of (γ,β) in a multi-task fashion to learn different tasks with a single network. CNGRAD brings this concept to the meta-learning and (meta-)tailoring settings and adapts the affine parameters (γ,β) to each query. For meta-tailoring, the inner loop minimizes the tailoring loss at an input x by adjusting the affine parameters and the outer optimization adapts the rest of the network. Similar to MAML [19], we implement a first-order version, which does not backpropagate through the optimization, and a second-order version, which does. CNGRAD efficiently parallelizes computations of multiple tailored models because the adapted parameters only require element-wise multiplications and additions. See Alg. 2 for the pseudo-code.
CNGRAD is widely applicable since the adaptable affine parameters can be added to any hidden layer and only represent a tiny portion of the network (empirically, around 1%). Moreover, we can see that, under realistic assumptions, we can minimize the inner tailoring loss using only the affine parameters. To analyze properties of these adaptable affine parameters, let us decompose θ into θ = (w, γ,β), where w contains all the weight parameters (including bias terms), and the (γ,β) contains all the affine parameters. Given an arbitrary function (fθ(x), x) �→ �tailor(fθ(x), x), let Ltailor(x, θ) = �ngi=1 �tailor(fθ(g(i)(x)), x), where g(1):(ng) are arbitrary input augmentation functions at prediction time.
Corollary 1 states that for any given ŵ, if we add any non-degenerate Gaussian noise δ as ŵ + δ with zero mean and any variance on δ, the global minimum value of Ltailor w.r.t. all parameters (w, γ,β) can be achieved by optimizing only the affine parameters (γ,β), with probability one. In other words, the CN parameters (γ,β) have enough capacity to optimize optimize the inner tailoring loss.
Corollary 1. Under the assumptions of Theorem 2, for any ŵ ∈ Rd, with probability one over randomly sampled δ ∈ Rd accordingly to any non-degenerate Gaussian distribution, the following holds: infw,γ,β Ltailor(x,w, γ,β) = infγ,β Ltailor(x, ŵ + δ, γ,β) for any x ∈ X . The assumption and condition in theorem 2 are satisfied in practice (see App. A). Therefore, CNGRAD is a practical and computationally efficient method to implement (meta-)tailoring.

heading: 5 Experiments

heading: 5.1 Tailoring to impose symmetries and constraints at prediction time
text: Exploiting invariances and symmetries is an established strategy for increasing performance in ML. During training, we can regularize networks to satisfy specific criteria; but this does not guarantee they will be satisfied outside the training dataset [45]. (Meta-)tailoring provides a general solution to this problem by adapting the model to satisfy the criteria at prediction time. We demonstrate the use of tailoring to enforce physical conservation laws for predicting the evolution of a 5-body planetary system. This prediction problem is challenging, as m-body systems become chaotic for m > 2. We generate a dataset with positions, velocities, and masses of all 5 bodies as inputs and the changes in position and velocity as targets. App. E further describes the dataset.
Our model is a 3-layer feed-forward network. We tailor it by taking the original predictions and adapting the model using the tailoring loss given by the L1 loss between the whole system’s initial and final energy and momentum. Note that ensuring this conservation does not guarantee better performance: predicting the input as the output conserves energy and momentum perfectly, but it is not correct.
While tailoring adapts some parameters in the network to improve the tailoring loss, an alternative for enforcing conservation would be to adapt the output y value directly. Table 1 compares the predictive accuracy of inductive learning, direct output optimization, and both tailoring and meta-tailoring, using varying numbers of gradient steps. Tailoring is more effective than adapting the output, as the parameters provide a prior on what changes are more natural. For meta-tailoring, we try both first-order and second-order versions of CNGRAD. The first-order gave slightly better results, possibly because it was trained with a higher tailor learning rate (10−3) with which the second-order version was unstable (we thus used 10−4). More details can be found in App. E.
Finally, meta-tailoring without any query-time tailoring steps already performs much better than the original model, even though both have almost the same number of parameters and can overfit the dataset. We conjecture meta-tailoring training adds an inductive bias that guides optimization towards learning a more generalizable model. Fig. 2 shows prediction-time optimization paths.

heading: 5.2 Tailoring to softly encourage inductive biases
text: A popular way of encoding inductive biases is with clever network design to make predictions translation equivariant (CNNs), permutation equivariant (GNNs), or conserve energy [23]. However, if an inductive bias is only partially satisfied, such approaches overly constrain the function class. Instead, tailoring can softly impose this bias by only fine-tuning the tailoring loss for a few steps.
We showcase this in the real pendulum experiment used by Hamiltonian Neural Networks (HNNs) [23]. HNNs have energy conservation built-in and easily improve a vanilla MLP. We meta-tailor this vanilla MLP with energy conservation without changing its architecture. Meta-tailoring significantly improves over the baseline and HNNs, since it can encode the imperfect energy conservation of real systems. We compare results in Fig. 3 and provide extra details in App. F. Note that, with inexact losses, fully enforcing them provides
sub-optimal results. Thus, we pick the tailoring learning rate that results in the lowest long-term prediction loss during training.
5.3 Tailoring with a contrastive loss for image classification
Following the setting described in section 3.2, we provide experiments on the CIFAR-10 dataset [31] by building on SimCLR [13]. SimCLR trains a ResNet-50 [25] fθ(·) coupled to a small MLP g(·) such that the outputs of two augmentations of the same image xi, xj ∼ T (x) agree; i.e. g(fθ(xi)) ≈ g(fθ(xj)). This is done by training g(f(·)) to recognize one augmentation from the other among a big batch of candidates with the cross-entropy loss. To show that the unsupervised training of fθ provides a useful representation, SimCLR trains a single linear layer on top of it, φ(fθ(·)), achieving good classification results. We now observe that we can tailor fθ at prediction-time by optimizing g(fθx(x)), which quantifies the agreement between different augmentations of the same input; thus ’learning’ about its particularities. To make the image classification prediction, we feed the final tailored representation to the linear layer: φ(fθx(x)). To match the evaluation from SimCLR, we do not redo SimCLR’s un-
supervised learning, which provides θ. The meta-tailoring outer loop trains φ to take the tailored representations fθx(x) instead of the original fθ(x). Thus, θ is unsupervisedly fine-tuned in the prediction function leading to θx, but never supervisedly trained as this would break the evaluation protocol (in meta-tailoring’s favor). We also implement a TTT [46] baseline with their original rotation-prediction loss. Moreover, TTT modifies θx at test time, but does not take this adaptation into account when training φ (see App. G for more details). TTT worsened base SimCLR despite significant hyper-parameter tuning. We conjecture this is because TTT was designed for OOD generalization, not in-distribution. In contrast, as shown in Fig. 4, we observe that meta-tailoring provides improvements over base SimCLR equivalent to doubling the amount of labeled data.

heading: 5.4 Tailoring for robustness against adversarial examples
text: Neural networks are susceptible to adversarial examples [8, 47]: targeted small perturbations of an input can cause the network to misclassify it. One approach is to make the prediction function smooth via adversarial training [34]; however, this only ensures smoothness in the training points. Constraining the model to be smooth everywhere makes it lose capacity. Instead, (meta-)tailoring asks for smoothness a posteriori, only on a specific query.
We apply meta-tailoring to robustly classifying CIFAR-10 [31] and ImageNet [15] images, tailoring predictions so that they are locally smooth. This is similar to VAT [36] but instead optimizes the loss within the prediction function, not as an auxiliary loss. Inspired by the notion of adversarial examples being caused by predictive, but non-robust, features [29], we meta-tailor our model by enforcing smoothness on the vector of features of the penultimate layer (denoted gθ(x)):
Ltailor(x, θ) = E[cos_dist(gθ(x), gθ(x+ δ))], δ ∼ N(0, ν2),
We build on Cohen et al. [14], who developed a method for certifying the robustness of a model via randomized smoothing (RS). RS samples points from a Gaussian N(x,σ2) around the query and, if there is enough agreement in classification, it provides a certificate that a small perturbation cannot adversarially modify the query to have a different class. We show that meta-tailoring improves the original RS method, testing for σ = 0.25, 0.5, 1.0. We use ν = 0.1 for all experiments. We initialized with the weights of Cohen et al. [14] by leveraging that CNGRAD can start from a pre-trained model by initializing the extra affine layers to the identity. Finally, we use σ� = √ σ2 − ν2 ≈ 0.23, 0.49, 0.995 so that the points used in our tailoring loss come from N(x,σ2).
Table 7 shows our results on CIFAR-10 where we improve the average certification radius (ARC) by 8.6%, 10.4%, 19.2% respectively. In table 2, we show results on Imagenet where we improve the ARC by 5.1%, 13.8%, 19.6% respectively. We chose to meta-tailor the RS method because it represents a strong standard in certified adversarial defenses, but we note that there have been advances on RS that sometimes achieve better results than those presented here [53, 43], see App. I. However, it is likely that meta-tailoring could also improve these methods.
These experiments only scratch the surface of what tailoring allows for adversarial defenses: usually, the adversary looks at the model and gets to pick a particularly bad perturbation x+ δ. With tailoring, the model responds, by changing to weights θx+δ. This leads to a game, where both weights and inputs are perturbed, similar to max|δ|<�x min|Δ|<�θ Lsup (fθ+Δ(x+ δ), y). However, since we don’t get to observe y; we optimize the weight perturbation by minimizing Ltailor instead.

heading: 6 Discussion

heading: 6.1 Broader Impact
text: Improving adversarial robustness: having more robust and secure ML systems is mostly a positive change. However, improving adversarial defenses could also go against privacy preservation, like the use of adversarial patches to gain anonymity from facial recognition. Encoding desirable properties: By optimizing an unsupervised loss for the particular query we care about, it is easier to have guarantees on the prediction. In particular, there could be potential applications for fairness, where the unsupervised objective could enforce specific criteria at the query or related inputs. More research needs to be done to make this assertion formal and practical. Potential effect on privacy: tailoring specializes the model to each input. This could have an impact on privacy. Intuitively, the untailored model can be less specialized to each input, lowering the individual information from each training point contained in the model. However, tailored predictions extract more information about the queries, from which more personal information could be leaked.

heading: 6.2 Limitations
text: Tailoring provides a framework for encoding a wide array of inductive biases, but these need to be specified as a formula by the user. For instance, it would be hard to programatically describe tailoring losses in raw pixel data, such as mass conservation in pixel space. Tailoring also incurs an extra time cost at prediction time, since we make an inner optimization inside the prediction function. However, as shown in Table 1, meta-tailoring often achieves better results than inductive learning even without adaptation at test-time, enabling better predictions at regular speed during test-time. This is due to meta-tailoring leading to better training. Moreover, optimization can be sped up by only tailoring the last layers, as discussed in App. D. Finally, to the best of our knowledge using MAMmoTh for meta-tailoring would be hard to parallelize in PyTorch [38] and Tensorflow [1]; we
proposed CNGRAD to make it easy and efficient. JAX[10], which handles per-example weights, makes parallelizing tailoring effortless.
Theory in Sec. 3 applies only to meta-tailoring. Unlike tailoring (and test-time training), metatailoring performs the same computations at training and testing time, which allows us to prove the results. Theorem 2 proves that optimizing the CN layers in CNGRAD has the same expressive power as optimizing all the layers for the inner (not outer) loss. However, it does not guarantee that gradient descent will find the appropriate optima. The study of such guarantee is left for future work.

heading: 6.3 Conclusion
text: We have presented tailoring, a simple way of embedding a powerful class of inductive biases into models, by minimizing unsupervised objectives at prediction time. Tailoring leverages the generality of auxiliary losses and improves them in two ways: first, it eliminates the generalization gap on the auxiliary loss by optimizing it on the query point; second, tailoring only minimizes task loss in the outer optimization and the tailoring loss in the inner optimization. This results in the model optimizing the only objective we care about in the outer loop, instead of a proxy loss. Beyond inductive biases, tailoring shows that model adaptation is useful even when test queries comes from the same distribution as the training data. This suggests one can improve models by performing prediction-time optimization, trading off large offline data&compute efforts with small online computations.
Tailoring is broadly applicable, as one can vary the model, the unsupervised loss, and the task loss. We show its applicability in three diverse domains: physics prediction time-series, contrastive learning, and adversarial robustness. We also provide a simple algorithm, CNGRAD, to make meta-tailoring practical with little additional code. Currently, most unsupervised or self-supervised objectives are optimized in task-agnostic ways; without taking into account the supervised downstream task. Instead, meta-tailoring provides a generic way to make these objectives especially useful for each application. It does so by learning how to best leverage the unsupervised loss to perform well on the final task we care about.
Acknowledgments and Disclosure of Funding
We would like to thank Kelsey Allen, Marc de la Barrera, Jeremy Cohen, Dylan Doblar, Chelsea Finn, Sebastian Flennerhag, Jiayuan Mao, Josh Tenenbaum, and Shengtong Zhang for insightful discussions. We would also like to thank Clement Gehring for his help with deploying the experiments and Lauren Milechin for her help with leveraging the MIT supercloud platform [42].
We gratefully acknowledge support from NSF grant 1723381; from AFOSR grant FA9550-17-1-0165; from ONR grant N00014-18-1-2847; from the Honda Research Institute, from MIT-IBM Watson Lab; and from SUTD Temasek Laboratories. We also acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPC resources that have contributed to the reported research results. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of our sponsors.

Review summaries:
The paper proposes a new method to adapt neural networks at inference time to the given input such that the model minimizes some given unsupervised loss. This method is called 'tailoring'. Additionally, 'meta-tailoring' is proposed that trains the model using 'tailoring', too. This way the gap between training and inference distribution is removed that 'tailoring' introduces.

This paper proposed tailoring, which is a general framework that could help to finetune the prediction on each test sample according to some specific inductive biases. Tailoring provides a different perspective that avoids involving extra loss function in the proxy fashion.  Besides, the authors proposed meta-tailoring, which integrates the unsupervised loss in a way similar to meta-learning. The theoretical discussion and empirical results demonstrate the effectiveness of the proposed tailoring. 

In this paper, the authors propose optimizing an unsupervised loss function at test time as an inductive prior on a neural network. They propose two schemes for achieving this: training a network in a regular fashion and then applying the unsupervised loss at inference time only, or applying the unsupervised loss during training as a meta-learning scheme. The authors claim that this method can improve robustness, as well as improve train-test generalization gap in cases where there is a known inductive prior on the expected output (e.g. when modelling physical systems)