PRIMA: Planner-Reasoner Inside a Multi-task Reasoning Agent

1 INTRODUCTION . Multi-task learning ( MTL ) ( Zhang & Yang , 2021 ; Zhou et al. , 2011 ) demonstrates superior sample complexity and generalizability compared with the conventional “ one model per task ” style to solve multiple tasks . Recent research has additionally leveraged the great success of deep learning ( LeCun et al. , 2015 ) to empower learning deep multi-task models ( Zhang & Yang , 2021 ; Crawshaw , 2020 ) . Deep MTL models either learn a common multi-task feature representation by sharing several bottom layers of deep neural networks ( Zhang et al. , 2014 ; Liu et al. , 2015 ; Zhang et al. , 2015a ; Mrksic et al. , 2015 ; Li et al. , 2015 ) , or learn task-invariant and task-specific neural modules ( Shinohara , 2016 ; Liu et al. , 2017 ) via generative adversarial networks ( Goodfellow et al. , 2014 ) . Although MTL is successful in many applications , a major challenge is the often impractically large MTL models . Although still smaller than piling up all models across different tasks , existing MTL models are significantly larger than a single model for tackling a specific task . This results from the intrinsic conflict underlying all MTL algorithms : balancing across-task generalization capability to perform different tasks with single-task efficiency in executing a specific task . On one hand , good generalization ability requires an MTL agent to be equipped with a large set of skills that can be combined to solve many different tasks . On the other hand , solving one particular task does not require all these skills . Instead , the agent needs to compose only a ( small ) subset of these skills into an efficient solution for a specific task . This conflict often hobbles existing MTL approaches . This paper focuses on multi-task reasoning ( MTR ) , a subarea of MTL that uses logic reasoning to solve multiple tasks . MTR is ubiquitous in human reasoning , where humans construct different reasoning paths for multiple tasks from the same set of reasoning skills . Conventional deep learning , although capable of strong expressive power , falls short in reasoning capabilities ( Bengio , 2019 ) . Considerable research has been devoted to endowing deep learning with logic reasoning abilities , the results of which include Deep Neural Reasoning ( Jaeger , 2016 ) , Neural Logic Reasoning ( Besold et al. , 2017 ; Bader et al. , 2004 ; Bader & Hitzler , 2005 ) , Neural Logic Machines ( Dong et al. , 2019 ) , and other approaches ( Besold et al. , 2017 ; Bader et al. , 2004 ; Bader & Hitzler , 2005 ) . However , these approaches consider only single-task reasoning rather than a multi-task setting , and applying 1The code will be released after acceptance . existing MTL approaches to learning these neural reasoning models leads to the same conflict between across-task generalization capability and single-task efficiency . To strike a balance between reasoning capability and efficiency in MTR , we develop a PlannerReasoner architecture Inside a Multi-task reasoning Agent ( PRIMA ) ( Section 2 ) , wherein the reasoner defines a set of neural logic operators for modeling reusable reasoning meta-rules ( “ skills ” ) across tasks ( Section 2.2 ) . When defining the logic operators , we focus on first-order logic because of its simplicity and wide applicability to many reasoning problems , such as automated theorem proving ( Fitting , 2012 ; Gallier , 2015 ) and knowledge-based systems ( Van Harmelen et al. , 2008 ) . A separate planner module activates only a small subset of the meta-rules necessary for a given task and composes them into a deduction process ( Section 2.3 ) . Thus , our planner-reasoner architecture features the dual capabilities of composing and pruning a logic deduction process , achieving a graceful capability-efficiency trade-off in MTR ( Section 2.4 ) . The model architecture is trained in an end-to-end manner using deep reinforcement learning ( Section 3 ) , and experimental results on several benchmarks demonstrate that this framework leads to a more principled predicate space search and reduces reasoning complexity ( Section 4 ) . We discuss related works in Section 5 , and conclude our paper in Section 6 . 2 PLANNER-REASONER FOR MULTI-TASK REASONING . This section proposes the Planner-Reasoner framework for MTR . To that end , we first formally state the logic reasoning problem in Section 2.1 , and then in Sections 2.2 and 2.3 , we describe the Planner-Reasoner Inside a Single Reasoning Agent ( PRISA ) framework . This is a plannerreasoner architecture , which is a neural-logic architecture for traversing the first-order predicate space . Typically , logical reasoning can be reduced to learning to search for a reasoning path with logical operators and then derive a logical consequence from premises . Therefore , a reasoning problem can be addressed in two steps : ( i ) constructing the elementary logical operators and ( ii ) selecting a reasoning path that chains these logical operators together . This key observation motivates the Reasoner ( Section 2.2 ) and Planner ( Section 2.3 ) modules in our framework . In Section 2.4 , the PRISA framework is extended to the MTR setting , which results in the Planner-Reasoner Inside a Multi-task Reasoning Agent ( PRIMA ) framework . 2.1 PROBLEM FORMULATION . Logic reasoning We begin with a brief introduction of the logic reasoning problem . Specifically , we consider a special variant of the First-Order Logic ( FOL ) system , which only consists of individual variables , constants and up to r-ary predicate variables . That is , we do not consider functions that map individual variables/constants into terms . An r-ary predicate p ( x1 , . . . , xr ) can be considered as a relation between r constants , which takes the value of True or False . An atom p ( x1 , · · · , xr ) is an r-ary predicate with its arguments x1 , · · · , xr being either variables or constants . A well-defined formula in our FOL system is a logical expression that is composed from atoms , logical connectives ( e.g. , negation ¬ , conjunction ^ , disjunction _ , implication ) , and possibly existential 9 and universal 8 quantifiers according to certain formation rules ( see Andrews ( 2002 ) for the details ) . In particular , the quantifiers 9 and 8 are only allowed to be applied to individual variables in FOL . In Fig . 1 , we give an example from the AdjacentToRed task ( Graves et al. , 2016 ) and show how it could be formulated as a logical reasoning problem . Specifically , we are given a random graph along with the properties ( i.e. , the color ) of the nodes and the relations ( i.e. , connectivity ) between nodes . In our context , each node i in the graph is a constant and an individual variable x takes values in the set of constants { 1 , . . . , 5 } . The properties of nodes and the relations between nodes are modeled as the unary predicate IsRed ( x ) ( 5⇥ 1 vector ) and the binary predicate HasEdge ( x , y ) ( 5⇥ 5 matrix ) , respectively . The objective of logical reasoning is to deduce the value of the unary predicate AdjacentToRed ( x ) ( i.e. , whether a node x has a neighbor of red color ) from the base predicates IsRed ( x ) and HasEdge ( x , y ) ( see Fig . 1 for an example of the deduction process ) . Multi-Task Reasoning Next , we introduce the definition of MTR . With a slight abuse of notations , let { p ( x1 , . . . , xr ) : r 2 [ 1 , n ] } be the set of input predicates sampled from any of the k different reasoning tasks , where x1 , . . . , xr are the individual variables and n is the maximum arity . A multi-task reasoning model takes p ( x1 , . . . , xr ) as its input and seeks to predict the corresponding ground-truth output predicates q ( x1 , . . . , xr ) . The aim is to learn multiple reasoning tasks jointly in a single model so that the reasoning skills in a task can be leveraged by other tasks to improve the general performance of all tasks at hand . 2.2 REASONER : TRANSFORMING LOGIC RULES INTO NEURAL OPERATORS . The Reasoner module conducts logical deduction using a set of neural operators constructed from first-order logic rules ( more specifically , a set of “ learnable ” Horn clauses ) . Its architecture is inspired by NLM ( Dong et al. , 2019 ) ( details about the difference can be found in Appendix D ) . Three logic rules are considered as essential meta-rules : BooleanLogic , Expansion , and Reduction . BooleanLogic : expression ( x1 , x2 , · · · , xr ) ! p̂ ( x1 , x2 , · · · , xr ) , where expression is composed of a combination of Boolean operations ( AND , OR , and NOT ) and p̂ is the output predicate . For a given r-ary predicate and a given permutation 2 Sn , we define p ( x1 , · · · , xr ) = p ( x ( 1 ) , · · · , x ( r ) ) where Sn is the set of all possible permutations as the arguments to an input predicate . The corresponding neural implementation of BooleanLogic is ( MLP ( p ( x1 , · · · , xr ) ) ; ✓ ) , where is the sigmoid activation function , MLP refers to a multilayer perceptron , a Permute ( · ) neural operator transforms input predicates to p ( x1 , · · · , xr ) , and ✓ is the learnable parameter within the model . This is similar to the implicit Horn clause with the universal quantifier ( 8 ) , e.g. , p1 ( x ) ^ p2 ( x ) ! p̂ ( x ) implicitly denoting 8x p1 ( x ) ^ p2 ( x ) ! p̂ ( x ) . The class of neural operators can be viewed as “ learnable ” Horn clauses . Expansion , and Reduction are two types of meta-rules for quantification that bridge predicates of different arities with logic quantifiers ( 8 and 9 ) . Expansion introduces a new and distinct variable xr+1 for a set of r-ary predicates with the universal quantifier ( 8 ) . For this reason , Expansion creates a new predicate q from p. Expansion : p ( x1 , x2 , · · · , xr ) ! 8xr+1 , q ( x1 , x2 , · · · , xr , xr+1 ) , where xr+1 /2 { xi } ri=1 . The corresponding neural implementation of Expansion , denoted by Expand ( · ) , expands the r-ary predicates into the ( r + 1 ) -ary predicates by repeating the r-ary predicates and stacking them in a new dimension . Conversely , Reduction removes the variable xr+1 in a set of ( r + 1 ) -ary predicates via the quantifiers of 8 or 9 . Reduction : 8xr+1 p ( x1 , x2 , · · · , xr , xr+1 ) ! q ( x1 , x2 , · · · , xr ) , or 9xr+1 p ( x1 , x2 , · · · , xr , xr+1 ) ! q ( x1 , x2 , · · · , xr ) . The corresponding neural implementation of Reduction , denoted by Reduce ( · ) , reducing the ( r+1 ) ary predicates into the r-ary predicates by taking the minimum ( resp . maximum ) along the dimension of xr+1 due to the universal quantifier 8 ( resp . existential quantifier 9 ) .

Summary: The paper introduces a neuro-symbolic framework for multi-task reasoning that builds on NLM. Specifically, the method tries to achieve higher efficiency at the inference time, by using an RL component that controls the operations performed as part of the reasoning process. They adapt the system for multi-task settings by introducing nullary predicates that capture the specific task details. The method is compared to multiple baselines, including an adaptation of NLM to the multi-task setting. While the methods perform similar to the baselines in terms of the overall accuracy, the approach does result in a lower number of floating-point operations at the inference time.