Title: TASK-MEDIATED REPRESENTATION LEARNING

Abstract: Traditionally, unsupervised representation learning is used to discover underlying regularities from raw sensory data without relying on labeled data. A great number of algorithms in this field resorts to utilizing proxy objectives to facilitate learning. Further, learning how to act upon these regularities is left to a separate algorithm. Neural encoding in biological systems, on the other hand, is optimized to represent behaviorally relevant features of the environment in order to make inferences that guide successful behavior. Evidence suggests that neural encoding in biological systems is shaped by such behavioral objectives. In our work, we propose a model of inference-driven representation learning. Rather than following some auxiliary, a priori objective (e.g. minimization of reconstruction error, maximization of the fidelity of a generative model, etc.) and indiscriminately encoding information present in an observation, our model learns to build representations that support accurate inferences. Given a set of observations, our model encodes underlying regularities that de facto are necessary to solve the inference problem in hand. Rather than labeling the observations and learning representations that portray corresponding labels or learning representation in a self-supervised manner and learning explicit features of the input observations, we propose a model that learns representations that implicitly shaped by the goal of correct inference.

1. INTRODUCTION:

One of the central problems in unsupervised representation learning traditionally has been to find a way of learning useful representations when target values are not available or simply do not exist (Hinton & Sejnowski, 1999; LeCun et al., 2015) . A common approach to this problem is to essentially compress high dimensional input s.t. it can be later reproduced from a low dimensional representation with minimum deviation from either the source or any other proxy objective. For instance, autoencoders learn latent representations by minimizing the reconstruction error (Hinton & Salakhutdinov, 2006) ; variational autoencoders additionally constrain the representation which allows for learning the factors of variation (Kingma & Welling, 2013; Higgins et al., 2017a) ; generative models can learn representations by maximizing the fidelity of generated outputs (Donahue et al., 2016) . Once representations are learned, their utility is usually evaluated on a separate downstream task (Hsu et al., 2018; Higgins et al., 2017b; Metz et al., 2019; Laversanne-Finot et al., 2018) . On the other hand, there is a significant amount of evidence that suggests that perceptual neural encoding in biological systems is shaped by the process of inference supporting effective behavior (Janzen & Turennout, 2004) . Theory of embodied cognition suggests that the process of information internalization requires physical interaction with the environment (Calvo & Gomila, 2008) . In our work, we propose one of the possible ways of learning representations of abstract concepts from unlabeled data that in turn support the inference. We argue that an abstract surrogate objective such as correct inference not only yields adequate representations but also a biologically valid way of learning representations.

We organized the paper by first defining and describing our dataset and the model in Section 2 and Section 3 respectively. In Section 4, we describe the results and provide an analysis of the network behavior and performance, demonstrating the robustness of learned representations. Then, in order to contextualize our work, we provide a summary of related research and empirical evidence that supports our approach.

2. DATASET:

In our work, we use a synthetic dataset that consists of several types of inference tasks and observations O. In general, the task is to make accurate inference based on a problem P which requires correct interpretation of given observations O. Every observation is generated such that it follows a certain pattern or regularity (i.e. regularities are implicitly encoded in the observations). Based on the regularities encoded in the tasks, they can be divided into three categories: (1) binary operations, (2) contextual 2-armed bandit, and (3) associative recall. In the case of binary operations, a single observation can be often ambiguous hence we use a set of three observations to reduce ambiguity. Finally, we use the results of the inference problems as targets. Learning algorithm does not use labels of the regularities placed in the observations. The data is presented in the form of tensors of binary activations. For the simplicity of data interpretation, we address each element as if it was a number. Here, the underlying regularity is based on binary function f (x 1 , x 2 ) is: x 3 = x 1 + x 2 . (Bottom left) Structure of Associative Recall task: x 1 and x 2 together constitute a key and x 3 is the target value. (Bottom right) Structure of Contextual 2-Armed Bandit task: x 3 = 0 means that the left arm can be associated with the reward, if x 3 = 1, then it is the right arm (here, the arm 13 should be preferred.)

2.1. BINARY OPERATIONS:

This part of the dataset includes 7 functions:

1. x 3 = x 1 + x 2 + c, where c ∈ [0, 1, 2, 3]

2. x 3 = sd(x 1 ) + sd(x 2 ), where sd is digit sum of an integer 3.

x 3 = x 1 > x 2 , x 3 ∈ 0, 1 4. x 3 = x 1 < x 2 , x 3 ∈ 0, 1 5. x 3 = x 1 -x 2 , s.t. x 3 > 0 6. x 3 = x 2 -x 1 , s.t. x 3 > 0 7. x 3 = x 2 div x 1 ,

where div is integer division Each of the observations is always structured such that the third number in it is a result of a mathematical or logic binary operation applied to the first two numbers:

x 3 = f (x 1 , x 2 ).

To make an inference, the model needs to determine the function f based on the observations O and apply it to the given problem P: ŷ = f (P). Within each function included in the dataset, each pair of x 1 and x 2 is unique.

2.2. CONTEXTUAL 2-ARMED BANDIT:

The problem of the contextual 2-armed bandit is a variant of a well-known multi-armed bandit problem (Auer et al., 2002) . The 2-armed bandit problem is usually used in meta-learning research. We, however, reframe it as an inference problem. Given a set of observations, the task is to pick either left or right arm. Each arm has a number associated with it. In the contextual bandit problem, it is not necessarily the left or the right arm that is more likely to score but the one with a label (number in our dataset) that, in the observations, corresponds with the right decision. The goal is therefore to pick the arm which is more likely to score. In our dataset we use the following reward distribution: 66 2 3 % or 33 1 3 %.

2.3. ASSOCIATIVE RECALL:

Associative recall task is reminiscent of a simple short-term memory task: given a dictionary (i.e. a set of (key, value) tuples), the task is to recall the value that corresponds to the given key in the dictionary. Every key in the dataset is unique, hence, every association is also a unique combination. One of the major goals of this work is to introduce a representation learning mechanism that combines several observations and is focused on discovering regularities underlying them (implicit information) rather than on compressing input data into more interpretable features of explicitly presented information. We suggest that one of the possible ways to do that is to learn representations not independently but in the process of observation facilitated inference. Hence, our dataset only contains the results for the inference problems but does not label the underlying regularities. In fact, in some cases, it is possible to assign several labels to a single regularity (for details see Section 4) which makes labeling counterproductive. Such an approach has its own limitations, yet it is based on certain aspects of representation learning in biological systems (see Section 5.2 for details).

3. MODEL:



3.2.1. OBJECTIVE FUNCTION:

Our model is trained to minimize the discrepancy between inferred result and actual result. We also use a regularization term R to make representation r more efficient, as suggested in (Kingma & Welling, 2013; Higgins et al., 2017a) . The loss function is the following:

EQUATION

We use multivariate KullbackLeibler divergence (D KL ) (2) between a diagonal multivariate normal, and a standard normal distribution (Kullback and Leibler, 1951) for representation regularization purposes:

EQUATION

Thus, the objective is the following:

EQUATION

3.3. IMPLEMENTATION:

We implemented our model in Tensorflow (Abadi et al., 2016 ) and used the Adam optimization algorithm for training (Kingma & Ba, 2014) . We use a synthetic dataset generated according to Section 2. The size of the dataset depends on a set of variables and can be customized. The content of the dataset can also be customized by excluding and including certain problems. Our model was trained on all problems described in Section 2 together. We split the dataset 50/50 for training and testing respectively. In this paper, we report only test performance. 

4. RESULTS:

In our work, we demonstrate a model of representation learning mediated by an inference process. We show that certain abstract concepts can be learned from observations only retrospectively in the process of making inferences. Empirical evidence suggests that, when it comes to representing abstract concepts, biological systems, rather than representing a maximum number of features, facilitate successful behavior with the selective neural encoding of features relevant at the behavioral level. Since it is difficult to imagine an a priori rule that can extract relevant information, we trained our model using the objective that optimizes for the accuracy of the inference. Overall test performance of our model (98.6±0.1%) suggests that it was able to learn representations that facilitate inference.

To visualize the representations learned by our model we reduced their dimensionality to 2 via an autoencoder and plotted the results. We also plotted the centers of the manifolds of representations and Figures 2 and 3 suggest that our model learned certain relationships between the concepts that these representations depict. For instance, virtually complete overlap between representations of

x 3 = x 1 -x 2 and x 3 = x 2 -x 1 , s.t.

x 3 > 0 shows that our model found that these operations are equivalents of the absolute difference between x 1 and x 2 . Close proximity and order in the group of representations for x 3 = x 1 + x 2 + c, where c ∈ [0, 3] shows that the model discovered the relationship between these operations. Overlap between special cases of x 3 = x 1 > x 2 and x 3 = x 1 div x 2 can be interpreted as well: for all x 1 and x 2 : x 1 > x 2 both functions will output 0. One of the major goals of representation learning has been an ability to encode raw data in the way such an encoding could be useful in subsequent learning of a downstream task (Kingma & Ba, 2014; Rasmus et al., 2015; Hsu et al., 2018) or have useful qualities, such as interpretability, smoothness of the manifold, explanatory power, sparsity, disentanglement of the underlying factors of variation, etc. (Kingma & Welling, 2013; Higgins et al., 2017a; Donahue et al., 2016; Bengio et al., 2014) . Hinton & Salakhutdinov (2006) showed that autoencoders (AE) can transform high dimensional data into low dimensional representations that minimize reconstruction error. Kingma & Welling (2013) presented Variational Auto Encoders (VAE) an approach to learning representations that encode high-level factors of variation. More recent, Higgins et al. (2017b) pre-trained β-VAE to encode basic visual concepts of the visual environment, effectively disentangling such variational factors. In turn, latent representations produced by β-VAE supported a reinforcement learning model and have proven to be particularly useful for transfer learning. Similarly, Laversanne-Finot et al. (2018) utilized β-VAE to learn representations that contain independent features. Then, a separate model used pre-trained β-VAE to support curiosity-driven exploration in a robotic arm task. Hsu et al. (2018) used several unsupervised embedding algorithms to learn representations useful in subsequent object discrimination procedure. Their results suggest that the use of embeddings improved the performance of the discrimination model. Goodfellow et al. (2014) demonstrated the ability of Generative Adversarial Networks (GANs) to learn a mapping from latent distribution to data distribution and captures semantic variation features. Further work, Bidirectional Generative Adversarial Networks (BiGAN) (Donahue et al., 2016) , made it possible to learn a mapping from data to latent representations useful for a supervised discrimination task.

5.2. TASK-MEDIATED NEURAL ENCODING:

A significant body of empirical evidence and theoretical work suggests that neuronal encoding emerges and further develops mediated by the behavioral objectives (Oxenham (ed.) & Oxenham, 2005; Kuchibhotla & Bathellier, 2018; Kraus et al., 2014; Janzen & Turennout, 2004) . Biological systems, in order to succeed in achieving these objectives, need to make inferences based on sensory input. In turn, the efficiency of the inference is contingent on the ability to encode relevant information. Sustained activity in PFC support context-relevant representations when a task requires different behavior based on the context (Rikhye et al., 2018; Bolkan et al., 2018) . The visual system of a fly is tuned to process and represent optic flow, which is perhaps the most important visual feature at the behavioral level of a relatively simple animal (Egelhaaf et al., 2002) . Recent research suggests that musical training improves auditory processing by enhancing neural encoding of meaningful acoustic features which, in turn, benefits language and cognition (Kraus et al., 2014; Tierney et al., 2013; Parbery-Clark et al., 2009; Kraus et al., 2010) . Implicit relevance of objects placed at key points was associated with neural activity in the parahippocampal gyrus (area of the brain implicated in visual navigation) whereas explicit relevance improved object recognition processing in the visual cortex (Janzen & Turennout, 2004) . Similarly, in auditory cortex, encoding has been shown to be mediated by ascending neuromodulation and top-down attention (Kuchibhotla & Bathellier, 2018; Caras & Sanes, 2017) .

6. DISCUSSION:

The ability to perceive regularities in the raw perceptual data one the most important skills that biological systems develop very early and continue to develop and which facilitates learning and other cognitive abilities. Development of this ability does not require explicit training as it is observed in various biological systems and perhaps one of the fundamental mechanisms in learning. Likewise, unsupervised learning, a discipline in machine learning, has been focused on problems that cannot and do not have explicit training or labels. We showed that the absence of labels does not automatically mean that the process of learning of the regularities is driven by a single universal objective (e.g. minimization of reconstruction error). Rather, learning of representations of relevant information can be shaped by behavioral objectives and enables learning of abstract concepts that cannot be learned solely from observations.

The utility of the representations learned in such a manner will depend on their behavioral relevance and experience. We showed that simple model that is given a reasonable objective, can learn representations that even very clever encoding algorithms could not in a self-supervised manner. Hence, in conclusion, we would like to emphasize the role of intrinsic motivation (i.e. objective) and the environment in learning useful representations.

:



Figures and Tables:

Figure 1: (Top) Structure and possible interpretation of a training entry.Here, the underlying regularity is based on binary function f (x 1 , x 2 ) is:x 3 = x 1 + x 2 . (Bottom left) Structure of Associative Recall task:x 1 and x 2 together constitute a key and x 3 is the target value. (Bottom right) Structure of Contextual 2-Armed Bandit task: x 3 = 0 means that the left arm can be associated with the reward, if x 3 = 1, then it is the right arm (here, the arm 13 should be preferred.)

MODEL ARCHITECTUREOur model is a two-tailed neural network. It takes a set of observations O and inference problem P and makes an inference ŷ : ŷ = H(O , P ; θ H ). The first tail takes observations one-by-one, consolidates them in the internal state of a recurrent cell, passes combined representation through a bottleneck r, and finally combines its output O = F (O; θ F ) with the output of the second tail P = G(P; θ G ). Second tail and the head components of the network are stacks of fully connected layers.

Figure 2: Model architecture

Figure 3: Learned representations of mathematical problems, associative recall, and contextual 2-armed bandit

Figure 4: Distance between manifolds of learned representations

The number of observations is a variable. In this work, we use 3 observations o i , yet our model shows adequate results for both fewer and more observations. For recurrent layer we use Gated Recurrent Unit cell(Cho et al., 2014) with an internal size of 96. The size of the bottleneck r is 8.

Table content: <table/>

Peer review (iclr2020):
The paper is driven by the argument that neural encoding in biological systems is shaped by behavioral objectives and, hence, proposes to make use of an abstract surrogate objective. This is said to result in a model that learns representations that are implicitly shaped by the goal of correct inference. Indeed, closing the gap between neural encodings in biological and artificial systems is an important goal. Unfortunately, the paper is not well presented. There are a number grammar flows. Or consider e.g. Fig. 1. Here, the variable Y is not explained. It is also not easy to understand *the arm 13* can be found. Or consider the introduction of the model. It follows very basic differentiable programming units, while the abstract claims that *rather than labeling the observations ... we propose a model that learns representations that are implicitly shaped by the goals of correct inference*. Ok, now the inference is the label. Fine. But what is the benefit of doing so? Unfortunately, this is not discussed and also the experimental evaluation does not show this. A comparison to VAE on a single dataset is not enough to convince the reader that this generally holds. Moreover, it is not shown that the representation now helps better in some applications than existing approaches. It would be much more convincing the new models beats VAE in a task VAE are good at. In any case, evaluations on more tasks/datasets are definitely missing to show that the method generally works fine. To summarize, interesting direction but the paper is not ready for publication at ICLR yet.