PRIMA: P LANNER -R EASONER I NSIDE A M ULTI - TASK R EASONING A GENT
Anonymous authors
Paper under double-blind review
A BSTRACT
We consider the problem of multi-task reasoning (MTR), where an agent can solve multiple tasks via (ﬁrst-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 conﬂict between reasoning capability and efﬁciency.
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 efﬁcient speciﬁc-task performance?
To address this problem, we propose a
Planner-Reasoner framework capable of state-of-the-art MTR capability and high efﬁciency.
The Reasoner models shareable (ﬁrst-order) logic deduction rules, from which the Planner selects a subset to compose into efﬁcient 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.
I NTRODUCTION
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-speciﬁc 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 signiﬁcantly larger than a single model for tackling a speciﬁc task.
This results from the intrinsic conﬂict underlying all MTL algorithms: balancing across-task generalization capability to perform different tasks with single-task efﬁciency in executing a speciﬁc 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 efﬁcient solution for a speciﬁc task.
This conﬂict 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 conﬂict between across-task generalization capability and single-task efﬁciency .
To strike a balance between reasoning capability and efﬁciency in MTR, we develop a P lannerR easoner architecture I nside a M ulti-task reasoning
A gent (PRIMA) (Section 2), wherein the reasoner deﬁnes a set of neural logic operators for modeling reusable reasoning meta-rules (“skills”) across tasks (Section 2.2).
When deﬁning the logic operators, we focus on ﬁrst-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-efﬁciency 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.
P LANNER -R EASONER FOR M ULTI - TASK R EASONING
This section proposes the Planner-Reasoner framework for MTR.
To that end, we ﬁrst formally state the logic reasoning problem in Section 2.1, and then in Sections 2.2 and 2.3, we describe the P lanner- R easoner I nside a S ingle Reasoning A gent (PRISA) framework.
This is a plannerreasoner architecture, which is a neural-logic architecture for traversing the ﬁrst-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 P lanner- R easoner I nside a
M ulti-task Reasoning A gent (PRIMA) framework.
2. P ROBLEM F ORMULATION
Logic reasoning We begin with a brief introduction of the logic reasoning problem.
Speciﬁcally, 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 ( x 1 , . . .
, x r ) can be considered as a relation between r constants, which takes the value of True or False .
An atom p ( x 1 , · · · , x r ) is an r -ary predicate with its arguments x 1 , · · · , x r being either variables or constants.
A well-deﬁned 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 ∃ and universal ∀ quantiﬁers according to certain formation rules (see Andrews (2002) for the details).
In particular, the quantiﬁers ∃ and ∀ 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.
Speciﬁcally, 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 , . . .
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 deﬁnition of MTR.
With a slight abuse of notations, let { p ( x 1 , . . .
, x r ) : r ∈ [1 , n ] } be the set of input predicates sampled from any of the k different reasoning tasks, where x 1 , . . .
, x r are the individual variables and n is the maximum arity.
A multi-task reasoning model takes p ( x 1 , . . .
, x r ) as its input and seeks to predict the corresponding ground-truth output predicates q ( x 1 , . . .
, x r ) .
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. R EASONER : TRANSFORMING LOGIC RULES INTO NEURAL OPERATORS
The Reasoner module conducts logical deduction using a set of neural operators constructed from ﬁrst-order logic rules (more speciﬁcally, 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 .
[Equation]
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 ψ ∈ S n , we deﬁne p ψ ( x 1 , · · · , x r ) =
p ( x ψ (1) , · · · , x ψ ( r ) ) where S n is the set of all possible permutations as the arguments to an input predicate.
The corresponding neural implementation of BooleanLogic is σ (MLP ( p ψ ( x 1 , · · · , x r )) ; θ ) , where σ is the sigmoid activation function, MLP refers to a multilayer perceptron, a Permute( · ) neural operator transforms input predicates to p ψ ( x 1 , · · · , x r ) , and θ is the learnable parameter within the model.
This is similar to the implicit Horn clause with the universal quantiﬁer( ∀ ), e.g., p 1 ( x ) ∧ p 2 ( x ) → ˆ p ( x ) implicitly denoting ∀ x p 1 ( x ) ∧ p 2 ( 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 quantiﬁcation that bridge predicates of different arities with logic quantiﬁers ( ∀ and ∃ ).
Expansion introduces a new and distinct variable x r +1 for a set of r -ary predicates with the universal quantiﬁer( ∀ ).
For this reason, Expansion creates a new predicate q from p .
Expansion :
[Equation]
where x r +1 / ∈ { x i } 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 x r + in a set of ( r + 1) -ary predicates via the quantiﬁers of ∀ or ∃ .
Reduction : ∀ x r +1 p ( x 1 , x 2 , · · · , x r , x r +1 )
→ q ( x 1 , x 2 , · · · , x r ) , or ∃ x r +1 p ( x 1 , x 2 , · · · , x r , x r +1 )
→ q ( x 1 , x 2 , · · · , x r ) .
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 x r +1 due to the universal quantiﬁer ∀ (resp.
existential quantiﬁer ∃ ).
2. P LANNER : ACTIVATING AND FORWARD - CHAINING LEARNABLE H ORN CLAUSES
The Planner is our key module to address the capability-efﬁciency 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., 2020a; 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 ﬁrst 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-efﬁciency tradeoff.
Rather than conducting an exhaustive search over all possible predicate candidates as in NLM, the Planner prunes all the unnecessary operators and identiﬁes 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 — [ I 0 . . .
I 6 ] (termed operator footprint ) — that corresponds to the neural operators of
Expand , Reduce and DirectInput 2 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 deﬁned 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.
2. T HE PRIMA F RAMEWORK FOR MULTI - TASK REASONING
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 ( O 0 in
Fig. 2) to inform the planner of the current task.
Based on the value of O 0 , the planner will learn to activate the appropriate subset of neural operators, including a subset of shared operators and a few task-speciﬁc 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.
L EARNING - TO -R EASON VIA R EINFORCEMENT L EARNING
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.
3. A N MDP FORMULATION OF THE LEARNING - TO - REASON PROBLEM
An MDP is deﬁned as the tuple ( S , A , P ass (cid:48) , R, γ ) , where S and A are ﬁnite sets of states and actions, the transition kernel P ass (cid:48) speciﬁes the probability of transition from state s ∈ S to state s (cid:48) ∈ S by taking action a ∈ 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 satisﬁes
[Equation]
where T max 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 s t corresponds to the set of predicates, s t =
[ O 0 t , O 1 t , O 2 t ] , with the superscript denoting the corresponding arity.
The action a t =
[ I 0 t , I 1 t , . .
, I K − 1 t ] (e.g., a 0 = [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 r t is deﬁned to be
[Equation]
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 deﬁnition), 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 a ss (cid:48) corresponds to the function modeled by one Reasoner layer; each Reasoner layer will take the state (predicates) s t −
=
[ O 0 t − 1 , O 1 t − 1 , O 2 t − 1 , ] and the action (operator footprint) a t =
[ I 0 t , I 1 t , . .
, I K − 1 t ] as its input and then generate the next state (predicates) s t
= [ O 0 t , O 1 t , O 2 t , ] .
This also implies that the Reasoner layer deﬁnes a deterministic transition kernel, i.e., given s t − 1 and a t the next state s t is determined.
[Table]
3. P OLICY NETWORK : MODELING OF PLANNER
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 speciﬁc indicators by reﬂecting 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 Conﬁdence 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 ass (cid:48) deﬁned 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.
3. O VERALL LEARNING FRAMEWORK
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 ass (cid:48) 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 subﬁgure).
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.
E XPERIMENTAL R ESULTS AND A NALYSIS
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 ﬂoating-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-efﬁciency tradeoff of PRIMA in multi-task reasoning.
4. E XPERIMENTAL S ETUPS
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-deﬁned 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-speciﬁc 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
[Table]
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
[Table]
4. O VERALL PERFORMANCE
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 conﬁrms 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 speciﬁcally 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 modiﬁcation 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 conﬁrms that
our
Planner-Reasoner architecture is more suitable for MTR.
We conjecture that the beneﬁt comes from using a Planner to explicitly select the necessary neural operators for each task, avoiding potential conﬂicts 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 ﬂops when the problem sizes for training and testing are different.
Reasoning efﬁciency To measure the reasoning efﬁciency 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 signiﬁcantly reduce the reasoning complexity by intelligently selecting ops using a planner.
Overall, PRIMA strikes a satisfactory capability-efﬁciency tradeoff in comparison with all available multi-tasking baselines.
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Operator/Path Sharing in MTR To take a closer look into how PRIMA achieves such a better capability-efﬁciency tradeoff (in Fig. 4), we examine the reasoning paths on three different graph tasks: 1-Outdegree , AdjacentToRed , and 4-Connectivity .
Speciﬁcally, 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-speciﬁc.
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-speciﬁc 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.
inference stage
Compared to NLM and NLM-MTR, our PRIMA signiﬁcantly reduces the reasoning complexity by intelligently selecting ops (short for neural operators) using a planner.
Although the FLOPs of MemNN seem low in most cases of testing, its testing accuracy is bad and cannot achieve accurate predictions (see Table 3)
R ELATED W ORK
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 ﬁrst 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., 2020b).
The ﬁrst 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 predeﬁned set of task-speciﬁc logic clauses.
This approach reduces the task to a subset-selection problem by selecting a subset of clauses from the predeﬁned 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 efﬁciency 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 ﬁnely-grained units.
We refer readers for more detailed related work to Appendix D.
C ONCLUSION
A long-standing challenge in multi-task learning is the intrinsic conﬂict between capability and efﬁciency.
In this paper, we propose a Planner-Reasoner framework for multi-task reasoning.
To maintain broad capability but efﬁcient speciﬁc performance, the Reasoner extracts shareable metarules via ﬁrst-order logic, from which the Planner efﬁciently 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.
R EPRODUCIBILITY S TATEMENT
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.
Speciﬁcally, 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.
E THICS S TATEMENT
This work is about the methodology of achieving a capability-efﬁciency trade-off in multi-task ﬁrst-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 beneﬁcial to the multi-task learning research community.
We have not considered speciﬁc applications or practical scenarios as the goal of this work.
Hence, it does not have any direct ethical consequences.
A. D ETAILS 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) a s 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
[Equation]
where C ( s ) denotes the set of all child nodes for s , the ﬁrst term V s (cid:48) is an estimate of the longterm cumulative reward that can be received when starting from the state represented by node s (cid:48) , and the second term represents the uncertainty (conﬁdence-interval size) of that estimate.
The conﬁdence interval is calculated based on the upper conﬁdence bound (UCB) (Auer et al., 2002; Auer, 2002) using N s and N s (cid:48) , which denote the number of times that nodes s and s (cid:48) have been visited (respectively).
The key idea of UCT policy (2) is to select the best action according to an optimistic estimation (the upper conﬁdence bound) of the expected return, which balances the exploitation (ﬁrst 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 V s and N s from the leaf node s T to the root node s 0 of the selected path by recursively performing the following update (i.e., from t =
T − 1 to t = 0 ):
[Equation]
where ˆ V s T is the simulation return of s T , and a t denotes the action selected following (2) at state s t .
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 conﬁdence 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
[Equation]
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 c 1 , c 2 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 satisﬁed) by choosing actions (edges) a s 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. D ETAILS OF C ALCULATING R EASONING A CCURACY
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) ).
B E XPERIMENTAL D ETAILS B. R EASONING T ASKS I NTRODUCTION
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.
Speciﬁcally, 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.
Speciﬁcally, 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. H YPER - PARAMETER S ETTINGS
For simplicity, we set discount factor γ = 1 in all experiments.
The T max 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 .
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 .
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 .
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. C OMPUTING INFRASTRUCTURE
We conducted our experiments on a CPU server where the CPU is “Intel(R) Xeon(R)
Silver
CPU” with 40 cores and 64 GB memory in total.
[Table] [Table]
[Table]
C A DDITIONAL E XPERIMENTS C. R EASONING COSTS OF PRISA
To demonstrate the improved reasoning efﬁciency 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. P REDICATE 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. A NALYSIS OF MULTI - TASKING CAPABILITIES
Why/how do different tasks share blocks (of ops) as in Fig.
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 deﬁned as min( P a ( I j ) ,
P b ( I j ))
where P a
I j ) refers to the probability of ops
I j ( j ∈ { 0 , . . .
, 9 } ) being active for task a and P b ( I j ) 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, I 8 is shared across different tasks.
D A DDITIONAL R ELATED W ORK
This section provides complementary contents to Sec. 5.
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; Omidshaﬁei 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.
originally proposed relational RL, Tadepalli et al.
(2004) surveyed relational
RL, Guestrin et al.
(2003) introduced relational MDPs, and Diuk et al.
(2008) introduced objectedoriented 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 classiﬁcation (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 efﬁciency.
Compared with NLM, our Reasoner can greatly improve the efﬁciency by using a Planner module.
Second , our Reasoner has more explicitness in the reasoning path.
Let [ O r − 1 i − 1 , O ri − 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
[Equation]
where O ri is the output predicate, Z ri = Concat (cid:0) Expand (cid:0)
O r − 1 i − 1 (cid:1) , O ri − 1 , Reduce (cid:0)
O r
+1 i − 1 (cid:1)(cid:1) , σ is the sigmoid nonlinearity and θ r i denotes learnable parameters.
On the contrary, in the Reasoner module of our PRISA/PRIMA, the implementation is
[Equation]
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.