DRAGNN: A TRANSITION-BASED FRAMEWORK FOR DYNAMICALLY CONNECTED NEURAL NETWORKS
In this work, we present a compact, modular framework for constructing new recurrent neural architectures. Our basic module is a new generic unit, the Transition Based Recurrent Unit (TBRU). In addition to hidden layer activations, TBRUs have discrete state dynamics that allow network connections to be built dynamically as a function of intermediate activations. By connecting multiple TBRUs, we can extend and combine commonly used architectures such as sequence-tosequence, attention mechanisms, and recursive tree-structured models. A TBRU can also serve as both an encoder for downstream tasks and as a decoder for its own task simultaneously, resulting in more accurate multi-task learning. We call our approach Dynamic Recurrent Acyclic Graphical Neural Networks, or DRAGNN. We show that DRAGNN is significantly more accurate and efficient than seq2seq with attention for syntactic dependency parsing and yields more accurate multi-task learning for extractive summarization tasks.
1 INTRODUCTION
To apply deep learning models to structured prediction, machine learning practitioners must address two primary issues: (1) how to represent the input, and (2) how to represent the output. The seq2seq encoder/decoder framework (Kalchbrenner & Blunsom, 2013; Cho et al., 2014; Sutskever et al., 2014) proposes solving these generically. In its simplest form, the encoder network produces a fixed-length vector representation of an input, while the decoder network produces a linearization of the target output structure as a sequence of output symbols. Encoder/decoder is state of the art for several key tasks in natural language processing, such as machine translation (Wu et al., 2016).
However, fixed-size encodings become less competitive when the input structure can be explicitly mapped to the output. In the simple case of predicting tags for individual tokens in a sentence, stateof-the-art taggers learn vector representations for each input token and predict output tags from those (Ling et al., 2015; Huang et al., 2015; Andor et al., 2016). When the input or output is a syntactic parse tree, networks that explicitly operate over the compositional structure of the network typically outperform generic representations (Dyer et al., 2015; Li et al., 2015; Bowman et al., 2016). Implictly learned mappings via attention mechanisms can significantly improve the performance of sequence-to-sequence (Bahdanau et al., 2015; Vinyals et al., 2015), but require runtime that’s quadratic in the input size.
In this work, we propose a modular neural architecture that generalizes the encoder/decoder concept to include explicit structure. Our framework can represent sequence-to-sequence learning as well as models with explicit structure like bi-directional tagging models and compositional, tree-structured models. Our core idea is to define any given architecture as a series of modular units, where connections between modules are unfolded dynamically as a function of the intermediate activations produced by the network. These dynamic connections represent the explicit input and output structure produced by the network for a given task.
We build on the idea of transition systems from the parsing literature (Nivre, 2006), which linearize structured outputs as a sequence of (state, decision) pairs. Transition-based neural networks have recently been applied to a wide variety of NLP problems; Dyer et al. (2015); Lample et al. (2016);
Kiperwasser & Goldberg (2016); Zhang et al. (2016); Andor et al. (2016), among others. We generalize these approaches with a new basic module, the Transition-Based Recurrent Unit (TBRU), which produces a vector representation for every transition state in the output linearization (Figure 1). These representations also serve as the encoding of the explicit structure defined by the states. For example, a TBRU that attaches two sub-trees while building a syntactic parse tree will also produce the hidden layer activations to serve as an encoding for the newly constructed phrase. Multiple TBRUs can be connected and learned jointly to add explicit structure to multi-task learning setups and share representations between tasks with different input or output spaces (Figure 2).
This inference procedure will construct an acyclic compute graph representing the network architecture, where recurrent connections are dynamically added as the network unfolds. We therefore call our approach Dynamic Recurrent Acyclic Graphical Neural Networks, or DRAGNN.
DRAGNN has several distinct modeling advantages over traditional fixed neural architectures. Unlike generic seq2seq, DRAGNN supports variable sized input representations that may contain explicit structure. Unlike purely sequential RNNs, the dynamic connections in a DRAGNN can span arbitrary distances in the input space. Crucially, inference remains linear in the size of the input, in contrast to quadratic-time attention mechanisms. Dynamic connections thus establish a compromise between pure seq2seq and pure attention architectures by providing a finite set of long-range inputs that ‘attend’ to relevant portions of the input space. Unlike recursive neural networks (Socher et al., 2010; 2011) DRAGNN can both predict intermediate structures (such as parse trees) and utilize those structures in a single deep model, backpropagating downstream task errors through the intermediate structures. Compared to models such as Stack-LSTM (Dyer et al., 2015) and SPINN Bowman et al. (2016), TBRUs are a more general formulation that allows incorporating dynamically structured multi-task learning (Zhang & Weiss, 2016) and more varied network architectures.
In sum, DRAGNN is not a particular neural architecture, but rather a formulation for describing neural architectures compactly. The key to this compact description is a new recurrent unit—the TBRU—which allows connections between nodes in an unrolled compute graph to be specified dynamically in a generic fashion. We utilize transition systems to provide succinct, discrete representations via linearizations of both the input and the output for structured prediction. We provide a straightforward way of re-using representations across NLP tasks that operate on different structures.
We demonstrate the effectiveness of DRAGNN on two NLP tasks that benefit from explicit structure: dependency parsing and extractive sentence summarization (Filippova & Altun, 2013). First, we show how to use TBRUs to incrementally add structure to the input and output of a “vanilla” seq2seq dependency parsing model, dramatically boosting accuracy over seq2seq with no additional computational cost. Second, we demonstrate how the same TBRUs can be used to provide structured intermediate syntactic representations for extractive sentence summarization. This yields better accuracy than is possible with the generic multi-task seq2seq (Dong et al., 2015; Luong et al., 2016) approach. Finally, we show how multiple TBRUs for the same dependency parsing task can be stacked together to produce a single state-of-the-art dependency parsing model.
2 TRANSITION SYSTEMS
We use transition systems to map inputs x into a sequence of output symbols, d1 . . . dn. For the purposes of implementing DRAGNN, transition systems make explicit two desirable properties. First, we stipulate that the output symbols represent modifications of a persistent, discrete state, which makes book-keeping to construct the dynamic recurrent connections easier to express. Second, transition systems make it easy to enforce arbitrary constraints on the output, e.g. the output should produce a valid tree.
Formally, we use the same setup as Andor et al. (2016), and define a transition system T = {S,A, t} as:
• A set of states S(x). • A special start state s† ∈ S(x). • A set of allowed decisions A(s, x) for all s ∈ S. • A transition function t(s, d, x) returning a new state s′ for any decision d ∈ A(s, x).
For brevity, we will drop the dependence of x in the functions given above. Throughout this work we will use transition systems in which all complete structures for the same input x have the same number of decisions n(x) (or n for brevity), although this is not necessary.
A complete structure is then a sequence of decision/state pairs (s1, d1) . . . (sn, dn) such that s1 = s†, di ∈ A(si) for i = 1 . . . n, and si+1 = t(si, di). We will now define recurrent network architectures that operate over these linearizations of input and output structure.
3 TRANSITION BASED RECURRENT NETWORKS
We now formally define how to combine transition systems with recurrent networks into what we call a transition based recurrent unit (TBRU). A TBRU consists of the following:
• A transition system T , • An input function m(s) that maps states to fixed-size vector representations, for example,
an embedding lookup operation for features from the discrete state:
m(s) : S 7→ RK
• A recurrrence function r(s) that maps states to a set of previous time steps: r(s) : S 7→ P{1, . . . , i− 1},
where P is the power set. Note that in general |r(s)| is not necessarily fixed and can vary with s. We use r to specify state-dependent recurrent links in the unrolled computation graph. • A RNN cell that computes a new hidden representation from the fixed and recurrent inputs: hs ← RNN(m(s), {hi | i ∈ r(s)}).
Example 1. Sequential tagging RNN. Let the input x = {x1, . . . ,xn} be a sequence of word embeddings, and the output be a sequence of tags d1, . . . , dn. Then we can model a simple LSTM tagger as follows:
• T sequentially tags each input token, where si = {1, . . . , di−1}, andA is the set of possible tags. We call this the tagger transition system. • m(si) = xi, the word embedding for the next token to be tagged. • r(si) = {i− 1} to connect the network to the previous state. • RNN is a single instance of the LSTM cell.
Example 2. Parsey McParseface. The open-source syntactic parsing model of Andor et al. (2016) can be defined in our framework as follows:
• T is the arc-standard transition system (Figure 3), so the state contains all words and partially built trees on the stack as well as unseen words on the buffer. • m(si) is the concatenation of 52 feature embeddings extracted from tokens based on their positions in the stack and the buffer. • r(si) = {} is empty, as this is a feed-forward network. • RNN is a feed-forward multi-layer perceptron (MLP).
Inference with TBRUs. Given the above, inference in the TBRU proceeds as follows:
1. Initialize s1 = s†. 2. For i = 1, . . . , n:
(a) Update the hidden state: hi ← RNN(m(si), {hj | j ∈ r(si)}). (b) Update the transition state: di ← argmaxd∈A(si) w>d hi, si+1 ← t(si, di).
A schematic overview of a single TBRU is presented in Figure 3. By adjusting RNN, r, and T , TBRUs can represent a wide variety of neural architectures.
3.1 CONNECTING MULTIPLE TBRUS TO LEARN SHARED REPRESENTATIONS
While TBRUs are a useful abstraction for describing recurrent models, the primary motivation for this framework is to allow new architectures by combining representations across tasks and compo-
sitional structures. We do this by connecting multiple TBRUs with different transition systems via the recurrence function r(s). We formally augment the above definition as follows:
1. We execute a list of T TBRU components, one at a time, so that each TBRU advances a global step counter. Note that for simplicity, we assume an earlier TBRU finishes all of its steps before the next one starts execution. 2. Each transition state from the τ ’th component sτ has access to the terminal states from every prior transition system, and the recurrence function r(sτ ) for any given component can pull hidden activations from every prior one as well.
Example 3. “Input” transducer TBRUs via no-op decisions. We find it useful to define TBRUs even when the transition system decisions don’t correspond to any output. These TBRUs, which we call no-op TBRUs, transduce the input according to some linearization. The simplest is the shiftonly transition system, in which the state is just an input pointer si = {i}, and there is only one transition which advances it: t(si, ·) = {i + 1}. Executing this transition system will produce a hidden representation hi for every input token.
Example 4. Encoder/decoder networks with TBRUs. We can reproduce the encoder/decoder framework for sequence tagging by using two TBRUs: one using the shift-only transition system to encode the input, and the other using the tagger transition system. For input x = {x1, . . . ,xn}, we connect them as follows:
• For shift-only TBRU: m(si) = xi, r(si) = {i− 1}. • For tagger TBRU: m(sn+i) = ydn+i−1 , r(si) = {n, n+ i− 1}.
We observe that the tagger TBRU starts at step n after the shift-only TBRU finishes, that yj is a fixed embedding vector for the output tag j, and that the tagger TBRU has access to both the final encoding vector hn as well as its own previous time step hn+i−1.
Example 4. Bi-directional LSTM tagger. With three TBRUs, we can implement a simple bidirectional tagger. The first two run the shift-only transition system, but in opposite directions. The final TBRU runs the tagger transition system and concatenates the two representations:
• Left to right: T = shift-only, m(si) = xi, r(si) = {i− 1}. • Right to left: T = shift-only, m(sn+i) = xn−i, r(sn+i) = {n+ i− 1}. • Tagger: T = tagger, m(s2n+i) = {}, r(s2n+i) = {i, 2n− i}.
We observe that the network cell in the tagger TBRU takes recurrences only from the bi-directional representations, and so is not recurrent in the traditional sense. See Figure 1 for an unrolled example.
Example 5. Multi-task bi-directional tagging. Here we observe that it’s possible to add additional annotation tasks to the bi-directional TBRU stack from Example 4 simply by adding more instances of the tagger TBRUs that produce outputs from different tag sets, e.g. parts-of-speech vs. morphological tags. Most important, however, is that any additional TBRUs have access to all three earlier TBRUs. This means that we can support the “stack-propagation” (Zhang & Weiss, 2016) style of multi-task learning simply by changing r for the last TBRU:
• Traditional multi-task: r(s3n+i) = {i, 2n− i} • Stack-prop: r(s3n+i) = { i︸︷︷︸
Left-to-right , 2n− i︸ ︷︷ ︸ Right-to-left , 2n+ i︸ ︷︷ ︸ Tagger TBRU }
Remark: the raison d’être of DRAGNN. This example highlights the primary advantage of our formulation: a TBRU can serve as both an encoder for downstream tasks and as a decoder for its own task simultaneously. This idea will prove particularly powerful when we consider syntactic parsing, which involves compositional structure over the input. For example, consider a no-op TBRU that traverses an input sequence x1, . . . ,xn in the order determined by a binary parse tree: this transducer can implement a recursive tree-structured network in the style of Tai et al. (2015), which computes representations for sub-phrases in the tree. In contrast, with DRAGNN, we can
use the arc-standard parser directly to produce the parse tree as well as encode sub-phrases into representations.
Example 6. Compositional representations from arc-standard dependency parsing. We use the arc-standard transition system (Nivre, 2006) to model dependency trees. The system maintains two data structures as part of the state s: an input pointer and a stack (Figure 3). Trees are built bottom up via three possible attachment decisions. Assume that the stack consists of S = {A,B}, with the next token being C. We use S0 and S1 to refer to the top two tokens on the stack. Then the decisions are defined as:
• Shift: Push the next token on to the stack: S = {A,B,C}, and advance the input pointer. • Left arc + label: Add an arc A←label B, and remove A from the stack: S = {B}. • Right arc + label: Add an arc A→label B, and remove B from the stack: S = {A}.
For a given parser state si, we compute two types of recurrences:
• rINPUT(si) = {INPUT(si)}, where INPUT returns the index of the next input token. • rSTACK(si) = {SUBTREE(si, S0), SUBTREE(s, S1)}, where SUBTREE(S,I) is a function
returning the index of the last decision that modified the i’th token:
SUBTREE(s, i) = argmax j {dj s.t. dj shifts or adds a new child to token i}
We show an example of the links constructed by these recurrences in Figure 4, and we investigate variants of this model in Section 4. This model is recursively compositional according to the decision taken by the network: when the TBRU at step si decides to add an arc A → B for state, the activations hi will be used to represent that new subtree in future decisions.1
Example 7. Extractive summarization pipeline with parse representations. To model extractive summarization, we follow Andor et al. (2016) and use a tagger transition system with two tags: “Keep” and “Drop.” However, whereas Andor et al. (2016) use discrete features of the parse tree, we can utilize the SUBTREE recurrence function to pull compositional, phrase-based representations of tokens as constructed by the dependency parser. This model is outlined in Figure 2. A full specification is given in the Appendix.
3.2 HOW TO TRAIN A DRAGNN
Given a list of TBRUs, we propose the following learning procedure. We assume training data consists of examples x along with gold decision sequences for one of the TBRUs in the DRAGNN.
1This composition function is similar to that in the constituent parsing SPINN model (Bowman et al., 2016), but with several key differences. Since we use TBRUs, we compose new representations for “Shift” actions as well as reductions, we take inputs from other recurrent models, and we can utilize subtree representations in downstream tasks.
Note that, at a minimum, we need such data for the final TBRU. Assuming given decisions d1 . . . dN from prior components 1 . . . T−1, we define a log-likelihood objective to train the T ’th TBRU along its gold decision sequence d?N+1, . . . , d ? N+n, conditioned on prior decisions:
L(x, d?N+1:N+n; θ) = ∑ i logP (d?N+i | d1:N , d?N+1:N+i−1; θ) (1)
where θ are the combined parameters across all TBRUs. We observe that this objective is locally normalized (Andor et al., 2016), since we optimize the probabilities of the individual decisions in the gold sequence.
The remaining question is where do the decisions d1 . . . dN come from. There are two options here: they can either come as part of the gold annotation (e.g. if we have joint tagging and parsing data), or they will be predicted by unrolling the previous components (e.g. when training stacked extractive summarization model, the parse trees will be predicted by the previously trained parser TBRU).
When training a given TBRU, we unroll an entire input sequence and then use backpropagation through structure (Goller & Kuchler, 1996) to optimize (1). To train the whole system on a set of C datasets, we use a similar strategy to (Dong et al., 2015; Luong et al., 2016); we sample a target task c, 1 ≤ c ≤ C, from a pre-defined ratio, and take a stochastic optimization step on the objective of that task’s TBRU. In practice, task sampling is usually preceded by a deterministic number of pretraining steps, allowing, for example, to schedule a certain number of tagger training steps before running any parser training steps.
4 EXPERIMENTS
In this section, we evaluate three aspects of our approach on two NLP tasks: English dependency parsing and extractive sentence summarization. For English dependency parsing, we primarily use the the Union Treebank setup from Andor et al. (2016). By evaluating on both news and questions domains, we can separately evaluate how the model handles naturally longer and shorter form text. On the Union Treebank setup there are 93 possible actions considering all arc-label combinations. For extractive sentence summarization, we use the dataset of Filippova & Altun (2013), where a large news collection is used to heuristically generate compression instances. The final corpus contains about 2.3M compression instances, but since we evaluated multiple tasks using this data, we subsampled the training set to be comparably sized to the parsing data (≈60K training sentences). The test set contains 160K examples. We implement our method in TensorFlow, using mini-batches of size 4 and following the averaged momentum training and hyperparameter tuning procedure of Weiss et al. (2015).
4.1 USING EXPLICIT STRUCTURE IMPROVES ENCODER/DECODER
We explore the impact of different types of recurrences on dependency parsing in Table 1. In this setup, we used relatively small models: single-layer LSTMs with 256 hidden units, taking
32-dimensional word or output symbol embeddings as input to each cell. In each case, the parsing TBRU takes input from a right-to-left shift-only TBRU. Under these settings, the pure encoder/decoder seq2seq model simply does not have the capacity to parse newswire text with any degree of accuracy, but the TBRU-based approach is nearly state-of-the-art at the same exact computational cost. As a point of comparison and an alternative to using input pointers, we also implemented an attention mechanism within DRAGNN. We used the dot-product formulation from Parikh et al. (2016), where r(si) in the parser takes in all of the shift-only TBRU’s hidden states and RNN aggregates over them.
4.2 UTILIZING PARSE REPRESENTATIONS IMPROVES SUMMARIZATION
We evaluate our approach on the summarization task in Table 2. We compare two single-task LSTM tagging baselines against two multi-task approaches: an adaptation of Luong et al. (2016) and the stack-propagation idea of Zhang & Weiss (2016). In both multi-task setups, we use a right-toleft shift-only TBRU to encode the input, and connect it to both our compositional arc-standard dependency parser and the “Keep/Drop” summarization tagging model.
In both setups we do not follow seq2seq, but utilize the INPUT function to connect output decisions directly to input token representations. However, in the stack-prop case, we use the SUBTREE function to connect the tagging TBRU to the parser TBRU’s phrase representations directly (Figure 2). We find that allowing the compressor to directly use the parser’s phrase representations significantly improves the outcome of the multi-task learning setup. In both setups, we pretrained the parsing model for 400K steps and tuned the subsequent ratio of parser/tagger update steps using a development set.
4.3 DEEP STACKED BI-DIRECTIONAL PARSING
Here we propose a continuous version of the bi-directional parsing model of Attardi & Dell’Orletta (2009): first, the sentence is parsed in the left-to-right order as usual; then a right-to-left transition system analyzes the sentence in reverse order using addition features extracted from the left-to-right parser. In our version, we connect the right-to-left parsing TBRU directly to the phrase representations of the left-to-right parsing TBRU, again using the SUBTREE function. Our parser has the significant advantage that the two directions of parsing can affect each other during training. During each training step the right-to-left parser uses representations obtained using the predictions of the left-to-right parser. Thus, the right-to-left parser can backpropagate error signals through the left-to-right parser and reduce cascading errors caused by the pipeline.
Dev Test Model UAS LAS UAS LAS
Our final model uses 5 TBRU units. Inspired by Zhang & Weiss (2016), a left-to-right POS tagging TBRU provides the first layer of representations. Next, we run two shift-only TBRUs, one in each direction, to provide representations to the parsers. Finally, we connect the left-to-right parser to the right-to-left parser using links defined via the SUBTREE function. The result (Table 3) is a state-ofthe-art dependency parser, yielding the highest published accuracy for a model trained solely on the Penn Treebank with no additional resources.
5 CONCLUSIONS
We presented a compact, modular framework for describing recurrent neural architectures. We evaluated our dynamically structured model and found it to be significantly more efficient and accurate than attention mechanisms for dependency parsing and extractive sentence summarization in both single- and multi-task setups. While we focused primarily on syntactic parsing, the framework provides a general means of sharing representations between tasks. There remains low-hanging fruit still to be explored: in particular, our approach can be globally normalized with multiple hypotheses in the intermediate structure. We also plan to push the limits of multi-task learning by combining many different NLP tasks, such as translation, summarization, tagging problems, and reasoning tasks, into a single model.
ACKNOWLEDGEMENTS
We thank Kuzman Ganchev, Michael Collins, Dipanjan Das, Slav Petrov, Aliaksei Severyn, Chris Dyer, and Noah Smith for their useful feedback and discussion while preparing this draft.

The authors present a general framework for defining a wide variety of recurrent neural network architectures, including seq2seq models, tree-structured models, attention, and a new family of dynamically connected architectures. The framework defines a new, general-purpose recurrent unit called the TBRU, which takes a transition system, defining and constraining its inputs and outputs, and input function which defines the mapping between raw inputs and fixed-width vector representations, and recurrence function that defines the inputs to each recurrent step as a function of the current state, and an RNN cell that computes the output from the input (fixed and recurrent). Many example instantiations of this framework are provided, including sequential tagging RNNs, Google’s Parsey McParseface parser, encoder/decoder networks, tree LSTMs and less familiar examples that demonstrate the power this framework.
The most interesting contribution of this work is the ease by which it can be used to incorporate dynamic recurrent connections through the definition of the transition system. In particular, this paper explores the application of these dynamic connections to syntactic dependency parsing, both as a standalone task, and by multitasking parsing with extractive summarization, using the same compositional phrase representations as features for the parser and summarization (previous work used discrete parse features), which is particularly simple/elegant in this framework. In experimental results, the authors demonstrate that such multitasking leads to more accurate summarization models, and using the framework to incorporate more structure into existing parsing models also leads to increased accuracy with no big-oh efficiency loss (compared with e.g. attention).
The “raison d’etre,” in particular the example, perhaps described even more thoroughly/explicitly, should be made as clear as possible as soon as possible. This is the most important contribution, but it gets lost in the description and presentation as a framework — emphasizing that attention, seq2seq, etc can be represented in the framework is distracting and makes it seem less novel than it is. AnonReviewer6 clearly missed this point, as did I in my first pass over the paper. To get this idea across and to emphasize the benefits of this representation, I’d love to see more detailed analysis of these representations and their importance to achieving your experimental results. I think it would also be helpful to emphasize the difference between a stack LSTM and Example 6.
Overall I think this paper presents a valuable contribution, though the exposition could be improved and analysis of experimental results expanded.


MAXIMUM ENTROPY FLOW NETWORKS
Maximum entropy modeling is a flexible and popular framework for formulating statistical models given partial knowledge. In this paper, rather than the traditional method of optimizing over the continuous density directly, we learn a smooth and invertible transformation that maps a simple distribution to the desired maximum entropy distribution. Doing so is nontrivial in that the objective being maximized (entropy) is a function of the density itself. By exploiting recent developments in normalizing flow networks, we cast the maximum entropy problem into a finite-dimensional constrained optimization, and solve the problem by combining stochastic optimization with the augmented Lagrangian method. Simulation results demonstrate the effectiveness of our method, and applications to finance and computer vision show the flexibility and accuracy of using maximum entropy flow networks.
1 INTRODUCTION
The maximum entropy (ME) principle (Jaynes, 1957) states that subject to some given prior knowledge, typically some given list of moment constraints, the distribution that makes minimal additional assumptions – and is therefore appropriate for a range of applications from hypothesis testing to price forecasting to texture synthesis – is that which has the largest entropy of any distribution obeying those constraints. First introduced in statistical mechanics by Jaynes (1957), and considered both celebrated and controversial, ME has been extensively applied in areas including natural language processing (Berger et al., 1996), ecology (Phillips et al., 2006), finance (Buchen & Kelly, 1996), computer vision (Zhu et al., 1998), and many more.
Continuous ME modeling problems typically include certain expectation constraints, and are usually solved by introducing Lagrange multipliers, which under typical assumptions yields an exponential family distribution (also called Gibbs distribution) with natural parameters such that the expectation constraints are obeyed. Unfortunately, fitting ME distributions in even modest dimensions poses significant challenges. First, optimizing the Lagrangian for a Gibbs distribution requires evaluating the normalizing constant, which is in general computationally very costly and error prone. Secondly, in all but the rarest cases, there is no way to draw samples independently and identically from this Gibbs distribution, even if one could derive it. Third, unlike in the discrete case where a number of recent and exciting works have addressed the problem of estimating entropy from discrete-valued data (Jiao et al., 2015; Valiant & Valiant, 2013), estimating differential entropy from data samples remains inefficient and typically biased. These shortcomings are critical and costly, given the common use of ME distributions for generating reference data samples for a null distribution of a test statistic. There is thus ample need for a method that can both solve the ME problem and produce a solution that is easy and fast to sample.
In this paper we develop maximum entropy flow networks (MEFN), a stochastic-optimization-based framework and algorithm for fitting continuous maximum entropy models. Two key steps are required. First, conceptually, we replace the idea of maximizing entropy over a density directly with maximizing, over the parameter space of an indexed function family, the entropy of the density induced by mapping a simple distribution (a Gaussian) through that optimized function. Modern
∗These authors contributed equally.
neural networks, particularly in variational inference (Kingma & Welling, 2013; Rezende & Mohamed, 2015), have successfully employed this same idea to generate complex distributions, and we look to similar technologies. Secondly, unlike most other objectives in this network literature, the entropy objective itself requires evaluation of the target density directly, which is unavailable in most traditional architectures. We overcome this potential issue by learning a smooth, invertible transformation that maps a simple distribution to an (approximate) ME distribution. Recent developments in normalizing flows (Rezende & Mohamed, 2015; Dinh et al., 2016) allow us to avoid biased and computationally inefficient estimators of differential entropy (such as the nearest-neighbor class of estimators like that of Kozachenko-Leonenko; see Berrett et al. (2016)). Our approach avoids calculation of normalizing constants by learning a map with an easy-to-compute Jacobian, yielding tractable probability density computation. The resulting transformation also allows us to reliably generate iid samples from the learned ME distribution. We demonstrate MEFN in detail in examples where we can access ground truth, and then we demonstrate further the ability of MEFN networks in equity option prices fitting and texture synthesis.
Primary contributions of this work include: (i) addressing the substantial need for methods to sample ME distributions; (ii) introducing ME problems, and the value of including entropy in a range of generative modeling problems, to the deep learning community; (iii) the novel use of constrained optimization for a deep learning application; and (iv) the application of MEFN to option pricing and texture synthesis, where in the latter we show significant increase in the diversity of synthesized textures (over current state of the art) by using MEFN.
2 BACKGROUND
2.1 MAXIMUM ENTROPY MODELING AND GIBBS DISTRIBUTION
We consider a continuous random variable Z ∈ Z ⊆ Rd with density p, where p has differential entropy H(p) = − ∫ p(z) log p(z)dz and support supp(p). The goal of ME modeling is to find, and then be able to easily sample from, the maximum entropy distribution given a set of moment and support constraints, namely the solution to:
p∗ = maximize H(p) (1) subject to EZ∼p[T (Z)] = 0
supp(p) = Z,
where T (z) = (T1(z), ..., Tm(z)) : Z → Rm is the vector of known (assumed sufficient) statistics, andZ is the given support of the distribution. Under standard regularity conditions, the optimization problem can be solved by Lagrange multipliers, yielding an exponential family p∗ of the form:
p∗(z) ∝ eη >T (z) 1(z ∈ Z) (2)
where η ∈ Rm is the choice of natural parameters of p∗ such that Ep∗ [T (Z)] = 0. Despite this simple form, these distributions are only in rare cases tractable from the standpoint of calculating η, calculating the normalizing constant of p∗, and sampling from the resulting distribution. There is extensive literature on finding η numerically (Darroch & Ratcliff, 1972; Salakhutdinov et al., 2002; Della Pietra et al., 1997; Dudik et al., 2004; Malouf, 2002; Collins et al., 2002), but doing so requires computing normalizing constants, which poses a challenge even for problems with modest dimensions. Also, even if η is correctly found, it is still not trivial to sample from p∗. Problemspecific sampling methods (such as importance sampling, MCMC, etc.) have to be designed and used, which is in general challenging (burn-in, mixing time, etc.) and computationally burdensome.
2.2 NORMALIZING FLOWS
Following Rezende & Mohamed (2015), we define a normalizing flow as the transformation of a probability density through a sequence of invertible mappings. Normalizing flows provide an elegant way of generating a complicated distribution while maintaining tractable density evaluation. Starting with a simple distribution Z0 ∈ Rd ∼ p0 (usually taken to be a standard multivariate
Gaussian), and by applying k invertible and smooth functions fi : Rd → Rd(i = 1, ..., k), the resulting variable Zk = fk ◦ fk−1 ◦ · · · ◦ f1(Z0) has density:
pk(zk) = p0(f −1 1 ◦ f −1 2 ◦ · · · ◦ f −1 k (zk)) k∏ i=1 |det(Ji(zi−1))|−1, (3)
where Ji is the Jacobian of fi. If the determinant of Ji can be easily computed, pk can be computed efficiently.
Rezende & Mohamed (2015) proposed two specific families of transformations for variational inference, namely planar flows and radial flows, respectively:
fi(z) = z + uih(w T i z + bi) and fi(z) = z + βih(αi, ri)(z− z′i), (4)
where bi ∈ R, ui,wi ∈ Rd and h is an activation function in the planar case, and where βi ∈ R, αi > 0, z′i ∈ Rd , h(α, r) = 1/(α + r) and ri = ||z − z′i|| in the radial. Recently Dinh et al. (2016) proposed a normalizing flow with convolutional, multiscale structure that is suitable for image modeling and has shown promise in density estimation for natural images.
3 MAXIMUM ENTROPY FLOW NETWORK (MEFN) ALGORITHM
3.1 FORMULATION
Instead of solving Equation 2, we propose solving Equation 1 directly by optimizing a transformation that maps a random variable Z0, with simple distribution p0, to the ME distribution. Given a parametric family of normalizing flows F = {fφ, φ ∈ Rq}, we denote pφ(z) = p0(f −1 φ (z))|det(Jφ(z))|−1 as the distribution of the variable fφ(Z0), where Jφ is the Jacobian of fφ. We then rewrite the ME problem as:
φ∗ = maximize H(pφ) (5) subject to EZ0∼p0 [T (fφ(Z0))] = 0
supp(pφ) = Z.
When p0 is continuous and F is suitably general, the program in Equation 5 recovers the ME distribution pφ exactly. With a flexible transformation family, the ME distribution can be well approximated. In experiments we found that taking p0 to be a standard multivariate normal distribution achieves good empirical performance. Taking p0 to be a bounded distribution (e.g. uniform distribution) is problematic for learning transformations near the boundary, and heavy tailed distributions (e.g. Cauchy distribution) caused similar trouble due to large numbers of outliers.
3.2 ALGORITHM
We solved Equation 5 using the augmented Lagrangian method. Denote R(φ) = E(T (fφ(Z0))), the augmented Lagrangian method uses the following objective:
L(φ;λ, c) = −H(pφ) + λ>R(φ) + c
2 ||R(φ)||2 (6)
where λ ∈ Rm is the Lagrange multiplier and c > 0 is the penalty coefficient. We minimize Equation 6 for a non-decreasing sequence of c and well-chosen λ. As a technical note, the augmented Lagrangian method is guaranteed to converge under some regularity conditions (Bertsekas, 2014). As is usual in neural networks, a proof of these conditions is challenging and not yet available, though intuitive arguments (see Appendix §A) suggest that most of them should hold. Due to the non rigorous nature of these arguments, we rely on the empirical results of the algorithm to claim that it is indeed solving the optimization problem.
For a fixed (λ, c) pair, we optimize L with stochastic gradient descent. Owing to our choice of network and the resulting ability to efficiently calculate the density pφ(z(i)) for any sample point
Algorithm 1 Training the MEFN 1: initialize φ = φ0, set c0 > 0 and λ0. 2: for Augmented Lagrangian iteration k = 1, ..., kmax do 3: for SGD iteration i = 1, ..., imax do 4: Sample z(1), ..., z(n) ∼ p0, get transformed variables z(i)φ = fφ(z(i)), i = 1, ..., n 5: Update φ by descending its stochastic gradient (using e.g. ADADELTA (Zeiler, 2012)):
∇φL(φ;λk, ck) ≈ 1
n n∑ i=1 ∇φ log pφ(z(i)φ ) + 1 n n∑ i=1 ∇φT (z(i)φ )λk + ck 2 n n 2∑ i=1 ∇φT (z(i)φ ) · 2 n n∑ i=n
2 +1
T (z (i) φ )
6: end for 7: Sample z(1), ..., z(ñ) ∼ p0, get transformed variables z(i)φ = fφ(z(i)), i = 1, ..., ñ 8: Update λk+1 = λk + ck 1ñ ∑ñ i=1 T (z (i) φ )
9: Update ck+1 ≥ ck (see text for detail) 10: end for
z(i) (which are easy-to-sample iid draws from the multivariate normal p0), we compute the unbiased estimator of H(pφ) with:
H(pφ) ≈ − 1
n n∑ i=1 log pφ(fφ(z (i))) (7)
R(φ) can also be estimated without bias by taking a sample average of z(i) draws. The resulting optimization procedure is detailed in Algorithm 1, of which step 9 requires some detail: denoting φk as the resulting φ after imax SGD iterations at the augmented Lagrangian iteration k, the usual update rule for c (Bertsekas, 2014) is:
ck+1 = { βck, if ||R(φk+1)|| > γ||R(φk)|| ck, otherwise
(8)
where γ ∈ (0, 1) and β > 1. Monte Carlo estimation of R(φ) sometimes caused c to be updated too fast, causing numerical issues. Accordingly, we changed the hard update rule for c to a probabilistic update rule: a hypothesis test is carried out with null hypothesis H0 : E[||R(φk+1)||] = E[γ||R(φk)||] and alternative hypothesis H1 : E[||R(φk+1)||] > E[γ||R(φk)||]. The p-value p is computed, and ck+1 is updated to βck with probability 1 − p. We used a two-sample t-test to calculate the p-value. What results is a robust and novel algorithm for estimating maximum entropy distributions, while preserving the critical properties of being both easy to calculate densities of particular points, and being trivially able to produce truly iid samples.
4 EXPERIMENTS
We first construct an ME problem with a known solution (§4.1), and we analyze the MEFN algorithm with respect to the ground truth and to an approximate Gibbs solution. These examples test the validity of our algorithm and illustrate its performance. §B and §4.3 then applies the MEFN to a financial data application (predicting equity option values) and texture synthesis, respectively, to illustrate the flexibility and practicality of our algorithm.
For §4.1 and §B, We use 10 layers of planar flow with a final transformation g (specified below) that transforms samples to the specified support, and use with ADADELTA (Zeiler, 2012). For §4.3 we use real NVP structure and use ADAM (Kingma & Ba, 2014) with learning rate = 0.001. For all our experiments, we use imax = 3000, β = 4, γ = 0.25. For §4.1 and §B we use n = 300, ñ = 1000, kmax = 10; For §4.3 we use n = ñ = 2, kmax = 8.
4.1 A MAXIMUM ENTROPY PROBLEM WITH KNOWN SOLUTION
Following the setup of the typical ME problem, suppose we are given a specified support S = {z = (z1, . . . , zd−1) : zi ≥ 0 and ∑d−1 k=1 zk ≤ 1} and a set of constraints E[logZk] = κk(k = 1, ..., d),
where Zd = 1− ∑d−1 k=1 Zk. We then write the maximum entropy program:
p∗ = maximize H(p) (9) subject to EZ∼p[logZk − κk] = 0 ∀k = 1, ..., d
supp(p) = S.
This is a general ME problem that can be solved via the MEFN. Of course, we have particularly chosen this example because, though it may not obviously appear so, the solution has a standard and tractable form, namely the Dirichlet. This choice allows us to consider a complicated optimization program that happens to have known global optimum, providing a solid test bed for the MEFN (and for the Gibbs approach against which we will compare). Specifically, given a parameter α ∈ Rd, the Dirichlet has density:
p(z1, . . . , zd−1) = 1
B(α) d∏ k=1 zαk−1k 1 ((z1, . . . , zd−1) ∈ S) (10)
where B(α) is the multivariate Beta function, and zd = 1 − ∑d−1 k=1 zk. Note that this Dirichlet is a distribution on S and not on the (d − 1)-dimensional simplex Sd−1 = {(z1, . . . , zd) : zk ≥ 0 and ∑d k=1 zk = 1} (an often ignored and seemingly unimportant technicality that needs to be correct here to ensure the proper transformation of measure). Connecting this familiar distribution to the ME problem above, we simply have to choose α such that κk = ψ(αk)−ψ(α0) for k = 1, ..., d, where α0 = ∑d k=1 αk and ψ is the digamma function. We then can pose the above ME problem to the MEFN and compare performance against ground truth. Before doing so, we must stipulate the transformation g that maps the Euclidean space of the multivariate normal p0 to the desired support S. Any sensible choice will work well (another point of flexibility for the MEFN); we use the standard transformation:
g(z1, ..., zd−1) = ( ez1∑d−1
k=1 e zk + 1
, ..., ezd−1∑d−1
k=1 e zk + 1
)> (11)
Note that the MEFN outputs vectors in Rd−1, and not Rd, because the Dirichlet is specified as a distribution on S (and not on the simplex Sd−1). Accordingly, the Jacobian is a square matrix and its determinant can be computed efficiently using the matrix determinant lemma. Here, p0 is set to the (d− 1)-dimensional standard normal.
We proceed as follows: We choose α and compute the constraints κ1, ..., κd. We run MEFN pretending we do not know α or the Dirichlet form. We then take a random sample from the fitted distribution and a random sample from the Dirichlet with parameter α, and compare the two samples using the maximum mean discrepancy (MMD) kernel two sample test (Gretton et al., 2012), which assesses the fit quality. We take the sample size to be 300 for the two sample kernel test. Figure 1 shows an example of the transformation from normal (left panel) to MEFN (middle panel), and comparing that to the ground truth Dirichlet (right panel). The MEFN and ground truth Dirichlet densities shown in purple match closely, and the samples drawn (red) indeed appear to be iid draws from the same (maximum entropy) distribution in both cases.
Additionally, the middle panel of Figure 1 shows an important cautionary tale that foreshadows our texture synthesis results (§4.3). One might suppose that satisfying the moment matching constraints is adequate to produce a distribution which, if not technically the ME distribution, is still interestingly variable. The middle panel shows the failure of this intuition: in dark green, we show a network trained to simply match the moments specified above, and the resulting distribution quite poorly expresses the variability available to a distribution with these constraints, leading to samples that are needlessly similar. Given the substantial interest in using networks to learn implicit generative models (e.g., Mohamed & Lakshminarayanan (2016)), this concern is particularly relevant and highlights the importance of considering entropy.
Figure 2 quantitatively analyzes these results. In the left panel, for a specific choice of α = (1, 2, 3), we show our unbiased entropy estimate of the MEFN distribution pφ as a function of the number of SGD iterations (red), along with the ground truth maximum entropy H(p∗) (green line). Note
that the MEFN stabilizes at the correct value (as a stochastic estimator, variance around that value is expected). In the middle panel, we show the distribution of MMD values for the kernel two sample test, as well as the observed statistic for the MEFN (red) and for a randomly chosen Dirichlet distribution (gray; chosen to be close to the true optimum, making a conservative comparison). The MMD test does not reject MEFN as being different from the true ME distribution p∗, but it does reject a Dirichlet whose KL to the true p∗ is small (see legend). In the right panel, for many different Dirichlets in a small grid around a single true p∗, the kernel two sample test statistic is computed, the MMD p-value is calculated, as is the KL to the true distribution. We plot a scatter of these points in grey, and we plot the particular MEFN solution as a red star. We see that for other Dirichlets with similar KL to the true distribution as the MEFN distribution, the p-values seem uniform, meaning that the KL to the true is indeed very small. Again this is conservative, as the grey points have access to the known Dirichlet form, whereas the MEFN considered the entire space (within its network capacity) of S supported distributions. Given this fact, the performance of MEFN is impressive.
4.2 RISK-NEUTRAL ASSET PRICING
We illustrate the flexibility and practicality of our algorithm extracting the risk-neutral asset price probability based on option prices, an active and interesting area for ME models. We find that MEFN and the classic Gibbs approach yield comparable performances. Owing to space limitations we have placed these results in Appendix §B.
4.3 MODELING IMAGES OF TEXTURES
Constructing generative models to generate random images with certain texture structure is an important task in computer vision. A line of texture synthesis research proceeds by first extracting a set
of features that characterizes the target texture and then generate images that match the features. The seminal work of Zhu et al. (1998) proposes constructing texture models under the ME framework, where features (or filters) of the given texture image are adaptively added in the model and a Gibbs distribution whose expected feature matches the target texture is learnt. One major difficulty with the method is that both model learning and image generation involve sampling from a complicated Gibbs distribution. More recent works exploit more complicated features (Portilla & Simoncelli, 2000; Gatys et al., 2015; Ulyanov et al., 2016). Ulyanov et al. (2016) propose the texture net, which uses a texture loss function by using the Gram matrices of the outputs of some convolutional layers of a pre-trained deep neural network for object recognition.
While the use of these complicated features does provide high-quality synthetic texture images, that work focuses exclusively on generating images that match these feature (moments). Importantly, this network focuses only on generating feature-matching images without using the ME framework to promote the diversity of the samples. Doing so can be deeply problematic: in Figure 1 (middle panel), we showed the lack of diversity resulting from only moment matching in that Dirichlet setting, and further we note that the extreme pathology would result in a point mass on the training image – a global optimum for this objective, but obviously a terrible generative model for synthesizing textures. Ideally, the MEFN will match the moments and promote sample diversity.
We applied MEFN to texture synthesis with an RGB representation of the 224 × 224 pixel images , z ∈ Z = [0, 1]d, where d = 224 × 224 × 3. We follow Ulyanov et al. (2016) (we adapted https://github.com/ProofByConstruction/texture-networks) to create a texture loss measure T (z) : [0, 1]d → R, and aim to sample a diverse set of images with small moment violation. For the transformation family F we use the real NVP network structure proposed in Dinh et al. (2016) (we adapted https://github.com/taesung89/real-nvp). We use 3 residual blocks with 32 feature maps for each coupling layer and downscale 3 times. For fair comparison, we use the same real NVP structure for both1, implemented in TensorFlow (Abadi et al., 2016).
As is shown in top row of figure 3, both methods generate visually pleasing images capturing the texture structure well. The bottom row of Figure 3 shows that texture cost (left panel) is similar for both methods, while MEFN generates figures with much larger entropy than the texture network formulation (middle panel), which is desirable (as previously discussed). The bottom right panel of figure 3 compares the marginal distribution of the RGB values sampled from the networks: we found that MEFN generates a more variable distribution of RGB values than the texture net. Further results are in Appendix §C.
1Ulyanov et al. (2016) use a quite different generative network structure, which is not invertible and is therefore infeasible for entropy evaluation, so we replace their generative network by the real NVP structure.
We compute in Table 1 the average pairwise Euclidean distance between randomly sampled images (dL2 = meani 6=j‖zi − zj‖22), and MEFN gives higher dL2 , quantifying diversity across images. We also consider an ANOVA-style analysis to measure the diversity of the images, where we think of the RGB values for the same pixel across multiple images as a group, and compute the within and between group variance. Specifically, denoting zki as the pixel value for a specific pixel k = 1, ..., d for an image i = 1, ...., n. We partition the total sum of square SST = ∑ i,k(z k i − z̄)2 as the within
group error SSW = ∑ i,k(z k i − z̄k)2 and between group error SSB = ∑ i,k n(z̄
k − z̄)2, where z̄ and z̄k are the mean pixel values across all data and for a specific pixel k. Ideally we want the samples to exhibit large variability across images (large SSW, within a group/pixel) and no structure in the mean image (small SSB, across groups/pixels). Indeed, the MEFN has a larger SSW, implying higher variability around the mean image, a smaller SSB, implying the stationarity of the generated samples, and a larger SST, implying larger total variability also. The MEFN produces images that are conclusively more variable without sacrificing the quality of the texture, implicating the broad utility of ME.
5 CONCLUSION
In this paper we propose a general framework for fitting ME models. This approach is novel and has three key features. First, by learning a transformation of a simple distribution rather than the distribution itself, we are able to avoid explicitly computing an intractable normalizing constant for the ME distribution. Second, by combining stochastic optimization with the augmented Lagrangian method, we can fit the model efficiently, allowing us to evaluate the ME density of any point simply and accurately. Third, critically, this construction allows us to trivially sample iid from a ME distribution, extending the utility and efficiency of the ME framework more generally. Also, accuracy equivalent to the classic Gibbs approach is in itself a contribution (owing to these other features). We illustrate the MEFN in both a simulated case with known ground truth and real data examples.
There are a few recent works encouraging sample diversity in the setting of texture modeling. Ulyanov et al. (2017) extended Ulyanov et al. (2016) by adding a penalty term using the Kozachenko-Leonenko estimator Kozachenko & Leonenko (1987) of entropy. Their generative network is an arbitrary deep neural network rather than a normalizing flow, which is more flexible but cannot give the probability density of each sample easily so as to compute an unbiased estimator of the entropy. Kozachenko-Leonenko is a biased estimator for entropy and requires a fairly large number of samples to get good performance in high-dimensional settings, hindering the scalability and accuracy of the method; indeed, our choice of normalizing flow networks was driven by these practical issues with Kozachenko-Leonenko. Lu et al. (2016) extended Zhu et al. (1998) by using a more flexible set of filters derived from a pre-trained deep neural networks, and using parallel MCMC chains to learn and sample from the Gibbs distribution. Running parallel MCMC chains results in diverse samples but can be computationally intensive for generating each new sample image. Our MEFN framework enables truly iid sampling with the ease of a feed forward network.
ACKNOWLEDGMENTS
We thank Evan Archer for normalizing flow code, and Xuexin Wei, Christian Andersson Naesseth and Scott Linderman for helpful discussion. This work was supported by a Sloan Fellowship and a McKnight Fellowship (JPC).
A AUGMENTED LAGRANGIAN CONDITIONS
We give a more thorough discussion of the regularity conditions which ensure that the Augmented Lagrangian method will work. The goal of this section is simply to state these conditions and give intuitive arguments about why some should hold in our case, not to attempt to prove that they indeed hold. The conditions (Bertsekas, 2014) are:
• There exists a strict local minimum φ∗ of the optimization problem of Equation 5: If the function class F is rich enough that it contains a true solver of the maximum entropy problem, then a global optimum exists. Although not rigorous, we would expect that even in the finite expressivity case that a global optimum remains, and indeed, recent theoretical work (Raghu et al., 2016; Poole et al., 2016) has gotten close to proving this.
• φ∗ is a regular point of the optimization problem, that is, the rows of∇φR(φ∗) are linearly independent: Again, this is not formal, but we should not expect this to cause any issues. This clearly depends on the specific form of T , but the condition basically says that there should not be redundant constraints at the optimum, so if T is reasonable this shouldn’t happen.
• H(pφ) and R(φ) are twice continuously differentiable on a neighborhood around φ∗: This holds by the smoothness of the normalizing flows.
• y>∇2φL(φ∗;λ∗, 0)y > 0 for every y 6= 0 such that ∇φR(φ∗)y = 0, where λ∗ is the true Lagrange multiplier: This condition is harder to justify. It would appear it is just asking that the Lagrangian (not the augmented Lagrangian) be strictly convex in feasible directions, but it is actually stronger than this and some simple functions might not satisfy the property. For example, if the function we are optimizing was x4 and we had no constraints, the Lagrangian’s Hessian would be 12x2, which is 0 at x∗ = 0 thus not satisfying the condition. Importantly, these conditions are sufficient but not necessary, so even if this doesn’t hold the augmented Lagrangian method might work (it certainly would for x4). Because of this and the nonrigorous justifications of the first two conditions, we left these conditions for the appendix and relied instead on the empirical performance to justify that we are indeed recovering the maximum entropy distribution.
If all of these conditions hold, the augmented Lagrangian (for large enough c and λ close enough to λ∗) has a unique optimum in a neighborhood around φ∗ that is close to φ∗ (as λ→ λ∗ it converges to φ∗) and its hessian at this optimum is positive-definite. Furthermore, λk → λ. This implies that gradient descent (with the usual caveats of being started close enough to the solution and with the right steps) will correctly recover φ∗ using the augmented Lagrangian method. This of course just guarantees convergence to a local optimum, but if there are no additional assumptions such as convexity, it can be very hard to ensure that it is indeed a global optimum. Some recent research has attempted to explain why optimization algorithms perform so well for neural networks (Choromanska et al., 2015; Kawaguchi, 2016), but we leave such attempts for our case for future research.
B RISK-NEUTRAL ASSET PRICE
We extract the risk-neutral asset price probability distribution based on option prices, an active and interesting area for ME models. We give a brief introduction of the problem and refer interested readers to see Buchen & Kelly (1996) for a more detailed explanation. Denoting St as the price of an asset at time t, the buyer of a European call option for the stock that expires at time te with strike price K will receive a payoff of cK = (Ste − K)+ = max(Ste − K, 0) at time te. Under the efficient market assumption, the risk-neutral probability distribution for the stock price at time te satisfies: cK = D(te)Eq[(Ste −K)+], (12) where D(te) is the risk-free discount factor and q is the risk-neutral measure. We also have that, under the risk-neutral measure, the current stock price S0 is the discounted expected value of Ste :
S0 = D(te)Eq(Ste). (13)
When givenm options that expire at time te with strikesK1, ...,Km and prices cK1 , ..., cKm , we get m expectation constraints on q(Ste) from Equation 12, together with Equation 13, we have m + 1 expectation constraints in total. With that partial knowledge we can approximate q(Ste), which is helpful for understanding the market expected volatility and identify mispricing in option markets, etc.
Inferring the risk-neutral density of asset price from a finite number of option prices is an important question in finance and has been studied extensively (Buchen & Kelly, 1996; Borwein et al., 2003; Bondarenko, 2003; Figlewski, 2008). One popular method proposed by Buchen & Kelly (1996) estimates the probability density as the maximum entropy distribution satisfying the expectation constraints and a positivity support constraint by fitting a Gibbs distribution, which results in a piece-wise linear log density:
p(z) ∝ exp { η0z +
m∑ i=1 ηi(z −Ki)+
} 1 (z ≥ 0) (14)
and optimize the distribution with numerical methods. Here we compare the performance of the MEFN algorithm with the method proposed in Buchen & Kelly (1996). To enforce the positivity constraint we choose g(z) = eaz+b, where a and b are additional parameters.
We collect the closing price of European call options on Nov. 1 2016 for the stock AAPL (Apple inc.) that expires on te = Jun. 16 2017. We usem = 4 of the options with highest trading volume as training data and the rest as testing data. On the left panel of figure 4, we show the fitted risk-neutral density of Ste by MEFN (red line) with that of the fitted Gibbs distribution result (blue line). We find that while the distributions share similar location and variability, the distribution inferred by MEFN is smoother and arguably more plausible. In the middle panel we show a Q-Q plot of the quantiles of the MEFN and Gibbs distributions. We can see that the quantile pairs match the identity closely, which should happen if both methods recovered the exact same distribution. This highlights the effectiveness of MEFN. There does exist a small mismatch in the tails: the distribution inferred by MEFN has slightly heavier tails. This mismatch is difficult to interpret: given that both the Gibbs and MEFN distributions are fit with option price data (and given that one can observe at most one value from the distribution, namely the stock price at expiration), it is fundamentally unclear which distribution is superior, in the sense of better capturing the true ME distribution’s tails. On the right panel we show the fitted option price for the two fitted distributions (for each strike price, we can recover the fitted option price by Equation 12). We noted that the fitted option price and strike price lines for both methods are very similar (they are mostly indiscernible on the right panel of figure 4). We also compare the fitted performance on the test data by computing the root mean square error for the fitted and test data. We observe that the predictive performances for both methods are comparable.
We note that for this specific application, there are practical concerns such as the microstructure noise in the data and inefficiency in the market, etc. Applying a pre-processing procedure and incorporating prior assumptions can be helpful for getting a more full-fledged method (see e.g. Figlewski (2008)). Here we mainly focus on illustrating the ability of the MEFN method to approximate the ME distribution for non-typical distributions. Future work for this application includes fitting a riskneutral distribution for multi-dimensional assets by incorporating dependence structure on assets.
C MODELING IMAGES OF TEXTURES
We tried our texture modeling approach with many different textures, and although MEFN samples don’t always exhibit more visual diversity than samples obtained from the texture network, they always have more entropy as in figure 3. Figure 5 shows two positive examples, i.e. textures in which samples from MEFN do exhibit higher visual diversity than those from the texture network, as well as a negative example, in which MEFN achieves less visual diversity than the texture network, regardless of the fact that MEFN samples do have larger entropy. We hypothesize that this curious behavior is due to the optimization achieving a local optimum in which the brick boundaries and dark brick locations are not diverse but the entropy within each brick is large. It should also be noted that among the experiments that we ran, this was the only negative example that we got, and that slightly modifying the hyperparameters caused the issue to disappear.

Much existing deep learning literature focuses on likelihood based models. However maximum entropy approaches are an equally valid modelling scenario, where information is given in terms of constraints rather than data. That there is limited work in flexible maximum entropy neural models is surprising, but  is due to the fact that optimizing a maximum entropy model requires (a) establishing the effect of the constraints on some distribution, and formulating the entropy of that complex distribution. There is no unbiased estimator of entropy from samples alone, and so an explicit model for the density is needed. This challenge limits approaches. The authors have identified that invertible neural models provide a powerful class of models for solving the maximum entropy network problem, and this paper goes on to establish this approach. The contributions of this paper are (a) recognising that, because normalising flows provide an explicit model for the density, they can be used to provide unbiased estimators for the entropy (b) that the resulting Lagrangian can be implemented as a relaxation of a augmented Lagrangian (c) establishing the practical issues in doing the augmented Lagrangian optimization. As far as the reviewer is aware this work is novel – this approach is natural and sensible, and is demonstrated on an number of models where clear evaluation can be done. Enough experiments have been done to establish this is an appropriate method, though not that it is entirely necessary – it would be good to have an example where the benefits of the flexible flow transformation were much clearer. Further discussion of the computational and scaling aspects would be valuable. I am guessing this approach is probably appropriate for model learning, but less appropriate for inferential settings where a known model is then conditioned on particular instance based constraints? Some discussion of appropriate use cases would be good. The issue of match to the theory via the regularity conditions has been brought up, but it is clear that this can be described well, and exceeds most of the theoretical discussions that occur regarding the numerical methods in other papers in this field.
Quality: Good sound paper providing a novel basis for flexible maximum entropy models.
Clarity: Good.
Originality: Refreshing.
Significance: Significant in model development terms. Whether it will be an oft-used method is not clear at this stage.
Minor issues
Please label all equations. Others might wish to refer to them even if you don’t.
Top of page 4: algorithm 1→ Algorithm 1.
The update for c to overcome stability appears slightly opaque and is mildly worrying.  I assume there are still residual stability issues? Can you comment on why this solves all the problems?
The issue of the support of p is glossed over a little. Is the support in 5 an additional condition on the support of p? If so, that seems hard to encode, and indeed does not turn up in (6). I guess for a Gaussian p0 and invertible unbounded transformations, if the support happens to be R^d, then this is trivial, but for more general settings this seems to be an issue that you have not dealt? Indeed in your Dirichlet example, you explicitly map to the required support, but for more complex constraints this may be non trivial to do with invertible models with known Jacobian? It would be nice to include this in the more general treatment rather than just relegating it to the specific example.
Overall I am very pleased to see someone tackling this question with a very natural approach.


We present an approach for learning state representations in multi-task reinforcement learning. Our method learns multiple low-dimensional state representations from raw observations in an unsupervised fashion, without any knowledge of which task is executed, nor of the number of tasks involved. The method is based on a gated neural network architecture, trained with an extension of the learning with robotic priors objective. In simulated experiments, we show that our method is able to learn better state representations for reinforcement learning, and we analyze why and when it manages to do so.
1 INTRODUCTION
(a) (b)
Figure 1: Slot car racing – the agent has learn how to drive any of the cars as far as possible (left), based on its raw observations (right).
In many reinforcement learning problems, the agent has to solve a variety of different tasks to fulfill its overall goal. A common approach to this problem is to learn a single policy for the whole problem, and leave the decomposition of the problem into subtasks to the learner. In many cases, this approach is successful (Mnih et al., 2015; Zahavy et al., 2016), but it comes at the expense of requiring large amounts of training data. Alternatively, multiple policies dedicated to different subtasks can be learned. This, however, requires prior knowledge about how the overal problem decomposes into subtasks. More-
over, it can run into the same issue of requiring large amounts of data, because the subtasks might overlap and thus afford shared computation to solve them.
A common approach to address overlapping problems is multi-task learning (Caruana, 1997): by learning a single policy with different subgoals, knowledge between the different tasks can be transferred. This not only allows to learn a compact representation more efficiently, but also improves the agent’s performance on all the individual subtasks (Rusu et al., 2016).
Multi-task learning, however, faces two problems: it requires the decomposition of the overall problem into subtasks to be given. Moreover, it is not applicable if the subtasks are unrelated, and are better solved without sharing computation. In this case, the single-policy approach results in an agent that does not perform well on any of the individual tasks (Stulp et al., 2014) or that unlearns
1The first two authors contributed equally to this work.
the successful strategy for one subtasks once it switches to another one, an issue known as catastrophic forgetting (McCloskey & Cohen, 1989).
In this work, we address the problem of identifying and isolating individual unrelated subtasks, and learning multiple separate policies in an unsupervised way. To that end, we present MT-LRP, an algorithm for learning state representations for multiple tasks by learning with robotic priors. MT-LRP is able to acquire different low-dimensional state representations for multiple tasks in an unsupervised fashion. Importantly, MT-LRP does not require knowledge about which task is executed at a given time or about the number of tasks involved. The representations learned with MT-LRP enable the use of standard reinforcement learning methods to compute effective policies from few data.
As explained before, our approach is orthogonal to the classical multi-task learning approach, and constitutes a problem of its own right due to the issues of underperformance and catastrophic forgetting. Therefore, we disregard the shared knowledge problem in this paper. However, any complete reinforcement learning system will need to combine both flavors of multi-task learning, for related and unrelated tasks, and future work will have to address the two problems together.
MT-LRP is implemented as two neural networks, coupled by a gating mechanism (Sigaud et al., 2015; Droniou et al., 2015) as illustrated in Figure 2. The first network, χ , detects which task is being executed and selects the corresponding state representation. The second network, ϕ , learns task-specific state representations. The networks are trained simultaneously using the robotic priors learning objective (Jonschkowski & Brock, 2015), exploiting physics-based prior knowledge about how states, actions, and rewards relate to each other. Both networks learn from raw sensor data, without supervision and solely based on the robot’s experiences.
In a simulated experimental scenario, we show that MT-LRP is able to learn multiple state representations and task detectors from raw observations and that these representations allow to learn better policies from fewer data when compared with other methods. Moreover, we analyze the contribution to this result of each the method’s individual components.
2 RELATED WORK
MT-LRP combines three ideas into a novel approach for task discovery and state representation learning: 1) extracting state representations for each task with robotic priors (Jonschkowski & Brock, 2015); 2) discovering discrete tasks and corresponding actions/policies in a RL context (Stulp et al., 2014; Höfer & Brock, 2016); 3) using gated networks to implement a “mixture of experts” (Jacobs et al., 1991; Droniou et al., 2015).
State Representation Learning: Learning from raw observations is considered a holy grail in reinforcement learning (RL). Deep RL has had major success in this, using model-free (Mnih et al., 2015) but also by combining model-free and model-based RL (Levine et al., 2015). These approaches apply end-to-end learning to get from raw input to value functions and policies. A different approach is to explicitly learn state representations using unsupervised learning, e.g. using auto-encoders (Lange et al., 2012). Recently, Watter et al. (2015) extended this idea to learn state representations jointly with dynamic models and apply optimal control to compute a policy. We use learning with robotic priors (Jonschkowski & Brock, 2015), a state representation learning method
that exploits information about temporal structure, actions, and rewards. We go beyond previous work by not only learning single state representations, but learning multiple state representations given raw data from multiple tasks.
Options and Parameterized Skills: A common approach to factorizing a RL problem into subtasks are macro-actions, often called options (Sutton et al., 1999; Hengst, 2002). The main difference with our approach is that options are used to hierarchically decompose one high-level task into subtasks (and learn sub-policies for these subtasks), whereas we learn task-specific state representations for different high-level tasks. However, options bear resemblance on a technical level, since they are often implemented by a high-level “selection” policy that parametrizes low-level policies (Daniel et al., 2012; Kupcsik et al., 2013; Stulp et al., 2014). Continuous versions of options, referred to as parametrized skills, have been proposed, too (Da Silva et al., 2012; Deisenroth et al., 2014; DoshiVelez & Konidaris, 2016). However, in all the work above, the state representation is given. To the best of our knowledge, state representation learning has not yet been considered in the context of RL with options or parameterized skills.
Gated Networks for Mixtures of Experts and Submanifold Learning: Gated networks are networks that contain gating connections, in which the outputs of at least two neurons are multiplied (Sigaud et al., 2015). This allows a gating neuron g to prohibit (or limit) the flow of information from one neuron x to another neuron y, similar to how transistors function. An early example of gated networks is the mixture of experts approach (Jacobs et al., 1991; Jacobs & Jordan, 1993; Haruno et al., 2001), where separate networks in a modular neural network specialize in predicting subsets of training examples from a database. Our contribution is to extend mixtures of experts by state representation learning (e.g. from raw images) and to the more difficult RL (rather than supervised learning) context. Our gated network architecture is similar to the one proposed by Droniou et al. (2015). Their network simultaneously learns discrete classes jointly with continuous class variations (called submanifolds) in an unsupervised way, e.g., discrete digit classes and shape variations within each class. We use a similar architecture, but in a different way: rather than learning discrete classes, we learn discrete tasks; class-specific submanifolds correspond to task-specific state representations; and finally, we consider a RL rather than an unsupervised learning context.
As mentioned in the introduction, our work is orthogonal to multi-task learning (Caruana, 1997) which has been extensively studied in recent reinforcement learning literature, too (Parisotto et al., 2016). Our approach can be trivially combined with multi-task learning by by prepending the gate and state extraction modules with a subnetwork that shares knowledge across tasks. Another interesting multi-task approach is policy distillation (Rusu et al., 2016). This method combines different policies for multiple tasks into a single network, which enables to share information between tasks and to learn a compact network that can even outperform the individual policies.
3 BACKGROUND: STATE REPRESENTATION LEARNING FOR REINFORCEMENT LEARNING
We formulate MT-LRP in a reinforcement learning (RL) setting using a Markov decision process (MDP) (S,A,T,R,γ): Based on the current state s ∈ S, the agent chooses and executes an action a ∈ A, obtains a new state s′ ∈ S (according to the transition function T ) and collects a reward r ∈ R. The agent’s goal is to learn a policy π : S→ A that maximizes the expected return IE(∑∞t=0 γ trt), with rt being the reward collected at time t and 0 < γ ≤ 1 the discount factor. We consider an episodic setting with episodes of finite length, a continuous state space S and a discrete action space A.
In this work, we assume that the agent cannot directly observe the state s but only has access to observations o ∈ O, which are usually high-dimensional and contain task-irrelevant distractors. This requires us to extract the state from the observations by learning an observation-state-mapping ϕ : O→ S, and use the resulting state representation S to solve the RL problem (assuming that a Markov state can be extracted from a single observation). To learn the state representation, we apply learning with robotic priors (Jonschkowski & Brock (2015), from now on referred to as LRP). This method learns ϕ from a set of temporally ordered experiences D = {(ot ,at ,rt)}dt=1 by optimizing the following loss:
LRP(D,ϕ) = ωtLtemp.(D,ϕ)+ωpLprop.(D,ϕ)+ωcLcaus.(D,ϕ)+ωrLrep.(D,ϕ). (1)
This loss consists of four terms, each expressing a different prior about suitable state representations for robot RL. We optimize it using gradient descent, assuming ϕ to be differentiable. We now explain the four robotic prior loss terms in Eq. (1).
Temporal Coherence enforces states to change gradually over time (Wiskott & Sejnowski, 2002): Ltemp.(D,ϕ) = IE [ ‖∆st‖2 ] ,
where ∆st = st+1− st denotes the state change. (To increase readability we replace ϕ(o) by s.) Proportionality expresses the prior that the same action should change the state by the same magnitude, irrespective of time and the location in the state space:
Lprop.(D,ϕ) = IE [ (‖∆st2‖−‖∆st1‖) 2 ∣∣∣ at1 = at2].
Causality enforces two states st1 ,st2 to be dissimilar if executing the same action in st1 generates a different reward than in st2 .
Lcaus.(D,ϕ) = IE [ e−‖st2−st1‖ 2 ∣∣∣ at1 = at2 ,rt1+1 6= rt2+1].
Repeatability requires actions to have repeatable effects by enforcing that the same action produces a similar state change in similar states:
Lrep.(D, ϕ̂) = IE [ e−‖st2−st1‖ 2‖∆st2 −∆st1‖ 2 ∣∣∣ at1 = at2].
Additionally, the method enforces simplicity by requiring s to be low-dimensional.
Note that learning with robotic priors only makes use of the actions a, rewards r, and temporal information t during optimization, but not at test time for computing ϕ(o) = s. Using a, r and t in this way is an instance of the learning with side information paradigm (Jonschkowski et al., 2015).
4 MULTI-TASK STATE REPRESENTATIONS: MT-LRP
Now consider a scenario in which an agent is learning multiple distinct tasks. For each task τ ∈ {1, . . . ,T}, the agent now requires a task-specific policy πτ : Sτ → A. We approach the problem by learning a task-specific state representation ϕτ : O→ Sτ for each policy, and a task detector χ which determines the task, given the current observation. We will consider a probabilistic task-detector χ : O→ [0,1]T that assigns a probability to each task being active. In order to solve the full multi-task RL problem, we must learn χ, {ϕτ}τ∈{1,...,T} and {πτ}τ∈{1,...,T}. We propose to address this problem by MT-LRP, a method that jointly learns χ and {ϕτ}τ∈{1,...,T} from raw observations, actions, and rewards. MT-LRP then uses the state representations {ϕτ} to learn task-specific policies {πτ}τ∈{1,...,T} (using standard RL methods), and switches between them using the task detector χ . To solve the joint learning problem, MT-LRP generalizes LRP (Jonschkowski & Brock, 2015) in the following regards: (i) we replace the linear observation-statemapping from the original method with a gated neural network, where the gates act as task detectors that switch between different task-specific observation-state-mappings; (ii) we extend the list of robotic priors by the prior of task coherence, which allows us to train multiple task-specific state representations without any specification (or labels) of tasks and states.
4.1 GATED NEURAL NETWORK ARCHITECTURE
We use a gated neural network architecture as shown schematically in Fig. 2. The key idea is that both the task detector χ as well as the state representation ϕ are computed from raw inputs. However, the output of the task detector gates the output of the state representation. Effectively, this means the output of χ(o) decides which task-specific state representation ϕτ is passed further to the policy, which is also gated by the output of χ(o).
Formally, χ(o) = σ(χpre(o)) is composed of a function χpre with T -dimensional output and a softmax σ(z) = e z j
∑k ezk . The softmax ensures that χ computes a proper probability distribution over tasks.
The probabilities are then used to gate ϕ . To do this, we decompose ϕ into a pre-gating function
ϕpre that extracts features shared across all tasks (i.e. ”multi-task” in the sense of Caruana (1997), unless set to the identity), and a T ×M×N gating tensor G that encodes the T (linear) observationstate mappings (M = dim(s) and N is the output dimension of ϕpre). The value of the state’s i-th dimension si computes as the expectation of the dot product of gating tensor and ϕpre(o) over the task probabilities χ(o):
si = ϕi(o) = T
∑ k=1 χk(o) 〈Gk,i,:,ϕpre(o)〉. (2)
4.2 LEARNING OBJECTIVE
To train the network, we extend the robotic prior loss LRP (Eq. 1), by a task-coherence prior Lτ :
L= LRP(D,ϕ)+ωτLτ(D,χ), (3)
where ωτ is a scalar weight balancing the influence of the additional loss term. Task coherence is the assumption that a task only changes between training episodes, not within the same episode. It does not presuppose any knowledge about the number of tasks or the task presented in an episode, but it exploits the fact that task switching weakly correlates with training episodes. Moreover, this assumption only needs to hold during training: since χ operates directly on the observation o, it can in principle switch the task at every point in time during execution. Task-coherence applies directly to the output of the task detector, χ(o), and consists of two terms:
Lcon+sepτ = Lconτ +L sep τ . (4)
The first term enforces task consistency during an episode: Lconτ = IE [ H(χ(ot1),χ(ot2)) ∣∣∣ episodet1 = episodet2], (5) where H denotes the cross-entropy H(p,q) = −∑x p(x) logq(x). It can be viewed as a measure of dissimilarity between probability distributions p and q. We use it to penalize χ if it assigns different task distributions to inputs ot1 , ot2 that belong to the same episode. Note that task-consistency can be viewed as a temporal coherence prior on the task level (Wiskott & Sejnowski, 2002).
The second term expresses task separation and encourages χ to assign tasks to different episodes: Lsepτ = IE [ e−H(χ(ot1 ),χ(ot2 )) ∣∣∣ episodet1 6= episodet2]. (6) This loss is complementary to task consistency, as it penalizes χ if it assigns similar task distributions to ot1 , ot2 from different episodes. Note that L sep τ will in general not become zero. The reason is that the number of episodes usually exceeds the number of tasks, and therefore two observations from different episodes sometimes do belong to the same task. We will evaluate the contribution of each of the two terms to learning success in Section 5.2.
5 EXPERIMENTS
We evaluate MT-LRP in two scenarios. In the multi-task slot-car racing scenario (inspired by Lange et al. (2012)), we apply MT-LRP to a linearly solvable problem, allowing us to easily inspect what and how MT-LRP learns. In slot-car racing, the agent controls one of multiple cars (Figure 1), with the goal of traversing the circuit as fast as possible without leaving the track due to speeding in curves. However, the agent does not know a priori which car it controls, and only receives the raw visual signal as input. Additionally, uncontrolled cars driving at random velocity, act as visual distractors. We turn this scenario into a multi-task problem in which the agent must learn to control each car, where controlling the different cars corresponds to separate tasks. We will now provide the technical details of our experimental set-up.
5.1 EXPERIMENTAL SET-UP: SLOT-CAR RACING
The agent controls the velocity of one car (see Fig. 1), receives a reward proportional to the car’s velocity, chosen from [0.01, 0.02, . . . , 0.1], and a negative reward of −10 if the car goes too fast
in curves. The velocity is subject to Gaussian noise (zero mean, standard deviation 10%) of the commanded velocity. All cars move on independent lanes and do not influence each other. The agent observes the scenario by getting a downscaled 16x16 RGB top-down view (dimension N = 16×16×3 = 768) of the car circuit (Fig. 1(b)). In our experiments, there are two or three cars on the track, and the agent controls a different one in every episode. To recognize the task, the agent must be able to extract a visual cue from the observation which correlates with the task. We study two types of visual cues: Static Visual Cue: The arrangement of cars stays the same in all episodes and a static visual cue (a picture of the controlled car) in the top-left image corner indicates which car is currently controlled. Dynamic Visual Cue: The agent always controls the same car (with a certain color), but in each task the car is located on a different lane (as in Fig. 1(b)).
Data Collection and Learning Procedure: The agent collects 40 episodes per task, each episode consisting of 100 steps. To select an action in each step, the agent performs ε-greedy exploration by picking a random action with probability ε = 0.3 and the best action according to its current policy otherwise. The agent computes a policy after every τ episodes, by first learning the observation-state mapping ϕ (state representation) and then computing policies π1, . . . ,πτ (based on the outcomes of the learned χ and ϕ). To monitor the agent’s learning progress, we measure the average reward the agent attains on T test episodes, i.e. one test episode of length 100 per task (using the greedy policy), amounting to 8000 experiences in total. To collect sufficient statistics, the whole experiment is repeated 10 times. Policy Learning: We consider the model-free setting with continuous states S, discrete actions A and solve it using nearest-neighbor Q-learning kNN-TD-RL (Martı́n H et al., 2009) with k = 10. More recent approaches to model-free RL would be equally applicable (Mnih et al., 2015). Learning Strategies and Baselines: We compare five strategies. We run a) MT-LRP with 5 gate units (two/three more than necessary), state dimensionality M = 2 and using Lcon+sepτ as taskcoherence prior. We compare MT-LRP to several state representation methods; for each method we evaluate different M and report only the best performing M: a) robotic priors without gated network, LRP (M = 4), b) principal components analysis (PCA) on the observations (M = 20) and c) raw observations (M = 768). Additionally, we evaluate d) a lower baseline in the form of a randomly moving agent and e) an upper baseline by applying RL on the known 2D-position of the slot car under control (M = 2). We use the same RL algorithm for all methods. To learn the state representations with robotic priors, we base our implementation on Theano and lasagne, using the Adam optimizer with learning rate 0.005, batch size 100, Glorot’s weight initialization and ωt = 1,ωp = 5,ωc = 1,ωr = 5,ωτ = 10. Moreover, we apply an L1 regularization of 0.001 on ϕ . Additionally, we analyze the contribution of task coherence priors by applying MT-LRP to the full set of 8000 experiences a) without task-coherence, b) with task consistency Lconτ only c) with task separation Lconτ only) and d) without task consistency and separation L con+sep τ .
5.2 RESULTS
We will now present the three main results of our experiments: (i) we show that MT-LRP enables the agent to extract better representations for RL; (ii) we provide insight in how the learner detects the task and encodes the state representations; and finally, (iii) we show the contribution of each of the task-coherence loss terms.
MT-LRP Extracts Better State Representations for RL Figure 3 shows the learning curves for RL based on state representations learned by the different methods in the two-slot-car scenario (static visual cue on the left, dynamic on the right). No method reaches the performance of the upper baseline, mainly due to aliasing errors resulting from the low image resolution.
The random baseline ranges around an average reward of −84.9 with standard error 0.72 and was omitted from the Figure. The state representation learning baselines without robotic priors perform poorly because they are unable to identify the taskirrelevant distractions. MT-LRP gets very close to the performance of the upper baseline, especially for very low amounts of training data (d < 2500), whereas LRP does not even attain this level of performance for the full training set d = 8000 in the static task. The gap between MT-LRP and LRP increases even more if we add another car (Figure 5) because LRP can only learn one state representation for all three tasks. Including the three slot cars in this representation results in distractions for the RL method. However, in the dynamic-visual-cue scenario LRP-4 performs on par with MT-LRP. Surprisingly, running LRP with only two dimensions suffices to achieve the performance of MT-LRP. We
will explain this phenomenon below. To conclude, MT-LRP allows to learn as good or better policies than the baselines in all slot-car scenarios.
MT-LRP Detects All Tasks and Learns Good State Representations To gain more insight into what is learned, we analyze the state representations extracted by MT-LRP and LRP. Figure 4 shows the state representation learned by MT-LRP for the static-visual-cue scenario. Each point in the figure corresponds to one observation, markers indicate the task and colors the most active gate unit. We see that the first gate unit (blue) is always active for task 1 (circle), and the second gate unit for task 2. This shows that the task is detected with high accuracy. The task detector χ is also highly cer-
tain which is reflected in the fact that its entropy evaluated on the data is close to zero. Moreover, the states reflect the circular structure of the slot car racing track. We thus conclude that MT-LRP has learned to identify the tasks and to represent the position of each car on the track.
The RL experiments raised the question why LRP manages to solve the dynamic, but not the static-visual-cue scenario as well as MT-LRP. We hypothesize that, for the dynamic cue, LRP is able to extract the position of the car on regardless of which lane it is in using a single linear mapping. Figure 6 confirms this hypothesis: LRP filters for the car’s color (blue) along the track and assigns increasing weights to these pixels which results in the extraction of its position. It also assigns constant weights along the track in the red channel using the lane change of the two cars as an offset. This results in a mapping to two circles similar to Fig. 4, where the state encodes both the position and the task. Such a mapping can be expressed by a linear
function precisely because the features that are relevant for one task do not reappear in another task (e.g. a blue slot car in track 1 does not appear in the task where the blue car is in track 2).
However, there exists no equivalent linear mapping for the static-visual-cue variant of the slotcar problem, because cars that are relevant for one task are also present in every other task.
We can generalize from this insight as follows. A single linear observation-state-mapping is sufficient for multiple tasks if the state representation for every task can be extracted by a linear function using only features that stay constant for all other tasks. If this is the case, than there is no need for decoupling the extraction of task and state.
Task-Consistency is Critical for Learning Performance To understand the influence of
the different task-coherence prior variants, we compared their performance in Figure 7. We see that relying solely on the robotic priors gives poor results, mainly because the gate units are not used properly: more than one gate unit is activated per task (χ has high entropy). Adding the taskseparation prior forces the network to use as many gates as possible (5 in our case), leading to bad state representations. Interestingly, using task consistency only gives roughly the same result as using task consistency and task separation.
Discussion The experiments showed that MT-LRP is able to solve the representation and reinforcement learning tasks better than the baselines. Important questions for future work concern: the necessity and influence of the task-separation loss, in particular for short episode lengths and if the number of expected tasks exceeds the number of actual tasks; and transferring knowledge by adding a shared neural network layers before gating.
6 CONCLUSION
We have presented MT-LRP, a method for multi-task state representation learning with robotic priors. The method learns in an unsupervised fashion, solely based on the robots own observations, actions, and rewards. Our experiments confirmed that MT-LRP is effective in simultaneously identifying tasks and learning task-specific state representations. This capability is beneficial for scaling reinforcement learning to realistic scenarios that require dedicated skills for different tasks.
ACKNOWLEDGMENTS
We gratefully acknowledge the funding provided by the German Research Foundation (DFG, Exploration Challenge, BR 2248/3-1), the Alexander von Humboldt foundation through an Alexander von Humboldt professorship (funded by the German Federal Ministry of Education and Research). Additionally, Antonin Raffin was supported by an Erasmus+ grant.

This paper builds upon the method of Jonschkowski & Brock to learn state representations for multiple tasks, rather than a single task. The research direction of learning representations for multiple tasks is an interesting one, and largely unexplored. The approach in the paper is to learn a different representation for each task, and a different policy for each task, where the task is detected automatically and built into the neural network.
The authors state that the proposed method is orthogonal to multi-task learning, though the end goal of learning to solve multiple tasks is the same. It would be interesting and helpful to see more discussion on this point in the paper, as discussed in the pre-review question phase. References to other multi-task learning works, e.g. policy distillation and actor-mimic (both ICLR ’16), may be appropriate as well.
The method proposes to jointly learn a task classifier with a state representation learner, by using a differentiable gating mechanism to control the flow of information. The paper proposes a task coherence prior for this gating mechanism to ensure that the learned task classifier is temporally coherent. Introducing this structure is what enables the method to improve performance over the standard, non-multitask approach.
The evaluation involves two toy experimental scenarios. The first involves controlling one of two cars to drive around a track. In this task, detecting the “task” is very easy, and the learned state representation is linear in the observation. The paper evaluates the performance of the policies learned with the proposed approach, and shows sufficient comparisons to demonstrate the usefulness of the approach over a standard non-multitask set-up.
In the second navigation scenario, only the state representation is qualitatively shown, not the resulting control policy nor any other learned state representations for comparison. Since the multi-task state representation learning approach is only useful if you can also learn control better, the paper should also evaluate on control, with the same comparisons as in the first experiment. Without this evaluation, the experiment is incomplete.
Lastly, to be on par with publications at a venue like ICLR, the method should be evaluated more thoroughly, on a wider range of set-ups, to demonstrate the generality of the approach and show that the method applies to more complex tasks. While in theory, the method should scale, the experiments do not demonstrate that it can handle more realistic scenarios, such as scaling beyond MNIST-level images, to 3D or real images, or higher-dimensional control tasks. Evaluating the method in this more complex scenario is important, because unexpected issues can come up when trying to scale. If scaling-up is straight-forward, then running this experiment (and including it in the paper) should be straight-forward.
In summary, here are the pros and cons of this paper:
Cons
- The approach does not necessarily share information across tasks for better learning, and requires learning a different policy for each task
- Only one experimental set-up that evaluates learned policy with multi-task state representation
- No experiments on more realistic scenarios, such 3D environments or high-dimensional control problems
Pros:
- This approach enables using the same network for multiple tasks, which is often not true for transfer and multi-task learning approaches
- Novel way to learn a single policy for multiple tasks, including a task coherence prior which ensures that the task classification is meaningful
- Experimentally validated on two toy tasks. One task shows improvement over baseline approaches
Thus, my rating would be higher if the paper included an evaluation of the control policy for navigation and included another more challenging and compelling scenario.
Lastly, here are some minor comments/questions on how I think the paper could be improved, but are not as important as the above:
Approach:
Could this approach be combined with other state representation learning approaches? e.g. approaches that use an autoencoder.
Experiments:
One additional useful comparison would be to evaluate performance in the single-task setting (e.g. only controlling the red car), as an upper bound on how well the policy should be able to perform. Does the learned multi-task policy reach the same level of performance? This upper bound will be tighter than the “known car position” baseline (which is also useful in its own right).
Does the “observations” baseline eventually reach the performance of the LRP approach? It would be useful to know if this approach simply speeds up learning (significantly) or if it enables better performance.
If there are aliasing issues with the images, why not just use higher resolution images?