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

1 Introduction . 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 can not 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 . 2 Related work . 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 .

Summary: 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)