Functional Indirection Neural Estimator for Better Out-of-distribution Generalization
Abstract
The capacity to achieve out-of-distribution (OOD) generalization is a hallmark of human intelligence and yet remains out of reach for machines.
This remarkable capability has been attributed to our abilities to make conceptual abstraction and analogy, and to a mechanism known as indirection, which binds two representations and uses one representation to refer to the other.
Inspired by these mechanisms, we hypothesize that OOD generalization may be achieved by performing analogymaking and indirection in the functional space instead of the data space as in current methods.
To realize this, we design FINE
(Functional Indirection Neural
Estimator), a neural framework that learns to compose functions that map data input to output on-the-ﬂy.
FINE consists of a backbone network and a trainable semantic memory of basis weight matrices.
Upon seeing a new input-output data pair, FINE dynamically constructs the backbone weights by mixing the basis weights.
The mixing coefﬁcients are indirectly computed through querying a separate corresponding semantic memory using the data pair.
We demonstrate empirically that FINE can strongly improve out-of-distribution generalization on
IQ tasks that involve geometric transformations.
In particular, we train FINE and competing models on IQ tasks using images from the MNIST, Omniglot and CIFAR100 datasets and test on tasks with unseen image classes from one or different datasets and unseen transformation rules.
FINE not only achieves the best performance on all tasks but also is able to adapt to small-scale data scenarios.
Introduction
Every computer science problem can be solved with a higher level of indirection.
—Andrew Koenig, Butler Lampson, David J. Wheeler
Generalizing to new circumstances is a hallmark of intelligence [16, 4, 11].
In some Intelligence
Quotient (IQ) tests–a popular benchmark for human intelligence–one must leverage their prior experience to identify the hidden abstract rules out of a concrete example (e.g., a transformation of an image) and then apply the rules to the next (e.g., a new set of images of totally different appearance).
These tasks necessitate several key capabilities, including conceptual abstraction and analogy-making [22].
Abstraction allows us to extend a concept to novel situations.
It is also driven by analogy-making, which maps the current situation to the previous experience stored in the memory.
Indeed, analogy-making has been argued to be one of the most important abilities of human cognition, or even further, “a concept is a package of analogies” [12]. The ability for humans to traverse seamlessly across concrete and abstraction levels suggests another mechanism known as indirection to bind two representations and use one representation to refer to the other [16, 20]. x ...
... backbone layerbackbone layer t y f un c t i o n c o m p
o s i
t i o n i nd i r e c t i o n
W t ..
y t x +t =
W qt ..
M value t,
M value t, M value t,s
M keyt,
M keyt,
M key t,s a n a l o g y - m a k
i ng
MEMORY encoder encoder
The function composer geometric transformation deep learning models have successfully utilized analogy and indirection.
The Transformer [30] and RelationNet [28] learn analogies between data, through self-attention or pairwise functions.
The ESBN [33] goes further by incorporating the indirection mechanism to bind an entity to a symbol and then reason on the symbols; and this has proved to be efﬁcient on tasks involving abstract rules, similar to those aforementioned IQ tests.
However, a common drawback of these approaches is that they operate on the data space, and thus are susceptible to out-of-distribution samples.
In this paper, we propose to perform analogy-making and indirection in functional spaces instead.
We aim to learn to compose a functional mapping from a given input to an output on-the-ﬂy .
By doing so, we achieve two clear advantages.
First, since the class of possible mappings is often restricted, it may not require a large amount of training data to learn the distribution of functions.
Second, more importantly, since this approach performs indirection in functional spaces, it avoids bindings between numerous entities and symbols in data spaces, thus may help improve the out-of-distribution generalization capability.
To this end, we introduce a new class of problems that requires functional analogy-making and indirection, which are deemed to be challenging for current data-oriented approaches.
The tasks are similar to popular IQ tasks in which the model is given hints about the hidden rules, then it has to predict the missing answer following the rules.
One reasonable approach is that models should be able to compare the current task to what they saw previously to identify the rules between appearing entities, and thus has to search on functional spaces instead of data spaces.
More concretely, we construct the IQ tasks by applying geometric transformations to images from MNIST dataset, handwritten Omniglot dataset and real-image CIFAR100 dataset, where the training set and test set contain disjoint image classes from the same or different datasets, and possibly disjoint transformation rules .
Second, we present a novel framework named Functional Indirection Neural Estimator (FINE) to solve this class of problems (see Fig.
1 for the overall architecture of FINE ).
FINE consists of (a) a neural backbone to approximate the functions and (b) a trainable key-value memory module that stores the basis of the network weights that spans the space of possible functions deﬁned by the backbone.
The weight basis memories allow FINE to perform analogy-making and indirection in the function space.
More concretely, when a new IQ task arrives, FINE ﬁrst (1) takes the hint images to make analogies with value memories, then (2) performs indirection to bind value memories with key memories and ﬁnally (3) computes the approximated functions based on key memories. Throughout
A
B
C
D
Rotation ?
A
B
C
D a comprehensive suite of experiments, we demonstrate that FINE outperforms the competing methods and adapts more effectively in small data scenarios.
Tasks
For concreteness, we will focus on Intelligence Quotient (IQ) tests, which have been widely accepted as one of the reliable benchmarks to measure human intelligence [26].
We will study a popular class of IQ tasks that provides hints following some hidden rules and requires the player to choose among given choices to ﬁll in a placeholder so that the ﬁlled-in entity obeys the rules of the current task.
In order to succeed in these tasks, the player must be able to ﬁgure out the hidden rules and perform analogy-making to select the correct choice.
Moreover, once ﬁguring out the rules for the current task, a human player can almost always solve tasks with similar rules regardless of the appearing entities given in the tasks.
This remarkable ability of out-of-distribution generalization indicates that humans treat objects and relations abstractly instead of relying on the raw sensory data.
We aim to solve IQ tasks that involve geometric transformations (e.g., see Fig.
2 for an example), which include afﬁne transformations (translation, rotation, shear, scale, and reﬂection), non-linear transformations (ﬁsheye, horizontal wave) and syntactic transformations (black-white, swap).
Details of transformations are given in Supplementary.
In a task, the models are given 3 images x, y and x (cid:48) , where y is the result of x after applying a geometric transformation.
The models are then asked to select y (cid:48) among 4 choices y 1 , y 2 , y 3 , y 4 so that ( x (cid:48) , y (cid:48) ) follows the same rule as ( x, y )
(i.e. if y = f ( x ) then y (cid:48) =
f ( x (cid:48) ) for transformation f ).
The 4 choices include (i) one with correct object/image and correct transformation (which is the solution), (ii) one with correct object/image and incorrect transformation, (iii) one with incorrect object/image and correct transformation, and (iv) one with both incorrect object/image and transformation.
Inspired by human capability, a reasonable approach to solve these tasks is that models should be able to ﬁgure out the transformation (or relation) between objects/images and apply the transformation to novel objects/images.
The datasets can be classiﬁed into two main categories: single-transformation datasets and multi-transformation datasets.
Single-transformation datasets are ones that only include a particular transformation, e.g. rotation.
Note that the individual transformations of the same type vary, e.g., rotations by different angles.
Multi-transformation datasets, on the other hand, consist of several transformation types.
To test the generalization capability of the models, we build testing sets including classes of images that have never been seen during training (see Section 4.1 and Section 4.2), or even more challenging tasks including unseen rules and unseen datasets (see Section 4.3).
Models must be able to leverage knowledge and memory gained from the training dataset to solve a new task.
Method 3. Functional Hypothesis
Let X
and Y be the data input and output spaces, respectively.
Denote by ( X train , Y train ) ⊂
( X , Y ) the training set, and ( X test , Y test ) ⊂
( X , Y ) the non-overlapping test set.
Classical ML assumes that
X train and X test are drawn from the same distribution.
Under this hypothesis, it is reasonable (for frequentists) to ﬁnd a function f : X → Y in the functional space F that ﬁts both the train and the test sets (i.e., the discrepancy between f ( x ) and corresponding y is small for all ( x, y ) in ( X train , Y train ) and ( X test , Y test ) ).
However, when X test is drawn from a different data distribution from that of X train , it has been widely reported that current deep learning models fail drastically [2, 8, 13, 32, 33].
This is because the models are inferred exclusively from the observed data distribution.
Moreover, it could be the case that relations between x and y in testing samples are unseen during training, which raises questions on the feasibility of learning a single function when dealing with out-of-distribution tasks.
A natural solution for this problem would be to train the model to learn the functions adaptively.
Formally speaking, the model will learn a function composer φ : X × Y → F that maps each pair of ( x, y ) ∈
X × Y (where y is the associated output of x ) to a function φ x,y in F so that φ x,y approximates the true relation between x and y .
As discussed, training models this way leads to two clear advantages: (1) it can help to handle the cases when there are multiple (and possibly disjoint) relations between inputs and outputs within the training and testing datasets; and (2) models are less dependent on data and thus can achieve more stable results on different training and testing sets.
Empirical evidence for these points will be given in Section 4.
Since there are an inﬁnite number of functions that can map an input to an output with any given degree of precision, the key is to design φ so that it can map a given input-output pair to a “good enough” function.
We humans may draw analogies between the current situation and our experiences and then work out the most suitable options [12].
For example, we can recognize a math problem in exam to be similar to a previous exercise in class with different variable names.
We may ﬁnd the presented idea in a paper related to that in another paper we read before, as someone even says “new ideas are just re-distribution of old ideas”.
All these examples illustrate that analogy-making is a powerful strategy in human thinking process.
Inspired by this mechanism, we equip the function composer φ with a semantic memory to store past knowledge, on which analogy-making is performed.
The memory also plays the role of constraining the searching region for φ , so that φ only looks for functions in the subspaces spanned by the memories instead of the whole functional space F .
Further analysis is given in Supplementary.
A remaining question is how to design the memory.
Let us be inspired again by human thinking process: When we see an animal, we compare its face, legs or tail with things we know and ﬁnally conclude it’s a “dog”.
Here “dog” is an abstract concept bound with primary characteristics (e.g. face, legs, tail, etc.)
; once a new entity arrives, we compare its properties with these characteristics, and if they are similar, we utilize the indirection mechanism to infer it is indeed the concept we are considering.
In our case, the concepts are functions, or more speciﬁcally, the geometric transformations.
We thus maintain a key-value memory structure, in which the keys represent abstract concepts and values the associated characteristics of the concepts.
A new input-output pair matches with some values, and the indirection enables us to compute the functions based on corresponding keys.
Note that we have swapped the role of keys and values of the traditional memory (where the query is matched against the key, not value).
This is to emphasize that we perform indirection to map concrete functional values to abstract functional keys.
3. Functional Indirection Neural Estimator
In this section, we present the Functional Indirection Neural Estimator (FINE), an neural architecture to realize the general idea laid out in Section 3.1.
FINE learns to compose a function mapping an input embedding x to an output embedding y on-the-ﬂy.
The function is drawn from a parametric family speciﬁed by a backbone neural network.
Coupled with the backbone is a function composer φ , which is trained to compute the parameters of the backbone.
More speciﬁcally, φ maps a data input-output pair to the weights of the neural networks.
FINE solves the proposed IQ tasks as follows:
(1) Given two images x, y from the hint, FINE uses φ to produce the function transforming x to y ; denoted by φ x,y .
(2) Then, we feed the third image x (cid:48) to φ x,y and get output y ∗ =
φ x,y ( x (cid:48) ) .
(3) We deﬁne a similarity metric to compare y ∗ with given choices.
The choice with the closest distance from y ∗ is model’s answer.
The function composer and the similarity metric are speciﬁed as follows.
Encoder : Images are encoded by a trainable encoder (see Fig. 1).
To effectively solve IQ tasks introduced in Section 2, we use the p 4 -CNN [7] encoder to serve as an inductive bias for geometric transformations.
Throughout this paper, we refer to the images by their embeddings.
The functional memory
Let us focus a weight matrix W t ∈ R d in t × d out t at the t -th layer of the backbone network.
In practice,
W t belongs to the huge d in t ×
d out t -dimensional real matrix space M d in t ,d out t .
To reduce the complexity of W t , we assume that W t only belongs to a s -dimensional subspace of M d in t ,d out t , where s (cid:28) d in t d out t .
This subspace has a basis of s matrices which are trainable and stored in FINE ’s memory, and W t will be written as a linear combination of these matrices.
Denote by M t the memory for the t -th backbone layer.
M t includes two sub-memories: the key memory M key t = { M key t, 1 , . . .
, M key t,s } and the corresponding value memory M value t = { M value t, 1 , . . .
, M value t,s } .
Elements of M key t and M value t are trainable matrices of size d in t × d out t .
We further let x t ∈ R d in t and y t ∈ R d out t be the associated input and output, respectively, where x t is output of the ( t − 1) -th layer and y t is the pseudo-output computed by a trainable 1-layer neural network y t
= γ t ( y ) .
With this design, we control the complexity of functional hypothesis space by either constraining the form of the backbone or the capacity of the functional memory.
Memory reading
By the virtue of simplicity, we aim to ﬁnd a simple query that can demonstrate the relation between x t and y t .
Although the exact relation may be non-linear, we found that a query induced from linear relation is enough to efﬁciently read from memory.
Formally, we want to ﬁnd a query W qt such that
W qt x t = y t .
The best-approximated solution is W qt =
y t x + t , where x + t is the pseudo-inverse of x t and can be efﬁciently approximated by the iterative Ben-Israel and Cohen algorithm [3].
This way of query computing requires no parameter as opposed to other methods, where the input is often fed into a trainable neural network to compute the query.
With the query in hand, the next step is to perform analogy-making.
In FINE , the query W qt represents for the current situation and the value memory M value t consists of past experiences.
The concrete query interacts with value memories to measure how close the current situation is to each of the experiences.
The similarities are computed as dot products between the query and value memories and normalized by a factor of (cid:112) d in t d out t :
[Equation]
where the flatten operator ﬂattens the matrix W qt of size d in t × d out t into a vector of size d in t d out t , while the fconcat operator ﬁrst ﬂattens all matrices in M value t , then concatenates them together to form the value matrix of size ( d in t d out t ) × s .
The resulting a t is a s -dimensional vector measuring the similarities between the query and the entries in the value memory.
Here we omit the softmax operator as in usual attention to allow the similarities with more freedom.
The same idea is shared in the ESBN [33], where the softmax similarities are scaled by a sigmoidal factor.
Finally, the value memories are bound with their associated key memories via indirection.
This can be understood as moving forward from the concrete space of data and value memories to the abstract space of functions and key memories.
With the key memories and the similarity vector a t , the weight
W t of current backbone layer can be computed as the linear combination of key memories:
[Equation]
where the reshape operator reshapes the vector of size d in t d out t to a matrix of size d in t × d out t .
Since the softmax is omitted when calculating the similarities, W t is not constrained to be in the convex hull of the key memories and indeed can lie anywhere in the subspace spanned by those key memories.
Memory
update
The key and value memories are updated using gradient descent:
[Equation]
where λ > 0 is the learning rate and L is the loss of the training step.
The similarity metric
After determining φ x,y , the model is given a new input x (cid:48) and being asked to select the correct associated output y (cid:48) among 4 choices y (cid:48) 1 , y (cid:48) 2 , y (cid:48) 3 , y (cid:48)
so that ( x (cid:48) , y (cid:48) ) follows the same transformation rule as ( x, y ) .
This problem can be cast as ﬁnding the choice that is the most similar with y ∗ =
φ x,y ( x (cid:48) ) .
We consider the weighted Euclidean metric that measures the distance between two vectors u, v ∈ R d :
[Equation]
where { α i } di =1 ≥ 0 are trainable scalars, i.e., each component of u and v contributes with different importance.
Finally, the probability to pick a choice is computed as:
[Equation]
Experiments
We conduct experiments to show the out-of-distribution generalization capability of FINE when performing tasks introduced in Section 2.
For non-invertible transformations (e.g. reﬂection), we use a simple 2-layer MLP as the backbone.
For invertible transformations, we use the NICE architecture [9] as backbone to serve as an inductive bias for invertibility.
Since in each NICE layer, only half of the input is transformed, we use the same memory for two consecutive layers, i.e., M 2 t =
M 2 t +1 .
To balance between the backbone complexity and computational cost, we use 4 NICE layers in all experiments.
We compare FINE with three major classes of models: (a) models that make analogies in the data space, including Transformer [30], PrediNet [29] and RelationNet [28]; (b) models that leverage indirection to bind feature vectors with associated symbols and reason on the symbols, including the
ESBN [33]; and (c) models that aim to learn a mapping from data to the functional space, including the HyperNetworks [14].
For HyperNetworks, we still use the NICE backbone and just apply their fast-weight generation method for fair comparisons.
Except for FINE and HyperNetworks, all models are trained with context normalization [32], which has been proved to be effective in improving the generalization ability.
Datasets & implementation : We generate data for IQ tasks described in Section 2 using images from
Omniglot dataset [5], which includes 1,623 handwritten characters, and real-image CIFAR100 dataset [17].
If not speciﬁed, models are trained with p 4 -CNN encoder [7].
Experiments are conducted using
PyTorch on a single GPU with Adam optimizer.
Reported results are averages of 10 runs.
4. Results on Omniglot Dataset
[Table]
We use 100 characters for training and other 800 characters for testing.
The train and test set size is 10,000 and 20,000, respectively.
For FINE , we use 4 NICE layers with 48 memories for each pair of NICE layers.
Experimental results for single-transformation tasks are shown in Table 1.
Overall,
FINE dominates others with large margins.
For example, the gap to the runner-up on the Rotation task is nearly 13%.
With test accuracy over 75% on all tasks, FINE shows a strong capability of out-of-distribution generalization.
We further conduct multi-afﬁne-transformation experiments.
In this case, the training set includes multiple types of afﬁne transformations, while other settings are similar to the single-transformation case.
Results are also reported in Table 1.
FINE continues to outperform other models.
This is because only FINE explicitly assumes the existence of multiple good functions that represent the transformation from input to output data.
We note that although HyperNet also makes a similar assumption by generating data-speciﬁc weights, it does not utilize analogy-making and indirection and thus, fails to generalize to unseen images.
4. Results on CIFAR100 Dataset
[Table]
Trans
[Table]
RelationNet
[Table] [Table]
PrediNet
[Table] [Table]
Transformer
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
For CIFAR100 dataset, we follow similar settings as in experiments on the Omniglot dataset, except that we use 50 classes for training and 50 remaining classes for testing.
We also conduct experiments on single-transformation and multi-afﬁne-transformation tasks.
Results for single-transformation tasks are reported in Table 2.
Again, FINE outperforms all other models on all tasks, especially on
Reﬂection where the gap is nearly 30%.
Although FINE does not show good performance on Swap task as in Omniglot experiments, it is still slightly better than other models.
On the remaining tasks,
FINE achieves test accuracy of more than 80%.
For the multi-afﬁne-transformation task, we report the performances of the models when trained with different encoder architectures, including the p 4 -CNN, 3-layer ResNet and 2-layer MLP, in
Table 3.
The results show two superior characteristics of FINE : ﬁrst, FINE is consistently better than other models across different encoders; second, FINE is more stable with small standard deviations.
This empirical result supports the functional hypothesis stated in Section 3.1, where we suggest that focusing on functions distribution instead of data distribution can boost model’s generalization capability and stability.
4. More Extreme OOD Tasks
Previous tasks only include unseen classes of objects during testing.
In this section, we further test
FINE and related models on more challenging OOD tasks: tasks with unseen rules during training and even ones with images from unseen datasets.
The training and testing sets are either images from CIFAR100, Omniglot or MNIST datasets, while the hidden rules are either translation, rotation or shear.
For translation, training problems consist of translation vectors ( a, b ) with | a | , | b | ≤ 3 ,
[Table]
We backbone
Table 4.
As expected, FINE continues to outperform other models on most of the tasks, even on extreme OOD tasks with unseen datasets and unseen rules where performances of all models drop signiﬁcantly.
This demonstrates FINE is capable of effectively learning the basis weights, which are stored in the memory, to represent novel rules.
4.4 Model Analysis and Ablation Study Clustering on functional spaces
[Figure]
We study how the transformations are distributed in the functional space.
We use the FINE model trained on multi-afﬁne transformation tasks.
For an input-output pair ( x, y ) , we ﬂatten and concatenate all weights of NICE layers to form a vector representing φ x,y .
We then use the UMAP [21] to project φ x,y ’s vectors onto the 2D plane.
Results are shown in Fig.
It is interesting to see that shears with the same horizontal or vertical angles are positioned closely; and scales are separated into 2 “big” clusters, one for smaller scale and one for bigger scale.
In contrast, reﬂection representations seem not to be clustered properly, which is worth further investigation in future work.
Number of memories and backbone layers
We do an ablation study to see the effect of number of memories and backbone layers in FINE .
For the limit case with 0 memory, we assign the query matrices to be the weights for NICE layers without the analogy-making and indirection process.
For the limit case with 0 NICE layer, we replace the
NICE backbone by a 2-layer MLP.
Results are shown in Fig. 4(a).
Overall, we can observe clear improvements when we increase the number of memories or number of NICE layers.
More interestingly, the more number of memories is, the more stable the results will be.
Increasing the number of memories is equivalent to enlarging the range of φ , and increasing the number of NICE layers is equivalent to enlarging the hypothesis space
F .
Enlarging F may help FINE approach the true functions while still being sufﬁciently constrained by the number of memories, thus still being able to maintain its generalization capability.
[Figure]
Number of training data points
We train models on multi-afﬁne-transformation task with training sets of different sizes.
Results are reported in Fig. 4(b).
FINE can adapt with small datasets of sizes 100 or 1000 and achieve fair accuracy (47.8% and 64.4%, respectively).
Moreover, FINE achieves average test accuracy of 81.1% on training set of size 10,000, which is quite close to the 85.3% accuracy on 100,000 data points training sets.
This shows FINE may obtain a near-optimal solution even with small number of training data points.
In contrast, ESBN and Transformer need 100,000 training data points to achieve 80% or higher test accuracy, while only achieving roughly 70% test accuracy on smaller datasets.
Related work
In recent years, there has been a strong interest in designing models that are capable of generalizing systematically.
The Module Networks [1], which dynamically compose neural networks out of trainable modules, have been shown to possess some degree of systematic generalization [2].
Parascandolo et al.
[24] proposed a method to train multiple competing experts to explain image transformations on MNIST and Omniglot datasets, yet the transformations are simpler than ones in our FINE dataset and it is not clear whether the proposed method can deal with unseen transformations.
The Neural
Interpreter [25] uses attention mechanism to recompose functional modules for each input-output pair and test their method on abstract reasoning tasks, yet tends to require a large amount of data to learn.
Switch Transformer [10] mitigates the communication and computational cost in mixture of experts models.
Recently, ESBN [33], which utilizes the indirection mechanism, shows a great promise on
OOD tasks.
Both methods achieve their degree of systematic generalization by injecting symbolic biases into the models.
Our model FINE follows the same strategy, but it performs analogy-making and indirection in functional spaces instead of data spaces, and this has proved to boost both the performance and stability.
IQ tests are powerful testbeds for visual and abstract reasoning.
Inspired by Raven’s Progressive
Matrices (RPM), the RAVEN dataset [34] was proposed as a testbed for visual reasoning models.
However, this dataset does not focus on testing the ability of OOD generalization.
Webb et al.
[33] propose a series of IQ tasks with Unicode characters to show the effectiveness of indirection in tasks involving abstract rules, however these tasks are relatively simple since they only require the models to understand the same-different relation.
The ARC dataset [6] aims to serve as a benchmark for general intelligence and includes various psychometric IQ tasks in the form of grid structures.
In this paper, we propose IQ tasks involving geometric transformations as introduced in Section 2.
These tasks are not only ﬂexible so that we can include images from different datasets or create/combine numerous transformations, but also challenging to test OOD generalization abilities of models.
The weight composition feature of FINE links back to the concept of fast weights [19], the idea of computing data-speciﬁc network weights on-the-ﬂy.
HyperNetworks [14] stylizes this idea by computing the fast weights using a separate trainable (slow weight) network.
The Meta-learned Neural
Memory [23] uses the pseudo-target technique (which we also leverage in FINE ) and appropriately updates the short-term memory once new input arrives.
The Neural Stored-Program Memory (NSM)
[18] proposes a hybrid approach between slow-weight and fast-weight to compute network weights on-the-ﬂy based on slow-weight key and value memories.
However, NSM only performs on sequential learning tasks, while FINE aims to solve OOD IQ tasks requiring abstract cognition.
Moreover, FINE computes the query based on the input and the (pseudo-) output, while NSM’s query is computed based on the input only.
Memories have been to be versatile in meta-learning and few-shot learning [15, 27, 31] due to the ability to rapidly store past examples and adapt to new situations.
In our case, an IQ task can be thought of as a one-shot learning scenario in which FINE has to make use of the long-term key-value memory to adapt to the current task.
Limitations
Operating in functional spaces requires FINE ’s memory to store several trainable weight matrices to infer the data-speciﬁc weights for the backbone.
Moreover, each layer of the backbone requires its own memory, thus FINE may need a large number of parameters for very deep backbones.
This could be addressed by parameter-sharing across layers and limiting the rank of the weight matrices.
It remains to design the backbone architecture optimally.
We have used the NICE architecture as backbone for invertible transformations and a 2-layer MLP backbone for non-invertible cases.
We further investigated our model’s performances with larger number of training points (up to 1M) and observed that FINE with the NICE backbone peaks at 100,000 training points with around 85% test accuracy, while FINE with MLP backbone continues to improve and achieves around 95% test accuracy at 1M training points.
This calls for further theoretical analysis to guide the architectural design of the backbone network.
Finally, testing FINE on IQ tasks in the visual space may limit its potential on IQ tasks involving other modalities.
For example, one may consider a textual IQ problem: “if abc → abd, then mmnnpp → ? ”.
It is worth to emphasize that the concepts of analogy-making and indirection in functional spaces are indeed general and thus the idea of FINE should be applicable to various scenarios.
Conclusion
To study the out-of-distribution (OOD) generalization capability of models, we have proposed IQ tasks that require models to rapidly recognize the hidden rules of geometric transformations between a pair of images and transfer the rules to a new pair of different image classes.
Such tasks would necessitate human-like abilities for conceptual abstraction, analogy-making, and utilizing indirection.
We put forward a hypothesis that these mechanisms should be performed in the functional space instead of data space as in current deep learning models.
To realize the hypothesis, we then proposed
FINE , a memory-augmented neural architecture that learns to compose functions mapping an input to an output on-the-ﬂy.
The memory has two trainable components: the value sub-memory and the binding key sub-memory, where the keys are basis weight matrices that span the space of functions.
For an IQ task, when given a hint in the form of an input-output pair, and FINE estimates the analogy between the pair and the values as mixing coefﬁcients.
These coefﬁcients are then used to mix the binding keys via indirection to generate the weights of the backbone neural net which computes the intended function.
For a test input of different class, the function is used to estimate the most compatible output, thus solving the IO task.
Through an extensive suite of experiments using images from the Omniglot and CIFAR100 datasets to construct the IQ tasks, FINE is found to be reliable in ﬁguring out the hidden relational pattern in each IQ task and thus is able to solve new tasks, even with unseen image classes.
Importantly, FINE outperforms other models in all experiments, and can generalize well in small data regimes.
Future works will include making FINE robust against OOD in transformations without catastrophic forgetting when new transformations are continually introduced.