Year: 2021

Abstract:
We consider the problem of multi-task reasoning (MTR), where an agent can solve multiple tasks via (first-order) logic reasoning. This capability is essential for human-like intelligence due to its strong generalizability and simplicity for handling multiple tasks. However, a major challenge in developing effective MTR is the intrinsic conflict between reasoning capability and efficiency. An MTR-capable agent must master a large set of “skills” to tackle diverse tasks, but executing a particular task at the inference stage requires only a small subset of immediately relevant skills. How can we maintain broad reasoning capability but efficient specific-task performance? To address this problem, we propose a Planner-Reasoner framework capable of state-of-the-art MTR capability and high efficiency. The Reasoner models shareable (first-order) logic deduction rules, from which the Planner selects a subset to compose into efficient reasoning paths. The entire model is trained in an end-to-end manner using deep reinforcement learning, and experimental studies over a variety of domains validate its effectiveness. 1

heading: 1 INTRODUCTION
text: 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.

heading: 2 PLANNER-REASONER FOR MULTI-TASK REASONING
text: 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.

heading: 2.1 PROBLEM FORMULATION
text: 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.

heading: 2.2 REASONER: TRANSFORMING LOGIC RULES INTO NEURAL OPERATORS
text: 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).

heading: 2.3 PLANNER: ACTIVATING AND FORWARD-CHAINING LEARNABLE HORN CLAUSES
text: The Planner is our key module to address the capability-efficiency tradeoff in the MTR problem; it is responsible for activating the neural operators in the Reasoner and chaining them into reasoning paths. Existing learning-to-reason approaches, which are often based on inductive logic programming (ILP) (Cropper et al., 2020; Cropper & Dumancic, 2020; Muggleton & De Raedt, 1994) and the correspondingly neural-ILP methods (Dong et al., 2019; Shi et al., 2020). Conventional ILP methods suffer from several drawbacks, such as heavy reliance on human-crafted templates and sensitivity to noise. On the other hand, neural-ILP methods (Dong et al., 2019; Shi et al., 2020), leveraging
the strength of deep learning and ILP, such as the Neural Logic Machine (NLM) (Dong et al., 2019), lack explicitness in the learning process of searching for reasoning paths. Let us take the learning process of NLM for example, which follows an intuitive two-step procedure. It first fully connects all the neural blocks and then searches all possible connections (corresponding to all possible predicate candidates) exhaustively to identify the desired reasoning path (corresponding to the desired predicate).
By using our proposed Planner module, we can strike a better capability-efficiency tradeoff. Rather than conducting an exhaustive search over all possible predicate candidates as in NLM, the Planner prunes all the unnecessary operators and identifies an essential reasoning path with low complexity for a given problem. Consider the following example. As shown in Fig. 1 and Fig. 2A (the highlighted orange-colored blocks), at each reasoning step, the Planner takes the input predicates and determines which neural operators should be activated. The decision is represented as a binary vector — [I0 . . . I6] (termed operator footprint) — that corresponds to the neural operators of Expand, Reduce and DirectInput2 at different arity groups. By chaining these sequential decisions, a sequence of operator footprints is formulated as a reasoning path, as the highlighted orange-colored path in Fig. 2A. Generally, the neural operators defined in the Reasoner can also be viewed as (learnable) Horn Clauses (Horn, 1951; Chandra & Harel, 1985), and the Planner forward-chains them into a reasoning path.

heading: 2.4 THE PRIMA FRAMEWORK FOR MULTI-TASK REASONING
text: By far, we have developed the PRISA architecture for the single-task reasoning setting. To extend PRISA to PRIMA for the multi-task reasoning setting, we add an extra nullary input predicate (O0 in Fig. 2) to inform the planner of the current task. Based on the value of O0, the planner will learn to activate the appropriate subset of neural operators, including a subset of shared operators and a few task-specific operators, for the current task. By doing so, PRIMA can learn to make the best use of reusable neural operators across different tasks. In Section 4, we will demonstrate such a capability of the planner on different tasks.

heading: 3 LEARNING-TO-REASON VIA REINFORCEMENT LEARNING
text: In our problem, the decision variables of the Planner are binary indicators of whether a neural operator module should be activated. Readers familiar with Markov decision processes (MDPs) (Puterman, 2014) might notice that the reasoning path of our model in Fig. 2A resembles a temporal rollout in MDP formulation (Sutton & Barto, 2018). Therefore, we frame learning-to-reason as a sequential decision-making problem and adopt off-the-shelf reinforcement learning (RL) algorithms.
2DirectInput is an identity mapping of the inputs where its following operators of Permute, MLP and sigmoid nonlinearity can directly approximate the BooleanLogic.

heading: 3.1 AN MDP FORMULATION OF THE LEARNING-TO-REASON PROBLEM
text: An MDP is defined as the tuple (S,A, P ass0 , R, ), where S and A are finite sets of states and actions, the transition kernel P ass0 specifies the probability of transition from state s 2 S to state s0 2 S by taking action a 2 A, R(s, a) : S ⇥A! R is the reward function, and 0   1 is a discount factor. A stationary policy ⇡ : S ⇥A! [0, 1] is a probabilistic mapping from states to actions. The primary objective of an RL algorithm is to identify a near-optimal policy ⇡⇤ that satisfies
⇡ ⇤ := argmax
⇡
⇢ J(⇡) := E  Tmax 1X
t=0
t r(st, at ⇠ ⇡) ,
where Tmax is a positive integer denoting the horizon length — that is, the maximum length of a rollout. The resemblance between Fig. 2 and an MDP formulation is relatively intuitive as summarized in Table 1. At the t-th time step, the state st corresponds to the set of predicates, st = [O0t , O1t , O2t ], with the superscript denoting the corresponding arity. The action at = [I0t , I1t , . . . , I K 1 t ] (e.g., a0 = [000011]) is a binary vector that indicates the activated neural operators (i.e., the operator footprint), where K is the number of operators per layer. The reward rt is defined to be
rt :=
( PK 1
i=0 I i t , if t < Tmax
Accuracy, t = Tmax . (1)
That is, the terminal reward is set to be the reasoning accuracy at the end of the reasoning path (see Appendix A.2 for its definition), and the intermediate reward at each step is chosen to be the negated number of activated operators (which penalizes the cost of performing the current reasoning step). The transition kernel P ass0 corresponds to the function modeled by one Reasoner layer; each Reasoner layer will take the state (predicates) st 1 = [O0t 1, O1t 1, O2t 1, ] and the action (operator footprint) at = [I0t , I 1 t , . . . , I K 1 t ] as its input and then generate the next state (predicates) st = [O0t , O1t , O2t , ]. This also implies that the Reasoner layer defines a deterministic transition kernel, i.e., given st 1 and at the next state st is determined.
Table 1: The identification between the concepts of PRIMA/PRISA and that of RL at the t-th time step.
RL State st Action at Reward rt Transition
Kernel P ass0 Policy Rollout
PRIMA/PRISA Predicates of different arities: [O0t , O 1 t , O 2 t ]t Operator footprint: [I0t . . . I K 1 t ] Eq. (1) One layer of Reasoner Planner Reasoning path

heading: 3.2 POLICY NETWORK: MODELING OF PLANNER
text: The Planner module is embodied in the policy network. As shown in Fig. 3A, the base planner is a separate module that has the same architecture of 1-layer (fully-activated) Reasoner followed
by a max-pooling layer. This architecture enables the reduction of input predicates to the specific indicators by reflecting whether the operations at the corresponding position of one layer of Reasoner are active or inactive. Further, we can also leverage Monte-Carlo Tree Search (MCTS) (Browne et al., 2012; Munos, 2014) to boost the performance, which leads to an Enhanced Planner (Fig. 3B). An MCTS algorithm such as the Upper Confidence Bound for Trees (UCT) method (Kocsis et al., 2006), is a model-based RL algorithm that plans the best action at each time step (Browne et al., 2012) by constructing a search tree, with states as nodes and actions as edges. The Enhanced Planner uses MCTS to exploit partial knowledge of the problem structure (i.e., the deterministic transition kernel P ass0 defined by the Reasoner layer) and construct a search tree to help identify the best actions (which ops to activate). Details of the MCTS algorithms used in the Enhanced Planner can be found in Appendix A.1.

heading: 3.3 OVERALL LEARNING FRAMEWORK
text: As illustrated in Fig. 3, we introduce concrete data-driven decision-making methods—that is, RL approaches (Sutton & Barto, 2018)—to address the learning-to-reason problem. To illustrate this, we apply the model-free RL method REINFORCE (Williams, 1992) and the model-based method MuZero (Schrittwieser et al., 2019). Compared with model-free reinforcement learning, model-based reinforcement learning (MBRL) more effectively handles large search-space problems, such as the game of Go (Silver et al., 2017b;b;a). MuZero (Schrittwieser et al., 2019), a recently proposed MBRL approach to integrating planning and learning, has achieved great success with a variety of complex tasks. Motivated by the success of MuZero, we propose an MBRL approach for neural-symbolic reasoning. The key insight behind adopting MuZero is that in real applications, we typically have partial structural knowledge of the transition kernel P ass0 and reward function r(s, a). As a result of the model-based module, testing complexity can be greatly reduced by adopting MCTS, which leads to a better set of predicates. Of course, MuZero is just one option in the model-based family of approaches. We leave it as future research to propose and compare other model-based alternatives.
The pipeline of the training process is illustrated in Fig.2.B (the right subfigure). After loading the model weights, the reasoning path rollouts are executed by the agent (or model instance), according to the current policy network. The performed reasoning path rollout is then stored in the replay buffer. The Planner-Reasoner is trained via rollouts sampled from the replay buffer.

heading: 4 EXPERIMENTAL RESULTS AND ANALYSIS
text: In this section, we evaluate the performance of different variants of PRISA (for single-task reasoning) and PRIMA (for multi-task reasoning) on eight tasks from the family tree and graph benchmarks (Graves et al., 2016), including 1-Outdegree, AdjacentToRed, HasFather, HasSister, 4-Connectivity, IsGrandparent, IsUncle, IsMGUncle. These tasks are widely used benchmark domains for inductive logic programming (Krötzsch, 2020; Calautti et al., 2015). Detailed descriptions about those tasks can be found in Appendix B.1. We evaluate their testing accuracy and reasoning cost (measured in FLOPs: the number of floating-point operations executed (Clark et al., 2020)) on these tasks and compare them to several baselines. Besides, the case study is conducted on the reasoning path, which indicates the operator sharing among different tasks. All the results demonstrate the graceful capability-efficiency tradeoff of PRIMA in multi-task reasoning.

heading: 4.1 EXPERIMENTAL SETUPS
text: Regarding the data generated in the single-task setting, we use the same methods as in NLM (Dong et al., 2019). A task will be randomly sampled in the multi-task setting according to a pre-defined probability distribution on different tasks. Compared to the generated data with that of the single task, it is augmented by using one-hot encoding for different tasks and wrapping it into nullary (background) predicates. Also, task-specific output heads are introduced to adapt to the multi-task setting. These adaptions apply to NLM-MTR, DLM-MTR, and PRIMA. The reasoning accuracy is used as the reward for both PRISA-MuZero and PRIMA. In the inference (or testing) stage, the Reasoner is combined with the learned Base Planner to perform the tasks for PRISA-MuZero and PRIMA, instead of an enhanced planner (MCTS), to reduce the extra computation. The problem
size for training is always 10 for graph tasks and 20 for family tree tasks across all single-task and multi-task settings, regardless of the sizes of testing problems.

heading: 4.2 OVERALL PERFORMANCE
text: Single-task reasoning capability First, we compare the performance of three different variants of PRISA (Section 3.3) for single-task reasoning, which learns their planners based on different reinforcement learning algorithms; PRISA-REINFORCE uses REINFORCE (Williams, 1992), PRISA-PPO uses PPO (Schulman et al., 2017), and PRISA-MuZero uses MuZero (Schrittwieser et al., 2019). We report the test accuracy and the Percentage of Successful Seeds (PSS) in Table 2 to measure the model’s reasoning capabilities, where the PSS reaches 100% of success rates (Matthieu et al., 2021). We note that PRISA-MuZero has the same 100% accuracy as NLM (Dong et al., 2019), DLM (Matthieu et al., 2021), and @ILP (Evans & Grefenstette, 2018) across different tasks, and outperforms MemNN (Sukhbaatar et al., 2015) (shown in the single-task part in Table 3). But it also has a higher successful percentage (PSS) in comparison with other methods. The results show that PRISA-MuZero achieves the best performance on single-task reasoning, which confirms the strength of the MCTS-based Enhanced Planner (Section 3.3) and the MuZero learning strategy. Therefore, we will use the Enhanced Planner and MuZero in PRISA and PRIMA in the rest of our empirical studies.
Multi-task reasoning capability Next, we evaluate the MTR capabilities of PRIMA. To the best of our knowledge, there is no existing approach that is designed specifically for MTR. Therefore, we adapt NLM and DLM into their multi-task versions, named NLM-MTR and DLM-MTR, respectively. NLM-MTR and DLM-MTR follow the same input and output modification as what we did to upgrade PRISA to PRIMA (Section 2.4). By this, we can examine the contribution of our proposed Planner-Reasoner architecture for MTR. As shown in Table 3, PRIMA (with MuZero as the Base Planner) performs well (perfectly) on different reasoning tasks. On the other hand, DLM-MTR experiences some performance degradation (on AdjacentToRed). This result confirms that our Planner-Reasoner architecture is more suitable for MTR. We conjecture that the benefit comes from using a Planner to explicitly select the necessary neural operators for each task, avoiding potential conflicts between different tasks during the learning process.
Experiments are also conducted to test the performance of PRIMA with different problem sizes. The problem size in training is 10 for all graph tasks and 20 for all family-tree tasks. In testing, we evaluate the methods on much larger problem sizes (50 for the graph tasks and 100 for the family tree tasks), which the methods have not seen before. Therefore, the Planner must dynamically activate a proper set of neural operators to construct a path to solve the new problem. As reported in Fig. 4 and Tables 2 and 3, PRIMA can achieve the best accuracy and lower flops when the problem sizes for training and testing are different.
Reasoning efficiency To measure the reasoning efficiency of the proposed methods at the inference stage, PRIMA is compared with NLM, MemNN, and NLM-MTR in terms of FLOPs. As shown in Fig. 4, NLM and NLM-MTR demonstrate a similar performance and suffer the highest reasoning cost when it is tested with large problem sizes, such as 50 for graph tasks and 100 for family tree tasks. For MemNN, although the FLOPs of it seem low in most cases of testing, its testing accuracy is bad and cannot achieve accurate predictions (Table 3). In contrast, PRIMA can significantly reduce the reasoning complexity by intelligently selecting ops using a planner. Overall, PRIMA strikes a satisfactory capability-efficiency tradeoff in comparison with all available multi-tasking baselines.
Operator/Path sharing in MTR To take a closer look into how PRIMA achieves such a better capability-efficiency tradeoff (in Fig. 4), we examine the reasoning paths on three different graph tasks: 1-Outdegree, AdjacentToRed, and 4-Connectivity. Specifically, we sample instances from these three tasks and feed them into PRIMA separately to generate their corresponding reasoning paths. The results are plotted in Fig. 5A, where the gray paths denote the ones shared across tasks, and the colored ones are task-specific. It clearly shows that PRIMA learns a large set of neural operators sharable across tasks. Given each input instance from a particular task, PRIMA activates a set of shared paths along with a few task-specific paths to deduce the logical consequences.
Generalizability of the Planner To demonstrate the generalizability of the Planner module in PRIMA, we generate input instances of AdjacentToRed with topologies that have not been seen during training. In Fig. 5B, we show the reasoning paths activated by the Planner for these different topologies, which demonstrates that different input instances share a large portion of reasoning paths. This fact is not surprising as solving the same task of AdjacentToRed should rely on a common set of skills. More interestingly, we notice that even for solving this same task, the Planner will have to call upon a few instance-dependent sub-paths to handle the subtle inherent differences (e.g., the graph topology) that exist between different input instances. For this reason, PRIMA maintains an instance-dependent dynamic architecture, which is in sharp contrast to the Neural Architecture Search (NAS) approaches. Although NAS may also use RL algorithms to seek for a smaller architecture (Zoph & Le, 2017), it only searches for a static architecture that will be applied to all the input instances.

heading: 5 RELATED WORK
text: Figure 5: The reasoning paths of PRIMA. The gray arrows denote the “shared path”. The colored arrows in sub-figures (A) and (B) denote task-specific paths and instance-specific paths, respectively. In (A), besides the AdjacentToRed task we have introduced earlier, we also consider the 1-Outdegree task, which reasons about whether the out-degree of a node is exactly equal to 1, and the 4-Connectivity task, which is to decide whether there are two nodes connected within 4 hops.
et al., 2008), which explores the sharing of features across different tasks using regularization techniques (Shinohara, 2016; Liu et al., 2017). The second approach assumes that tasks are intrinsically related, such as low-rank learning (Ando & Zhang, 2005; Zhang et al., 2005), learning with task clustering (Gu et al., 2011; Zhang et al., 2016), and task-relation learning (such as via task similarity, task correlation, or task covariance) (Goncalves et al., 2016; Zhang, 2013; Ciliberto et al., 2015). MTL has also been explored under other paradigms such as unsupervised learning, for example, via multi-task clustering in (Zhang & Zhang, 2013; Zhang et al., 2015b; Gu et al., 2011; Zhang, 2015).
Neural-symbolic reasoning Neural-symbolic AI for reasoning and inference has a long history (Besold et al., 2017; Bader et al., 2004; Garcez et al., 2008), and neural-ILP has developed primarily in two major directions (Cropper et al., 2020). The first direction applies neural operators such as tensor calculus to simulate logic reasoning (Yang et al., 2017; Dong et al., 2019; Shi et al., 2020) and (Manhaeve et al., 2018). This approach uses binary tensors over constant domains to represent the predicates and tensor chain products to simulate logic clauses. The second direction involves relaxed subset selection (Evans & Grefenstette, 2018; Si et al., 2019) with a predefined set of task-specific logic clauses. This approach reduces the task to a subset-selection problem by selecting a subset of clauses from the predefined set and using neural networks to search for a relaxed solution.
Our work is different from the most recent work, such as (Dong et al., 2019; Shi et al., 2020), in several ways. Compared with (Dong et al., 2019), the most notable difference is the improvement in efficiency by introducing learning-to-reason via reinforcement learning. A second major difference is that PRIMA offers more generalizability than NLM by decomposing the logic operators into more finely-grained units. We refer readers for more detailed related work to Appendix D.

heading: 6 CONCLUSION
text: A long-standing challenge in multi-task learning is the intrinsic conflict between capability and efficiency. In this paper, we propose a Planner-Reasoner framework for multi-task reasoning. To maintain broad capability but efficient specific performance, the Reasoner extracts shareable metarules via first-order logic, from which the Planner efficiently selects a subset of meta-rules to formulate the reasoning path. The model’s training follows a complete data-driven end-to-end approach via deep reinforcement learning, and the performance is validated across a variety of benchmark tasks. Future work could include extending the framework to high-order logic and investigating scenarios when meta-rules have a hierarchical structure.
REPRODUCIBILITY STATEMENT
We commit to ensuring that other researchers with reasonable background knowledge in our area can reproduce our theoretical and empirical results. The algorithm details can be seen in Appendix A. Specifically, benchmark details are provided in Appendix B.1, hyper-parameter settings are provided in Appendix B.2, and computing infrastructure are detailed in Appendix B.3. The experiment details can be seen in Appendix B.
ETHICS STATEMENT
This work is about the methodology of achieving a capability-efficiency trade-off in multi-task first-order logic reasoning. The potential impact of this work is likely to further extend the framework to high-order logic and investigate scenarios when meta-rules have a hierarchical structure, which should be generally beneficial to the multi-task learning research community. We have not considered specific applications or practical scenarios as the goal of this work. Hence, it does not have any direct ethical consequences.

heading: A ALGORITHM DETAILS
text: A.1 DETAILS OF MCTS: FOUR KEY STEPS
During each epoch, MCTS repeatedly performs four sequential steps: selection, expansion, simulation, and backpropagation. The selection step traverses the existing search tree until the leaf node (or termination condition) by choosing actions (edges) as at each node s according to a tree policy. One widely used tree policy is the UCT (Kocsis et al., 2006) policy, which conducts the node-selection via
as = argmax s02C(s)
⇢ Vs0 + r 2 logNs Ns0 , (2)
where C(s) denotes the set of all child nodes for s, the first term Vs0 is an estimate of the longterm cumulative reward that can be received when starting from the state represented by node s0, and the second term represents the uncertainty (confidence-interval size) of that estimate. The confidence interval is calculated based on the upper confidence bound (UCB) (Auer et al., 2002; Auer, 2002) using Ns and Ns0 , which denote the number of times that nodes s and s0 have been visited (respectively). The key idea of UCT policy (2) is to select the best action according to an optimistic estimation (the upper confidence bound) of the expected return, which balances the exploitation (first term) and exploration (second term) with controlling the trade-off. The second step (node expansion) is conducted according to a prior policy by adding a new child node if the selection process reaches a leaf node of the search tree. Next, the simulation step estimates the value function (cumulative reward) V̂s by running the environment simulator with a default (simulation) policy. Finally, the backpropagation step updates the statistics Vs and Ns from the leaf node sT to the root node s0 of the selected path by recursively performing the following update (i.e., from t = T 1 to t = 0):
Nst Nst + 1, V̂st r(st, at) + V̂st+1 , Vst (Nst 1)Vst + V̂st /Nst ,
where V̂sT is the simulation return of sT , and at denotes the action selected following (2) at state st. Generally, the expansion step and the simulation step are more time-consuming than the other two steps, as they involve a large number of interactions with the environment.
In our Enhanced Planner, we adopt a recently proposed variant of MCTS that is based on probabilistic upper confidence tree (PUCT) (Rosin, 2011). It conducts node-selection according to a so-called PUCT score with two components PUCT(s, a) = Q(s, a) +U(s, a). Q(s, a) is the mean state-action value calculated from the averaged game result that takes action a during current simulations, and U(s, a) is the exploration bonus calculated as
U(s, a) = ⇡(a|s) pP b N(s, b)
1 +N(s, a)
✓ c1 + log ⇣P b N(s, b) + c2 + 1
c2
⌘◆ ,
where a, b are possible actions, ⇡(a|s) is the output of the Base Planner, N(s, a) is the visit count of the (s, a) pair during current simulations, and constants c1, c2 are the exploration-controlling hyper-parameters.
An MCTS algorithm is a model-based RL algorithm that plans the best action at each time step (Browne et al., 2012) by constructing a search tree, with states as nodes and actions as edges. It uses the MDP model to identify the best action at each time step until the leaf node (or until other termination conditions are satisfied) by choosing actions (edges) as at each node s according to a tree policy. For implementations of Monte-Carlo Tree Search, we refer readers to https: //github.com/werner-duvaud/muzero-general for details.
A.2 DETAILS OF CALCULATING REASONING ACCURACY
Since the “target” predicates are available for all objects or pairs of objects, the reasoning accuracy refers to the accuracy evaluated on all objects (for properties such as AdjacentToRed(x) or 1-OutDegree(x)) or pairs of objects (for relations such as 4-Connectivity(x, y)).

heading: B EXPERIMENTAL DETAILS
text: B.1 REASONING TASKS INTRODUCTION
Graph For graph tasks, they have the same background knowledge (or background predicate): HasEdge(x, y), i.e., HasEdge(x, y) is True if there is an undirected edge between node x and node y. However, AdjacentToRed has an extra background predicate, Red(x). Red(x) is True if the color of node x is red. Specifically, these graph tasks seeks to predict the target concepts (or target predicates) shown as below.
• 1-OutDegree: 1-OutDegree(x) is True if the out-degree of node x is exactly 1.
• AdjacentToRed: AdjacentToRed(x) is True if the node x has an edge with a red node.
• 4-Connectivity: 4-Connectivity(x, y) is True if there are within 4 hops between node x and node y.
Family Tree For family tree tasks, they have the same background knowledge (or background predicates): IsFather(x, y), IsMother(x, y), IsSon(x, y) and IsDaughter(x, y). For instance, IsFather(x, y) is True when y is x’s father. Specifically, these family tree tasks seek to predict the target concepts (or target predicates) shown below.
• HasFather: HasFather(x) is True if x has father.
• HasSister: HasSister(x) is True if x has at least one sister.
• IsGrandparent: IsGrandparent(x, y) is True if y is x’s grandparent.
• IsUncle: IsUncle(x, y) is True if y is x’s uncle.
• IsMGUncle: IsMGUncle(x, y) is True if y is x’s maternal great uncle.
B.2 HYPER-PARAMETER SETTINGS
For simplicity, we set discount factor = 1 in all experiments. The Tmax depends on the depth of the Reasoner. In single-task setting, we use the same hyper-parameter settings as in the original papers for MemNN (Sukhbaatar et al., 2015), @ILP (Evans & Grefenstette, 2018), NLM (Dong et al., 2019), and DLM (Matthieu et al., 2021). In multi-task setting, the probability distribution of sampling a task for NLM-MTR is [0.09, 0.11, 0.11, 0.09, 0.09, 0.135, 0.135, 0.24] that is corresponding to the task list of [1-Outdegree, AdjacentToRed, 4-Connectivity, HasFather, HasSister, IsGrandparent, IsUncle, IsMGUncle]. That distribution for DLM-MTR is [0.1, 0.12, 0.12, 0.1, 0.1, 0.13, 0.13, 0.2]. The details of hyper-parameters in NLM-MTR and DMLMTR can be found in Table 4. For all MLPs inside NLM-MTR, we keep the same settings as NLM. Similarly, the “Internal Logic”, initial temperature and scale of the Gumbel distribution, and the dropout probability in DLM-MTR are also kept the same settings as those in DLM. Besides, the probability distribution of sampling a task for PRIMA is [0.1, 0.12, 0.12, 0.1, 0.1, 0.13, 0.13, 0.2]. Regarding the hyper-parameters of PRIMA, details can be found in Table 5. We also show the hyper-parameters of PRISA-REINFORCE, PRISA-PPO, PRISA-MuZero in Table 5.
We note that the problem size for training is always 10 for graph tasks and 20 for family tree tasks in single-task and multi-task settings.
B.3 COMPUTING INFRASTRUCTURE
We conducted our experiments on a CPU server where the CPU is “Intel(R) Xeon(R) Silver 4114 CPU” with 40 cores and 64 GB memory in total.
Table 4: Hyper-parameter settings for NLM-MTR and DML-MTR.
NLM-MTR DML-MTR
learning rate 0.005 0.005 epochs 50 200 epoch size 8000 2000 batch size (train) 4 4 breadth 3 3 depth 4 9
Table 5: Hyper-parameter settings for PRISA and PRIMA. “lr-reasoner” refers to the learning rate for the reasoner, “lr-policy” denotes the learning rate for the policy network, “lr-value” denotes the learning rate for the value network, “residual” refers to the residual connection in the reasoner, “NumWarmups” refers to the number of warm-ups before starting the training, “NumRollouts” refers to the number of roll-outs in MCTS, “RwdDecay” refers to the constant exponential decay applied
on the reward, “c1 and c2” refers to the same constants in PUCT formula in Section A.1, “RBsize” refers to the replay buffer size that is used to count the number of stored trajectories, “BatchSize” refers to the batch size for training, “TrainingSteps” refers to the number of training steps. For PRISA-REINFORCE, PRISA-PPO, PRISA-MuZero, the “TrainingSteps” denotes the training steps for each task, while that of PRIMA represents the training steps for all 8 tasks.
PRISA- PRISA- PRISA- PRIMA REINFORCE PPO MuZero
lr-reasoner 0.005 0.005 0.005 0.004 lr-policy 0.085 0.085 0.075 0.075 lr-value - 0.15 0.075 0.075 breadth 3 3 3 3 depth 4 4 4 4 residual False False False False NumWarmups - - 200 200 NumRollouts - - 1200 1200 RwdDecay 5 5 5 5 c1 - - 30 30 c2 - - 19652 19652 RBsize 16 32 400 400 BatchSize 16 16 16 16 TrainingSteps 7⇥ 104 7⇥ 104 7⇥ 104 39⇥ 104

heading: C ADDITIONAL EXPERIMENTS
text: C.1 REASONING COSTS OF PRISA
To demonstrate the improved reasoning efficiency in PRISA at the inference stage, three variants of PRISA (e.g., PRISA-REINFORCE, PRISA-PPO, and PRISA-MuZero) are compared with PRIMA in terms of FLOPs. As shown in Fig. 6, three variants of PRISA generally have lower reasoning complexity than PRIMA. The difference in FLOPs between PRIMA and PRISA-MuZero is primarily from the difference between multi-tasking setting and single-task setting.
C.2 PREDICATE DIMENSION ANALYSIS
Since the Reasoner only realizes a partial set of Horn clauses of FOL, the number of intermediate predicates will greatly determine the expressive power of the model. The performance of PRIMA degrades gracefully when the number of intermediate predicates decreases. This is shown in the additional experiments of examining the limiting performance of our PRIMA, where we examine the limiting performance of our PRIMA when the number of intermediate predicates at each layer decreases. The results in Table 6 show that the performance of PRIMA degrades gracefully when the number of intermediate predicates decreases; for example, it still performs reasonably well on most tasks even when the number of predicates becomes 4.
C.3 ANALYSIS OF MULTI-TASKING CAPABILITIES
Why/how do different tasks share blocks (of ops) as in Fig. 7? To quantitatively measure how different tasks reuses the “skills”, we show in Fig. 7, which is a heat-map showing the probability of overlapping about Ops reuse across different tasks, where red color denotes a high-level of Ops overlapping. Ops reuse factor is defined as min(Pa(Ij), Pb(Ij)) where Pa(Ij) refers to the probability of ops Ij(j 2 {0, . . . , 9}) being active for task a and Pb(Ij) denotes the probability of I j being active for task b. From Fig. 7, it can be inferred that some tasks share skills more than others. For example, in most of the actions, I8 is shared across different tasks.

heading: D ADDITIONAL RELATED WORK
text: This section provides complementary contents to Sec. 5.
Model-based symbolic planning Another large body of related work is symbolic planning (SP) (Van Harmelen et al., 2008; Hanheide et al., 2015; Chen et al., 2016; Khandelwal et al., 2017). In SP approaches, a planning agent carries prior symbolic knowledge of objects, properties, and how they are changed by executing actions in the dynamic system, represented in a formal, logic-based language such as PDDL (McDermott et al., 1998) or an action language (Gelfond & Lifschitz, 1998) that relates to logic programming under answer set semantics (answer set programming) (Lifschitz, 2008). The agent utilizes a symbolic planner, such as a PDDL planner FASTDOWNWARD (Helmert, 2006) or an answer set solver CLINGO (Gebser et al., 2012) to generate a sequence of actions based on its symbolic knowledge, executes the actions to achieve its goal. Compared with RL approaches, an SP agent does not require a large number of trial-and-error to behave reasonably well, yet requires predefined symbolic knowledge as the model prior. It should be noted that it remains questionable if PDDL-based methods can be applied directly to our problem, as there is no explicitly predefined prior knowledge under our problem setting, which is required by PDDL-based approaches. In contrast, we address the learning-to-reason problem via a complete data-driven deep reinforcement approach.
Multi-task RL Multi-task learning can help boost the performance of reinforcement learning, leading to multi-task reinforcement learning (Vithayathil Varghese & Mahmoud, 2020; Zhu et al., 2020). Some research (Wilson et al., 2007; Li et al., 2009; Lazaric & Ghavamzadeh, 2010; Calandriello et al., 2014; Andreas et al., 2017; Deshmukh et al., 2017; Saxe et al., 2017; Bräm et al., 2019; Vuong et al., 2019; Igl et al., 2020) has adapted the ideas behind multi-task learning to RL. For example, in (Bräm et al., 2019), a multi-task deep RL model based on attention can group tasks into sub-networks with state-level granularity. The idea of compression and distillation has been incorporated into multi-task RL as in (Parisotto et al., 2016; Rusu et al., 2016; Omidshafiei et al., 2017; Teh et al., 2017). For example, in (Parisotto et al., 2016), the proposed actor-mimic method combines deep reinforcement learning with model-compression techniques to train a policy network that can learn to act for multiple tasks. Other research in multi-task RL focuses on online and distributed settings (Bsat et al., 2017; Sharma & Ravindran, 2017; Sharma et al., 2018; Espeholt et al., 2018; Tutunov et al., 2018; Lin et al., 2019).
RL for logic reasoning Relational RL integrates RL with statistical relational learning and connects RL with classical AI for knowledge representation and reasoning. Most prominently, Džeroski et al. (2001) originally proposed relational RL, Tadepalli et al. (2004) surveyed relational RL, Guestrin et al. (2003) introduced relational MDPs, and Diuk et al. (2008) introduced objected-oriented MDPs (OO-MDPs). More recently, Battaglia et al. (2018) proposed to incorporate relational inductive bias, Zambaldi et al. (2018) proposed deep relational RL, Keramati et al. (2018) proposed strategic object-oriented RL, and some researchers have also adopted deep learning approaches for dealing with relations and/or reasoning, for example (Battaglia et al., 2016; Chen et al., 2018; Santoro et al., 2017; Santurkar et al., 2018; Palm et al., 2017; Xia et al., 2018; Yi et al., 2018).
Statistical relational reasoning Statistical relational reasoning (SRL) (Koller et al., 2007) often provides a better understanding of domains and predictive accuracy, but with more complex learning and inference processes. Unlike ILP, SRL takes a statistical and probabilistic learning perspective and extends the probabilistic formalism with relational aspects. Examples of SRL include text classification (Ganiz et al., 2010), recommendation systems (Cao, 2016), and wireless networks (Yoshida et al., 2020). SRL can be divided into two research branches: probabilistic relational models (PRMs) (Koller, 1999) and probabilistic logic models (PLMs) (Chen et al., 2008). PRMs start from probabilistic graphical models and then extend to relational aspects, while PLMs start from ILP and extend to probabilistic semantics. A related research area is graph relational reasoning, such as using graph neural networks (Zhou et al., 2020) to conduct reasoning. A discussion of this approach is beyond the scope of this paper, but we refer interested readers to (Chakrabarti & Faloutsos, 2006; Cook & Holder, 2006; Zhang et al., 2019) for a comprehensive review.
Comparisons with NLM in details The Reasoner shares the same functionality with NLM, in terms of approximating the meta-rules (BooleanLogic, Expansion, and Reduction). However,
our Reasoner has major differences from NLM. First, the targets are different. Our Reasoner targets the trade-off between capability and efficiency. Compared with NLM, our Reasoner can greatly improve the efficiency by using a Planner module. Second, our Reasoner has more explicitness in the reasoning path. Let [Or 1i 1 , O r i 1, O r+1 i 1 ] denote the output predicates (in tensor representation) from the previous layer, and Concate(·) denotes the concatenation operation. NLM performs the intra-group computation as
O r i = (MLP (Permute (Z r i ) ; ✓ r i )) ,
where Ori is the output predicate, Zri = Concat Expand O r 1 i 1 , O r i 1,Reduce O r+1 i 1 , is the sigmoid nonlinearity and ✓ri denotes learnable parameters. On the contrary, in the Reasoner module of our PRISA/PRIMA, the implementation is
O r i =
MLP(Permute Expand(Or 1i 1 ) ; ✓ri ) +MLP(Permute O r i 1); ✓ r i
+MLP(Permute Reduce(Or+1i 1 )); ✓ r i
! ,
which allows for the neural implementations of BooleanLogic, Expansion, and Reduction (at the same arity group) to be executed independently to some extent and reduce the unnecessary computations.
PRIMA: PLANNER-REASONER INSIDE A MULTI-TASK REASONING AGENT
Anonymous authors Paper under double-blind review
ABSTRACT
We consider the problem of multi-task reasoning (MTR), where an agent can solve multiple tasks via (first-order) logic reasoning. This capability is essential for human-like intelligence due to its strong generalizability and simplicity for handling multiple tasks. However, a major challenge in developing effective MTR is the intrinsic conflict between reasoning capability and efficiency. An MTR-capable agent must master a large set of “skills” to tackle diverse tasks, but executing a particular task at the inference stage requires only a small subset of immediately relevant skills. How can we maintain broad reasoning capability but efficient specific-task performance? To address this problem, we propose a Planner-Reasoner framework capable of state-of-the-art MTR capability and high efficiency. The Reasoner models shareable (first-order) logic deduction rules, from which the Planner selects a subset to compose into efficient reasoning paths. The entire model is trained in an end-to-end manner using deep reinforcement learning, and experimental studies over a variety of domains validate its effectiveness.

heading: 1 INTRODUCTION
text: 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 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 3), wherein the reasoner defines a set of neural logic operators for modeling reusable reasoning meta-rules (“skills”) across tasks (Section 3.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 3.3). Thus, our planner-reasoner architecture features the dual capabilities of composing and gpruning a logic deduction process, achieving a graceful capability-efficiency trade-off in MTR. The model architecture is trained in an end-toend manner using deep reinforcement learning (Section 4), and experimental results on several benchmarks demonstrate that this framework leads to a state-of-the-art balance between capability and efficency (Section 5). We discuss related works in Section 6, and conclude our paper in Section 7.

heading: 2 PROBLEM FORMULATION
text: 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.

heading: 3 PRIMA: PLANNER-REASONER FOR MULTI-TASK REASONING
text: In this section, we develop our PRIMA framework, which is a novel neural-logic model architecture for multi-task reasoning. We begin with an introduction to the overall architecture (Section 3.1) along with its key design insights, and then proceed to discuss its two modules in Sections 3.2–3.3.

heading: 3.1 THE OVERALL MODEL ARCHITECTURE OF PRIMA
text: Logic reasoning typically seeks to chain an appropriate sequence of logic operators to reach desirable logical consequences from premises. There are two types of chaining to conduct logic reasoning: (i) forward-chaining and (ii) backward-chaining. The forward-chaining strategy recursively deduces all the possible conclusions based on all the available facts and deduction rules until it reaches the answer. In contrast, backward-chaining starts from the desired conclusion (i.e., the goal) and then works backward recursively to validate conditions through available facts. This paper adopts the forward-chaining strategy since the goal is unavailable here beforehand. Note that, in the MTR setting, it generally requires a large set of “skills” (logic operators) to tackle diverse tasks. Therefore, it is unrealistic to recursively generate all possible conclusions by chaining all the available deduction rules as in the vanilla forward-chaining strategy. Therefore, pruning the forward-chaining process is essential to improve the reasoning efficiency in the MTR setting.
We note that reasoning via pruned forward-chaining could be implemented by: (i) constructing the elementary logic operators and (ii) selecting and chaining a subset of these logic operators together. Accordingly, we develop a novel neural model architecture, named Planner-Reasoner (see Figure 2A), to jointly accomplish these two missions with two interacting neural modules. Specifically, a “Reasoner” defines a set of neural-logic operators (representing learnable Horn clauses (Horn, 1951)). Meanwhile, a “Planner” learns to select and chain a small subset of these logic operators together towards the correct solution (e.g., the orange-colored path in Figure 2A). The inputs to the Planner-Reasoner model are the base predicates {p, p(x), p(x, y), . . .} that describe the premises. In particular, in the MTR setting, the nullary predicate p characterizes the current task to be solved. This design allows the Planner to condition its decisions on the input task type. We feed the outputs (i.e., the conclusion predicates) into multiple output heads, one for each task (just as in standard multitask learning), to generate predictions. The entire model is trained in a complete data-driven end-to-end approach via deep reinforcement learning (Figure 2B and Section 4): it trains the Planner to dynamically prune the forward-chaining process while jointly learning the logic operators in the Reasoner, leading to a state-of-the-art balance between MTR capability and efficiency (Section 5).

heading: 3.2 REASONER: TRANSFORMING LOGIC RULES INTO NEURAL OPERATORS
text: 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).

heading: 3.3 PLANNER: ACTIVATING AND FORWARD-CHAINING LEARNABLE HORN CLAUSES
text: The Planner is our key module to address the capability-efficiency tradeoff in the MTR problem; it is responsible for activating the neural operators in the Reasoner and chaining them into reasoning paths. Existing learning-to-reason approaches, which are often based on inductive logic programming (ILP) (Cropper et al., 2020; Cropper & Dumancic, 2020; Muggleton & De Raedt, 1994) and the correspondingly neural-ILP methods (Dong et al., 2019; Shi et al., 2020). Conventional ILP methods suffer from several drawbacks, such as heavy reliance on human-crafted templates and sensitivity to noise. On the other hand, neural-ILP methods (Dong et al., 2019; Shi et al., 2020), leveraging the strength of deep learning and ILP, such as the Neural Logic Machine (NLM) (Dong et al., 2019), lack explicitness in the learning process of searching for reasoning paths. Let us take the learning process of NLM for example, which follows an intuitive two-step procedure. It first fully connects all the neural blocks and then searches all possible connections (corresponding to all possible predicate candidates) exhaustively to identify the desired reasoning path (corresponding to the desired predicate).
By using our proposed Planner module, we can strike a better capability-efficiency tradeoff. Rather than conducting an exhaustive search over all possible predicate candidates as in NLM, the Planner prunes all the unnecessary operators and identifies an essential reasoning path with low complexity for a given problem. Consider the following example. As shown in Fig. 1 and Fig. 2A (the highlighted orange-colored blocks), at each reasoning step, the Planner takes the input predicates and determines which neural operators should be activated. The decision is represented as a binary vector — [I0 . . . I6] (termed operator footprint) — that corresponds to the neural operators of Expand, Reduce and DirectInput1 at different arity groups. By chaining these sequential decisions, a sequence of operator footprints is formulated as a reasoning path, as the highlighted orange-colored path in Fig. 2A. Generally, the neural operators defined in the Reasoner can also be viewed as (learnable) Horn Clauses (Horn, 1951; Chandra & Harel, 1985), and the Planner forward-chains them into a reasoning path.

heading: 4 LEARNING-TO-REASON VIA REINFORCEMENT LEARNING
text: In our problem, the decision variables of the Planner are binary indicators of whether a neural operator module should be activated. Readers familiar with Markov decision processes (MDPs) (Puterman,
1DirectInput is an identity mapping of the inputs where its following operators of Permute, MLP and sigmoid nonlinearity can directly approximate the BooleanLogic.
2014) might notice that the reasoning path of our model in Fig. 2A resembles a temporal rollout in MDP formulation (Sutton & Barto, 2018). Therefore, we frame learning-to-reason as a sequential decision-making problem and adopt off-the-shelf reinforcement learning (RL) algorithms.

heading: 4.1 AN MDP FORMULATION OF THE LEARNING-TO-REASON PROBLEM
text: An MDP is defined as the tuple (S,A, P ass0 , R, ), where S and A are finite sets of states and actions, the transition kernel P ass0 specifies the probability of transition from state s 2 S to state s0 2 S by taking action a 2 A, R(s, a) : S ⇥A! R is the reward function, and 0   1 is a discount factor. A stationary policy ⇡ : S ⇥A! [0, 1] is a probabilistic mapping from states to actions. The primary objective of an RL algorithm is to identify a near-optimal policy ⇡⇤ that satisfies
⇡ ⇤ := argmax
⇡
⇢ J(⇡) := E  Tmax 1X
t=0
t r(st, at ⇠ ⇡) ,
where Tmax is a positive integer denoting the horizon length — that is, the maximum length of a rollout. The resemblance between Fig. 2 and an MDP formulation is relatively intuitive as summarized in Table 1. At the t-th time step, the state st corresponds to the set of predicates, st = [O0t , O1t , O2t ], with the superscript denoting the corresponding arity. The action at = [I0t , I1t , . . . , I K 1 t ] (e.g., a0 = [000011]) is a binary vector that indicates the activated neural operators (i.e., the operator footprint), where K is the number of operators per layer. The reward rt is defined to be
rt :=
( PK 1
i=0 I i t , if t < Tmax
Accuracy, t = Tmax . (1)
That is, the terminal reward is set to be the reasoning accuracy at the end of the reasoning path (see Appendix A.2 for its definition), and the intermediate reward at each step is chosen to be the negated number of activated operators (which penalizes the cost of performing the current reasoning step). The transition kernel P ass0 corresponds to the function modeled by one Reasoner layer; each Reasoner layer will take the state (predicates) st 1 = [O0t 1, O1t 1, O2t 1, ] and the action (operator footprint) at = [I0t , I 1 t , . . . , I K 1 t ] as its input and then generate the next state (predicates) st = [O0t , O1t , O2t , ]. This also implies that the Reasoner layer defines a deterministic transition kernel, i.e., given st 1 and at the next state st is determined.

heading: 4.2 POLICY NETWORK: MODELING OF PLANNER
text: The Planner module is embodied in the policy network. As shown in Fig. 3A, the base planner is a separate module that has the same architecture of 1-layer (fully-activated) Reasoner followed by a max-pooling layer. This architecture enables the reduction of input predicates to the specific indicators by reflecting whether the operations at the corresponding position of one layer of Reasoner are active or inactive. Further, we can also leverage Monte-Carlo Tree Search (MCTS) (Browne et al., 2012; Munos, 2014) to boost the performance, which leads to an Enhanced Planner (Fig. 3B). An MCTS algorithm such as the Upper Confidence Bound for Trees (UCT) method (Kocsis et al., 2006), is a model-based RL algorithm that plans the best action at each time step (Browne et al., 2012) by constructing a search tree, with states as nodes and actions as edges. The Enhanced Planner uses MCTS to exploit partial knowledge of the problem structure (i.e., the deterministic transition kernel P ass0 defined by the Reasoner layer) and construct a search tree to help identify the best actions (which ops to activate). Details of the MCTS algorithms used in the Enhanced Planner can be found in Appendix A.1.

heading: 4.3 OVERALL LEARNING FRAMEWORK
text: As illustrated in Fig. 3, we introduce concrete data-driven decision-making methods—that is, RL approaches (Sutton & Barto, 2018)—to address the learning-to-reason problem. To illustrate this, we apply the model-free RL method REINFORCE (Williams, 1992) and the model-based method MuZero (Schrittwieser et al., 2019). Compared with model-free reinforcement learning, model-based reinforcement learning (MBRL) more effectively handles large search-space problems, such as the game of Go (Silver et al., 2017b;b;a). MuZero (Schrittwieser et al., 2019), a recently proposed MBRL approach to integrating planning and learning, has achieved great success with a variety of complex tasks. Motivated by the success of MuZero, we propose an MBRL approach for neural-symbolic reasoning. The key insight behind adopting MuZero is that in real applications, we typically have partial structural knowledge of the transition kernel P ass0 and reward function r(s, a). As a result of the model-based module, testing complexity can be greatly reduced by adopting MCTS, which leads to a better set of predicates. Of course, MuZero is just one option in the model-based family of approaches. We leave it as future research to propose and compare other model-based alternatives.
The pipeline of the training process is illustrated in Fig.2.B (the right subfigure). After loading the model weights, the reasoning path rollouts are executed by the agent (or model instance), according to the current policy network. The performed reasoning path rollout is then stored in the replay buffer. The Planner-Reasoner is trained via rollouts sampled from the replay buffer.

heading: 5 EXPERIMENTAL RESULTS AND ANALYSIS
text: In this section, we evaluate the performance of different variants of PRIMA on eight tasks from the family tree and graph benchmarks (Graves et al., 2016), including 1-Outdegree, AdjacentToRed, HasFather, HasSister, 4-Connectivity, IsGrandparent, IsUncle, IsMGUncle. These tasks are widely used benchmarks for inductive logic programming (Krötzsch, 2020; Calautti et al., 2015). Detailed descriptions about those tasks can be found in Appendix B.1. We evaluate their testing accuracy and reasoning cost (measured in FLOPs: the number of floating-point operations executed (Clark et al., 2020)) on these tasks and compare them to several baselines. Furthermore, detailed case studies are conducted on the reasoning path, which indicates the operator sharing among different tasks. All the results demonstrate the graceful capability-efficiency tradeoff of PRIMA in multi-task reasoning.

heading: 5.1 EXPERIMENTAL SETUPS
text: In the multi-task setting, we first randomly sample a task according to a pre-defined probability distribution (over different tasks), and then generate the data for the selected task using the same methods as in NLM (Dong et al., 2019). In addition, we also augment the generated data with one-hot encoding to indicate the selected task, which will be further converted into nullary (background) input predicates. Also, task-specific output heads are introduced to generate outputs for different tasks. These adaptions apply to NLM-MTR, DLM-MTR, and PRIMA. The reasoning accuracy is
used as the reward. In the inference (or testing) stage, the Reasoner is combined with the learned Base Planner to perform the tasks, instead of an enhanced planner (MCTS), to reduce the extra computation. The problem size for training is always 10 for graph tasks and 20 for family tree tasks across all settings, regardless of the sizes of testing problems.

heading: 5.2 OVERALL PERFORMANCE
text: Single-task reasoning capability Before we proceed to evaluate the MTR capabilities of PRIMA, we first adapt it to the single-task setting by letting the nullary input predicates to be zero, and then train a separate model for each individual task. We name this single-task version of PRIMA as PRISA (i.e., Planner-Reasoner Inside a Single Reasoning Agent), and compare it with existing (single-task) neural-logic reasoning approaches (e.g., NLM (Dong et al., 2019), DLM (Matthieu et al., 2021), and MemNN (Sukhbaatar et al., 2015)). First, we compare the performance of three different variants of PRISA (Section 4.3) for single-task reasoning, which learns their planners based on different reinforcement learning algorithms; PRISA-REINFORCE uses REINFORCE (Williams, 1992), PRISA-PPO uses PPO (Schulman et al., 2017), and PRISA-MuZero uses MuZero (Schrittwieser et al., 2019). We report the test accuracy and the Percentage of Successful Seeds (PSS) in Table 2 to measure the model’s reasoning capabilities, where the PSS reaches 100% of success rates (Matthieu et al., 2021). We note that PRISA-MuZero has the same 100% accuracy as NLM (Dong et al., 2019), DLM (Matthieu et al., 2021), and @ILP (Evans & Grefenstette, 2018) across different tasks, and outperforms MemNN (Sukhbaatar et al., 2015) (shown in the single-task part in Table 3). But it also has a higher successful percentage (PSS) in comparison with other methods. The results show that PRISA-MuZero achieves the best performance on single-task reasoning, which confirms the strength of the MCTS-based Enhanced Planner (Section 4.3) and the MuZero learning strategy. Therefore, we will use the Enhanced Planner and MuZero in PRISA and PRIMA in the rest of our empirical studies.
Multi-task reasoning capability Next, we evaluate the MTR capabilities of PRIMA. To the best of our knowledge, there is no existing approach that is designed specifically for MTR. Therefore, we adapt NLM and DLM into their multi-task versions, named NLM-MTR and DLM-MTR, respectively. NLM-MTR and DLM-MTR follow the same input and output modification as what we did to upgrade PRISA to PRIMA (Section 3.1). By this, we can examine the contribution of our proposed Planner-Reasoner architecture for MTR. As shown in Table 3, PRIMA (with MuZero as the Base Planner) performs well (perfectly) on different reasoning tasks. On the other hand, DLM-MTR experiences some performance degradation (on AdjacentToRed). This result confirms that our Planner-Reasoner architecture is more suitable for MTR. We conjecture that the benefit comes from using a Planner to explicitly select the necessary neural operators for each task, avoiding potential conflicts between different tasks during the learning process.
Experiments are also conducted to test the performance of PRIMA with different problem sizes. The problem size in training is 10 for all graph tasks and 20 for all family-tree tasks. In testing, we evaluate the methods on much larger problem sizes (50 for the graph tasks and 100 for the family tree tasks), which the methods have not seen before. Therefore, the Planner must dynamically activate a proper set of neural operators to construct a path to solve the new problem. As reported in Fig. 4 and Tables 2 and 3, PRIMA can achieve the best accuracy and lower flops when the problem sizes for training and testing are different.
Reasoning efficiency To measure the reasoning efficiency of the proposed methods at the inference stage, PRIMA is compared with NLM, MemNN, and NLM-MTR in terms of FLOPs. As shown in
Table 3: Testing Accuracy and PSS of PRIMA and other baselines on different reasoning tasks. The results of @ILP, NLM, and DLM are merged in one row due to space constraints and are presented in the same order if the
results are different. Note that PRIMA’s test efficiency is superior to NLM-MTR’s as shown in Fig. 4. “m”: the problem size. “PSS”: Percentage of Successful Seeds. Numbers in red denote < 100%.
te st
in g
ac cu
ra cy
Family Tree HasFather HasSister IsGrandparent IsUncle IsMGUncle Graph AdjacentToRed 4-Connectivity 1-OutDegree
Si ng
le Ta
sk MemNN m=20 99.9 86.3 96.5 96.3 99.7 m=10 95.2 92.3 99.8 m=100 59.8 59.8 97.7 96 98.4 m=50 93.1 81.3 78.6
PSS 0 0 0 0 0 PSS 0 0 0
@ILP/ NLM/ DLM m=20 100 100 100 100 100 m=10 100 100 100 m=100 100 100 100 100 100 m=50 100 100 100
PSS 100 100 100 100/ 90/ 100 100/ 20/ 70 PSS 100/ 90/ 90 100 100
M ul
tiTa
sk
NLM-MTR m=20 100 100 100 100 100 m=10 100 100 100
m=100 100 100 100 100 100 m=50 100 100 100 PSS 100 100 100 100 90 PSS 90 100 100
DLM-MTR m=20 100 100 100 100 100 m=10 96.7 100 100
m=100 100 100 100 100 100 m=50 97.2 100 100 PSS 100 100 100 100 100 PSS 0 100 100
PRIMA m=20 100 100 100 100 100 m=10 100 100 100
m=100 100 100 100 100 100 m=50 100 100 100 PSS 100 100 100 100 90 PSS 90 100 100
Fig. 4, NLM and NLM-MTR demonstrate a similar performance and suffer the highest reasoning cost when it is tested with large problem sizes, such as 50 for graph tasks and 100 for family tree tasks. For MemNN, although the FLOPs of it seem low in most cases of testing, its testing accuracy is bad and cannot achieve accurate predictions (Table 3). In contrast, PRIMA can significantly reduce the reasoning complexity by intelligently selecting ops using a planner. Overall, PRIMA strikes a satisfactory capability-efficiency tradeoff in comparison with all available multi-tasking baselines.
Operator/Path sharing in MTR To take a closer look into how PRIMA achieves such a better capability-efficiency tradeoff (in Fig. 4), we examine the reasoning paths on three different graph tasks: 1-Outdegree, AdjacentToRed, and 4-Connectivity. Specifically, we sample instances from these three tasks and feed them into PRIMA separately to generate their corresponding reasoning paths. The results are plotted in Fig. 5A, where the gray paths denote the ones shared across tasks, and the colored ones are task-specific. It clearly shows that PRIMA learns a large set of neural operators sharable across tasks. Given each input instance from a particular task, PRIMA activates a set of shared paths along with a few task-specific paths to deduce the logical consequences.
Generalizability of the Planner To demonstrate the generalizability of the Planner module in PRIMA, we generate input instances of AdjacentToRed with topologies that have not been seen during training. In Fig. 5B, we show the reasoning paths activated by the Planner for these different topologies, which demonstrates that different input instances share a large portion of reasoning paths. This fact is not surprising as solving the same task of AdjacentToRed should rely on a common set of skills. More interestingly, we notice that even for solving this same task, the Planner will have to call upon a few instance-dependent sub-paths to handle the subtle inherent differences (e.g., the graph topology) that exist between different input instances. For this reason, PRIMA maintains an instance-dependent dynamic architecture, which is in sharp contrast to the Neural Architecture Search (NAS) approaches. Although NAS may also use RL algorithms to seek for a smaller architecture (Zoph & Le, 2017), it only searches for a static architecture that will be applied to all the input instances.
most cases of testing, its testing accuracy is bad and cannot achieve accurate predictions (see Table 3). “OD” denotes 1-Outdegree, likewise, “AR”:AdjacentToRed, “4C”:4-Connectivity, “HF”:HasFather, “HS”:HasSister, “GP”:IsGrandparent, “UN”:IsUncle, “MG”:IsMGUncle.

heading: 6 RELATED WORK
text: Multi-task learning Multi-task Learning (Zhang & Yang, 2021; Zhou et al., 2011; Pan & Yang, 2009) has focused primarily on the supervised learning paradigm, which itself can be divided into several approaches. The first is feature-based MTL, also called multi-task feature learning, (Argyriou et al., 2008), which explores the sharing of features across different tasks using regularization techniques (Shinohara, 2016; Liu et al., 2017). The second approach assumes that tasks are intrinsically related, such as low-rank learning (Ando & Zhang, 2005; Zhang et al., 2005), learning with task clustering (Gu et al., 2011; Zhang et al., 2016), and task-relation learning (such as via task similarity, task correlation, or task covariance) (Goncalves et al., 2016; Zhang, 2013; Ciliberto et al., 2015). MTL has also been explored under other paradigms such as unsupervised learning, for example, via multi-task clustering in (Zhang & Zhang, 2013; Zhang et al., 2015b; Gu et al., 2011; Zhang, 2015).
Neural-symbolic reasoning Neural-symbolic AI for reasoning and inference has a long history (Besold et al., 2017; Bader et al., 2004; Garcez et al., 2008), and neural-ILP has developed primarily in two major directions (Cropper et al., 2020). The first direction applies neural operators such as tensor calculus to simulate logic reasoning (Yang et al., 2017; Dong et al., 2019; Shi et al., 2020) and (Manhaeve et al., 2018). The second direction involves relaxed subset selection (Evans & Grefenstette, 2018; Si et al., 2019) with a predefined set of task-specific logic clauses. This approach reduces the task to a subset-selection problem by selecting a subset of clauses from the predefined set and using neural networks to search for a relaxed solution.
Our work is different from the most recent work, such as (Dong et al., 2019; Shi et al., 2020), in several ways. Compared with (Dong et al., 2019), the most notable difference is the improvement in efficiency by introducing learning-to-reason via reinforcement learning. A second major difference is that PRIMA offers more generalizability than NLM by decomposing the logic operators into more finely-grained units. We refer readers for more detailed related work to Appendix D.

heading: 7 CONCLUSION
text: A long-standing challenge in multi-task reasoning (MTR) is the intrinsic conflict between capability and efficiency. Our main contribution is the development of a novel neural-logic architecture termed PRIMA, which learns to perform efficient multitask reasoning (in first-order logic) by dynamically chaining learnable Horn clauses (represented by the neural logic operators). PRIMA improves inference efficiency by dynamically pruning unnecessary neural logic operators, and achieves a stateof-the-art balance between MTR capability and inference efficiency. The model’s training follows a complete data-driven end-to-end approach via deep reinforcement learning, and the performance is validated across a variety of benchmarks. Future work could include extending the framework to high-order logic and investigating scenarios when meta-rules have a hierarchical structure.
REPRODUCIBILITY STATEMENT
We commit to ensuring that other researchers with reasonable background knowledge in our area can reproduce our theoretical and empirical results. The algorithm details can be seen in Appendix A. Specifically, benchmark details are provided in Appendix B.1, hyper-parameter settings are provided in Appendix B.2, and computing infrastructure are detailed in Appendix B.3. The experiment details can be seen in Appendix B.
ETHICS STATEMENT
This work is about the methodology of achieving a capability-efficiency trade-off in multi-task first-order logic reasoning. The potential impact of this work is likely to further extend the framework to high-order logic and investigate scenarios when meta-rules have a hierarchical structure, which should be generally beneficial to the multi-task learning research community. We have not considered specific applications or practical scenarios as the goal of this work. Hence, it does not have any direct ethical consequences.

Review summaries:
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.

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.

The paper presents PRIMA, a framework designed to solve the problem of multi-task reasoning. PRIMA consists of two parts: a reasoner, which produces logical deductions using a set of neural operators built by translating FOL rules; and a planner, which activates the reasoner operators and creates reasoning paths to solve the task. The paper presents several tests that achieve performance usually stable on 100% succesful seeds.