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: This paper presents PRIMA (“Planner-Reasoner Inside a Multi-Task Reasoning Agent”), a multi-task reasoning model that can be applied to several tasks that require first order logical reasoning. Notably, this paper seeks to address a problem that plagues multi-task learning more generally — the tradeoff between cross-task generalizability, and single-task effectiveness. Prior work in multi-task learning shows that it’s hard to strike this balance, and most work in building neural systems for logical reasoning focus on the single task setting.  The main contribution of this paper is first a single task planning-based agent that learns to reason for an individual logical task (e.g., a graph reasoning task like `4-Connectivity` or `1-Outdegree`) by routing input problem representations (predicates describing the problem) through a planning module; this planner acts as a high-level policy that (over multiple steps) routes the input through individual “reasoning” components that are “neural” versions of first-order logic meta-operators. This is a compact, two-party approach for solving complex reasoning problems; the reasoner defines the atomic operations, while the planner figures out what operations to apply, in what order.  The full system is learned via reinforcement learning, first via model-free methods such as REINFORCE and PPO, but with significantly better results using the structure of the problem (we can always roll-out through the planner) to use Model-based RL, specifically the MuZero formulation. Using MuZero/MCTS to learn the Planner/Reasoner, PRIMA shows competitive performance on single-tasks, and to generalize to multiple tasks, we simply add a “nullary predicate” that specifies the task ID. PRIMA performs well, comparably to other approaches for logical reasoning, even when prior approaches (such as the Neural Logic Machine) are extended to multi-task settings in the same way (adding this extra task conditioning information). Where PRIMA shows benefits over prior approaches are in aggregate FLOPs (showing that PRIMA may be learning to reason more efficiently), and in a single probing experiment showing how PRIMA can generalize to new input classes unseen at training.