SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tan Zhu
University of Connecticut tan.zhu@uconn.edu
Fei Dou
University of Georgia fei.dou@uga.edu
Xinyu Wang
University of Connecticut xinyu.wang@uconn.edu
Jin Lu
University of Georgia jin.lu@uga.edu
Jinbo Bi
University of Connecticut jinbo.bi@uconn.edu
Learning feature interactions can be the key for multivariate predictive modeling. ReLU-activated neural networks create piecewise linear prediction models. Other nonlinear activation functions lead to models with only high-order feature interactions, thus lacking of interpretability. Recent methods incorporate candidate polynomial terms of fixed orders into deep learning, which is subject to the issue of combinatorial explosion, or learn the orders that are difficult to adapt to different regions of the feature space. We propose a Polyhedron Attention Module (PAM) to create piecewise polynomial models where the input space is split into polyhedrons which define the different pieces and on each piece the hyperplanes that define the polyhedron boundary multiply to form the interactive terms, resulting in interactions of adaptive order to each piece. PAM is interpretable to identify important interactions in predicting a target. Theoretic analysis shows that PAM has stronger expression capability than ReLU-activated networks. Extensive experimental results demonstrate the superior classification performance of PAM on massive datasets of the click-through rate prediction and PAM can learn meaningful interaction effects in a medical problem.
Learning feature interactions provide insights into how predictors (features in  $\mathbf{x}$ ) interact to vary a dependent variable  $y$ , and could significantly improve prediction performance in a wide range of research areas, such as recommendation systems [1, 2, 3, 4], genetic analysis [5, 6] and neuroscience [7, 8], and help explain the decision-making of a complex model.
Interactions among features can significantly improve the model's expression capability. For a simple example, by incorporating the two-way interaction effect between gender (0/1: female and male) and age into the linear regression model height  $\sim w_{1} \cdot \text{gender} + w_{2} \cdot \text{age} + w_{3} \cdot \text{gender} \times \text{age} + w_{0}$  (where  $w_{0}, w_{1}, w_{2}, w_{3}$  are trainable parameters), the effects of female's age on the heights will be different from that of male ( $w_{2}$  v.s.  $w_{2} + w_{3}$ ). A predictive model  $f(\mathbf{x})$  has a  $k$ -way interaction effect if it satisfies [9]:
Definition 1 (k-way interaction effect) Let  $\mathbf{x}$  be the input of a function  $f(\mathbf{x}):\mathbb{R}^p\to \mathbb{R}$ , and  $x_{i}$  be the  $i^{th}$  feature of  $\mathbf{x}$ . Let  $\mathcal{I}\subseteq \{1,2,\dots,p\}$  be a set of feature indices and  $k$  ( $k > 2$ ) be the cardinality of  $\mathcal{I}$ . If
$$
\mathbb {E} _ {\mathbf {x}} \left[ \frac {\partial^ {| \mathcal {I} |} f (\mathbf {x})}{\partial x _ {i _ {1}} \partial x _ {i _ {2}} \dots \partial x _ {i _ {k}}} \right] ^ {2} > 0, \tag {1}
$$
$f(\mathbf{x})$  exhibits a  $k$ -way interaction effect among features indexed by  $i_1, i_2, \ldots, i_k \in \mathcal{I}$ .
A  $k$ -th order polynomial  $f(\mathbf{x})$  can define  $k$ -order interactions which may not be  $k$ -way. A  $k$ -way interaction must occur among  $k$  different features. For instance, a 2-way interaction between  $x_{i}$  and  $x_{j}$  may occur as  $x_{i}x_{j}$  or  $x_{i}^{2}x_{j}^{3}$ , so can be generally written as  $\sum_{m_i,m_j}w_{m_i,m_j}x_i^{m_i}x_j^{m_j}$  for some nonzero  $m_{i}$  and  $m_j$  if coefficients  $w$ 's are nonzero. However, high dimensional features can generate enormous candidate interactions due to the combinatorial explosion, causing the curse of dimensionality. For example,  $k$  binary features can have  $2^{k}$  possible  $k$ -way interactions. This challenge in learning feature interactions attracts tremendous research interest in various models ranging from logistic regression to deep neural networks (DNN).
Early methods learn the effects of feature interaction by multiplying features together and employing the multiplications as input to create an additive model. Two lines of methods have been used to reduce the quantity of candidate interactions. Penalization methods (e.g., Hierarchical LASSO [10, 11] and Group-LASSO [12]) assume a hierarchy of interactions, so high-way interactions may be eliminated by removing low-way interactive terms. Factorization Machines (FM) require fewer parameters to fit all predefined interaction terms, which implicitly selects among candidate interactions [13, 14, 15]. Only 2- and/or 3-way interactions were considered by these methods. To learn interactions among more features (higher-way), models such as FM-supported NN [16], product-based NN [17], AutoInt+ [18], DESTINE [19] and AFN [20] predict the target with sigmoidal-activated (such as the softmax or sigmoid function) DNNs. However, sigmoidal-activated DNN models only capture extremely high-order interactions among all features according to Definition 1, because sigmoidal activation functions have continuous (non-zero) derivatives up to infinite order with respect to all input features.
To capture interaction effects among a subset of features, an increasing number of works adopt ReLU-activated plain DNN which fits a piecewise linear function [21], as the backbone to develop interaction learning models. Models such as MaskNet [22] first form interaction effects and then use them as the input of a ReLU-activated DNN, also known as the stacked structure. Recently, more state-of-the-art performance has been achieved with models such as DeepIM [23], DeepFM [24], DCN [25], xDeepFM [26], AOANet [4], DCN-V2 [25], EDCN [27] and FinalMLP [28] in a parallel structure, which combines outputs of the proposed feature interaction learning models with those of a ReLU-activated plain DNN by operators such as summation, concatenation or Hadamard product.
By fusing ReLU-activated DNNs with models that explicitly incorporate interaction effects, piecewise polynomial functions can be formed. Rather than predicting the target with a highly non-linear function infinitely differentiable on the input space, these methods divide the input space into pieces (e.g., polyhedrons), map instances to different pieces, and predict the target with a piece-specific polynomial function. For all existing models, the piece-specific polynomial functions fit the same format of interactions (fixed  $\prod x_{i_1}^{m_1}x_{i_2}^{m_2}\cdot x_{i_k}^{m_k}$  terms for a specific  $k$ ) across all pieces once they are identified. However, data instances belonging to different pieces may endorse interactions in different ways. To address this problem, we propose a new type of attention module called Polyhedron Attention Module (PAM). The main contribution of this paper is summarized as follows:
- PAM splits the input space into polyhedrons and predicts the target using a piecewise polynomial function which contains an attention term for each piece. The attention captures interaction effects of adaptive orders for different polyhedrons according to their local structure.
- We propose a model interpretation approach for PAM to identify important interaction effects for any given data instance.
- We prove a universal approximation theorem for PAM, and show that DNNs incorporating PAMs need fewer parameters than ReLU-activated plain DNNs to maintain the same level of approximation error in fitting functions in the Sobolev space.
- Empirical studies are designed to validate PAM on benchmark datasets that are previously recognized to require feature interaction learning, such as the massive Criteo (33 million samples, 2.1 million features) and Avazu (28.3 million samples, 1.5 million features) click-through-rate datasets. Using a medical database with feature meanings, we explore the interpretation of PAM.
2 Motivation - Reinterpreting ReLU-based DNN
We argue that a plain fully-connected DNN with ReLU activation function may be interpreted as an attention mechanism. Such a neural network contains  $L$  consecutive blocks each constituting
a fully-connected layer and a ReLU-activation layer. Let  $\mathbf{x}^{\ell}$  be the input of the  $\ell$ -th block which outputs a nested function  $\mathrm{ReLU}(W^{\ell}\mathbf{x}^{\ell} + b^{\ell})$ . The  $\mathrm{ReLU}(z)$  returns  $z$  when activated (if  $z \geq 0$ ) or 0 otherwise. The consecutive  $L$  blocks together create piecewise linear functions (each output in layer  $L$  defines such a function). For a given raw input  $\mathbf{x}$ , let  $I^{\ell}$  contains the indices of the activated ReLU functions in layer  $\ell$  when  $\mathbf{x}$  passes through the layer, and  $W_{I^{\ell}}$  and  $b_{I^{\ell}}$  contain the original weights, respectively, in  $W^{\ell}$  and  $b^{\ell}$  for those rows indexed in  $I^{\ell}$  and 0 vectors for those in the complement set of  $I^{\ell}$  (denoted by  $I^{\ell -}$ ). Then, at a layer  $\ell$ ,  $W^{(\ell)}\mathbf{x} + b^{(\ell)} \geq 0$  where  $W^{(\ell)} = \prod_{j=1}^{\ell} W_{I^j}$  and  $b^{(\ell)} = \sum_{j=1}^{\ell} b_{I^j} \prod_{k=j+1}^{\ell} W_{I^k}$ . More precisely, it means that  $\mathbf{x}$  stays in a polyhedron that is defined by the intersection of all the half-spaces  $\{\mathbf{x}: W^{(\ell)}\mathbf{x} + b^{(\ell)} \geq 0\}$  and  $\{\mathbf{x}: W^{(\ell -)}\mathbf{x} + b^{(\ell -)} < 0\}$  for all  $\ell = 1, \dots, L-1$  where  $W^{(\ell -)} = W_{I^{\ell -}} \prod_{j=1}^{\ell - 1} W_{I^j}$  and  $b^{(\ell -)} = b_{I^{\ell -}} + \sum_{j=1}^{\ell - 1} b_{I^j} W_{I^{\ell -}} \prod_{k=j+1}^{\ell - 1} W_{I^k}$  specify those affine functions not activated in layer  $\ell$ . Here,  $W_{I^{\ell -}}$  and  $b_{I^{\ell -}}$  contain the original weights in  $W^{\ell}$  and  $b^{\ell}$  for those rows indexed in  $I^{\ell -}$  and 0 vectors for those in  $I^{\ell}$ . For all  $\mathbf{x}$  in this polyhedron, the  $L$ -th layer outputs affine functions  $W^{(L)}\mathbf{x} + b^{(L)}$ .
Precisely, the input space is divided into non-overlapping polyhedrons and each polyhedron  $\Delta$  is defined by a combination of activated ReLU functions across the layers (a formal proof is given in Appendix A). To ease the notation, we simply use  $\Delta$  to denote the set of indices of the activated ReLU across all layers that identify the polyhedron and  $\Delta^{-}$  denotes the index set of the inactivated ReLU in the layers. We use an indicator function  $\mathbb{1}(\mathbf{x}\in \Delta)$  to define the polyhedron which returns 1 for all vectors  $\mathbf{x}$  that satisfy  $W_{\Delta}^{(\ell)}\mathbf{x} + b_{\Delta}^{(\ell)}\geq 0$  and  $W_{\Delta}^{(\ell -)}\mathbf{x} + b_{\Delta}^{(\ell -)} < 0$  for  $\ell = 1,\dots ,L - 1$ , and 0 otherwise. The  $i$ -th activated ReLU in the  $L$ -th layer corresponds to a piecewise linear function that computes  $W_{\Delta ,i}^{(L)}\mathbf{x} + b_{\Delta ,i}^{(L)}$  on each piece  $\Delta$ . Note that  $W_{\Delta ,i}^{(L)}$  and  $b_{\Delta ,i}^{(L)}$  vary for different polyhedrons due to the differences in the ReLU activation in the early layers (illustrated in
Fig. 1). Denote the hyperplane  $\{\mathbf{x}:W_{\Delta ,i}^{(L)}\mathbf{x} + b_{\Delta ,i}^{(L)} = 0\}$  by  $H_{\Delta ,i,L}$ , which further splits  $\Delta$  into two polyhedrons  $\Delta_1 = \mathbb{1}(\mathbf{x}\in \Delta ,W_{\Delta ,i}^{(L)}\mathbf{x} + b_{\Delta ,i}^{(L)}\geq 0)$  and  $\Delta_2 = \mathbb{1}(\mathbf{x}\in \Delta ,W_{\Delta ,i}^{(L)}\mathbf{x} + b_{\Delta ,i}^{(L)} < 0)$ . Thus, this piecewise linear function can be written as:
Figure 1: An example of 2-layer ReLU-activated plain DNN with 1 output (the green shading shows the function value). Black lines are fitted in layer 1 (their intersection defines the polyhedrons) and the red line is fitted in layer 2 and varies on different polyhedrons.
$$
\begin{array}{l} f (\mathbf {x}) = \sum_ {\Delta} \mathbb {1} (\mathbf {x} \in \Delta) \cdot \operatorname {R e L U} \left(W _ {\Delta , i} ^ {(L)} \mathbf {x} + b _ {\Delta , i} ^ {(L)}\right) = \sum_ {\Delta} \mathbb {1} (\mathbf {x} \in \Delta) \cdot \operatorname {R e L U} \left(\frac {W _ {\Delta , i} ^ {(L)} \mathbf {x} + b _ {\Delta , i} ^ {(L)}}{\| W _ {\Delta , i} ^ {(L)} \|} \| W _ {\Delta , i} ^ {(L)} \|\right) \\ = \sum_ {\Delta} \underbrace {\mathbb {1} (\mathbf {x} \in \Delta_ {1}) \operatorname {d i s t} (\mathbf {x} , H _ {\Delta , i , L})} _ {\text {a t t e n t i o n}} \underbrace {| | W _ {\Delta , i} ^ {(L)} | |} _ {\text {v a l u e o n} \Delta_ {1}} + \underbrace {\mathbb {1} (\mathbf {x} \in \Delta_ {2}) \operatorname {d i s t} (\mathbf {x} , H _ {\Delta , i , L})} _ {\text {a t t e n t i o n}} \underbrace {0} _ {\text {v a l u e o n} \Delta_ {2}} \tag {2} \\ \end{array}
$$
where  $\mathrm{dist}(\mathbf{x},H_{\Delta ,i,L})$  means the distance from  $\mathbf{x}$  to the hyperplane  $H_{\Delta ,i,L}$ . The values on the two pieces  $\Delta_1$  and  $\Delta_{2}$  are constant  $||W_{\Delta ,i}^{(L)}||$  where  $||\cdot ||$  is the  $\ell_2$  vector norm or 0. The attention is a function of  $\mathbf{x}$  and depends on how far  $\mathbf{x}$  is from one of the hyperplanes that define the polyhedron boundary (see Fig. 1 for illustration).
We observe that a polyhedron  $\Delta$  is defined using a sequence of hyperplanes corresponding to the affine functions in different layers, but the attention of ReLU-activated DNNs is calculated based only on the hyperplane in the last layer for the polyhedron (piece). Although not all of the hyperplanes in early layers make the active boundary of a polyhedron (i.e., the polyhedron can locate in the interior of a half-space), using only one hyperplane in the attention is restrictive. An attention mechanism that allows multiple active boundary hyperplanes of a polyhedron to aid attention calculation may increase the model's expression power. Let  $\mathcal{H}_{\Delta}$  contain all of the active boundary hyperplanes  $H$  of  $\Delta$ . For convenient notation, and with mild relaxation, we rescale the  $W_{H}$  to  $-W_{H}$  for those inactivated affine functions in the DNN so the half-spaces can all be written in the form of  $W_{H}\mathbf{x} + b_{H} \geq 0$ . Then,  $\mathbb{1}(\mathbf{x} \in \Delta) = \prod_{H \in \mathcal{H}_{\Delta}} \mathbb{1}(W_{H}\mathbf{x} + b_{H} \geq 0)$ . To learn feature interaction effects, we multiply the distances from  $\mathbf{x}$  to each hyperplane in  $\mathcal{H}_{\Delta}$ . Given a hyperplane is linear in terms of  $\mathbf{x}$ , multiplying the distances from  $\mathbf{x}$  to two  $(m)$  hyperplanes creates quadratic  $(m$ -th order) terms. Thus, the number of active boundary hyperplanes of  $\Delta$  offers the upper bound on the order of the multiplicative terms
Figure 3: Overview description of the proposed method: (A) Split the input space into overlapping polyhedrons with the oblique tree. (B) Calculate attention on each polyhedron based on distances to the polyhedron's boundaries. (C) Calculate the model's output based on the attention and value vector. (D) Extract main and interaction effects from the model.
if each hyperplane is used only once. To allow the order to be adaptive to each  $\Delta$  within the upper limit, we further enclose a bounding constraint on the distance (other strategies may exist, which we leave for future investigation). Thus, the new attention formula can be written as follows:
$$
\begin{array}{l} a _ {\Delta} (\mathbf {x}) = \mathbb {1} (\mathbf {x} \in \Delta) \prod_ {H \in \mathcal {H} _ {\Delta}} \min  (\operatorname {d i s t} (\mathbf {x}, H), \frac {2 U _ {H}}{| | W _ {H} | |}) \\ = \prod_ {H \in \mathcal {H} _ {\Delta}} \mathbb {1} \left(W _ {H} \mathbf {x} + b _ {H} \geq 0\right) \min  (\operatorname {d i s t} (\mathbf {x}, H), \frac {2 U _ {H}}{| | W _ {H} | |}) \tag {3} \\ = \prod_ {H \in \mathcal {H} _ {\Delta}} \max  \left(\min  \left(\frac {W _ {H} \mathbf {x} + b _ {H}}{| | W _ {H} | |}, \frac {2 U _ {H}}{| | W _ {H} | |}\right), 0\right), \\ \end{array}
$$
where  $U_{H}$  determines an upper bound on the distance from  $\mathbf{x}$  to  $H$  and is a trainable parameter. Fig. 2 demonstrates how adding the upper bounds  $U_{H}$  allows the attention module to learn interaction effects of adaptive orders. For example, the instances in the area marked by 0 in the figure (left) are far from each boundary hyperplane  $H$  beyond their respective upper bounds, so the min operator returns a constant  $\frac{2U_{H}}{||W_{H}||}$ . These  $\mathbf{x}$  instances thus receive a constant attention that is the multiplication of these upper bounds each for an  $H \in \mathcal{H}_{\Delta}$ , which gives a 0-order interaction. When  $\mathbf{x}$  is close to two of
the hyperplanes (in those areas labeled by 2), two distance terms are used in Eq. 3, defining two-way interactions.
Early methods pre-specify a set of polynomial terms and use them as input features to a DNN [23] which limits the search space. In another line of methods [20], each feature is associated with an order parameter  $m_{i}$  and products in the form of  $\prod x_{i}^{m_{i}}$  are used as input (rather than raw features  $x_{i}$ ) to a ReLU-activated DNN. The DNN is required to also learn the proper values of  $m_{i}$ . It is an excellent approach to identify useful interactions, but once these  $m_{i}$ 's are determined, the orders of feature interaction are fixed and cannot adapt to different pieces because the piecewise linear model (linear in terms of the input products) learned by the DNN does not change interaction orders. Unlike
Figure 2: For a triangular polyhedron  $\Delta$  in this example, the attention  $a_{\Delta}$  can capture constant, additive, 2-way interactive, and 3-way interactive effects in the areas, respectively, marked by 0, 1, 2, and 3, by adjusting the upper bound parameter  $U$  which is smaller in the left figure than in the right.
these methods, our attention module automatically identifies the interactions (polynomial terms) of appropriate orders according to the topology of a piece (particularly, using the number of active boundary hyperplanes of the polyhedron). Using introduced parameters  $U_{H}$ , we can learn the appropriate order of the interactive terms within a polyhedron by adjusting  $U_{H}$  (larger values of  $U_{H}$  lead to higher-orders of feature interaction). Our theoretical analysis (Section 5) shows that a DNN incorporating our attention module can approximate any function represented by ReLU-activated DNN at any arbitrarily accurate level with fewer model parameters.
3 The Proposed Polyhedron Attention Module (PAM)
This section elaborates on the PAM with the attention score defined in Eq. 3. With input  $\mathbf{x} \in \mathbb{R}^p$ , a DNN using PAM defines a function  $f_{PAM}$ :
$$
f _ {P A M} (\mathbf {x}) = V (\mathbf {x}; \theta_ {G}) + \sum_ {\Delta} a _ {\Delta} (\mathbf {x}) V (\mathbf {x}; \theta_ {\Delta}), \tag {4}
$$
where  $V(\mathbf{x};\theta_G)$  is a global value function with trainable parameter  $\theta_{G}$ , and  $V(\mathbf{x};\theta_{\Delta})$  is the local value function on a piece  $\Delta$  with trainable parameters  $\theta_{\Delta}$ . We set both the global and local value
functions to be affine functions, so  $\theta$  contains the normal vector  $W$  and offset  $b$ . Sec 3.2 explains why affine functions are our choice. Instead of forming polyhedrons by intersecting hyperplanes as done in ReLU-activated DNNs, we use a tree search to partition the input space into overlapping polyhedrons. Note that the potential set of partitions created by tree search is a superset of that created by hyperplane intersections (see Appendix B). We introduce how we generate polyhedrons and then discuss the attention and value functions.
3.1 Generating polyhedrons via oblique tree
Let  $S_{\Delta}$  contain all the polyhedrons needed to form a partition of the input space. We adopt the oblique tree to generate  $S_{\Delta}$ . An oblique tree is a binary tree where each node splits the space by a hyperplane rather than by thresholding a single feature. The tree starts with the root of the full input space  $S$ , and by recursively splitting  $S$ , the tree grows deeper. For a  $D$ -depth  $(D \geq 3)$  binary tree, there are  $2^{D-1} - 1$  internal nodes and  $2^{D-1}$  leaf nodes. As shown in Fig. 3A, each internal and leaf node maintains a sub-space representing a polyhedron  $\Delta$  in  $S$ , and each layer of the tree corresponds to a partition of the input space into polyhedrons. Denote the polyhedron defined in node  $n$  by  $\Delta_{n}$ , and the left and right child nodes of  $n$  by  $n_{L}$  and  $n_{R}$ . Unlike classic oblique trees that partition  $S$  into non-overlapping sub-spaces, we perform soft partition to split each  $\Delta_{n}$  into  $\Delta_{n_{L}}$  and  $\Delta_{n_{R}}$  with an overlapping buffer. Let the splitting hyperplane be  $\{\mathbf{x} \in \mathbb{R}^{p} : W_{n}\mathbf{x} + b_{n} = 0\}$ . Then the two sub-spaces  $\Delta_{n_{L}}$  and  $\Delta_{n_{R}}$  are defined as follows:
$$
\Delta_ {n _ {L}} = \left\{\mathbf {x} \in \Delta_ {n} \mid W _ {n} \mathbf {x} + b _ {n} \geq - U _ {n} \right\},
$$
$$
\Delta_ {n _ {R}} = \left\{\mathbf {x} \in \Delta_ {n} \mid - W _ {n} \mathbf {x} - b _ {n} \geq - U _ {n} \right\}, \tag {5}
$$
where  $U_{n}$  indicates the width of the overlapping buffer. Eq. 5 shows that those instances satisfying  $|W_{n}\mathbf{x} + b_{n}| < U_{n}$  belong to the buffer in both  $\Delta_{n_L}$  and  $\Delta_{n_R}$ . This buffer creates a symmetric band around the splitting hyperplane.
Let  $\mathcal{P}_n$  be the node set containing all the ancestor nodes above  $n$ . We can group the nodes in  $\mathcal{P}_n$  into two subsets  $\mathcal{P}_n^l$  or  $\mathcal{P}_n^r$  where  $n$  appears in the left or right subtree of those ancestors. Let  $i$  index the nodes in  $\mathcal{P}_n^l$  and  $j$  index the nodes in  $\mathcal{P}_n^r$ . Then for any node  $n$  except the root node,  $\Delta_n$  can be expressed as the intersection of the half-spaces:
$$
\Delta_ {n} = \left\{\mathbf {x} \in \mathbb {R} ^ {p}: W _ {i} \mathbf {x} + b _ {i} \geq - U _ {i}, \forall i \in \mathcal {P} _ {n} ^ {l}, \text {a n d} - W _ {j} \mathbf {x} - b _ {j} \geq - U _ {j}, \forall j \in \mathcal {P} _ {n ^ {r}} \right\}. \tag {6}
$$
Based on the polyhedron defined by Eq. 6, the attention  $a_{\Delta_n}(\mathbf{x})$  (which we refer to as  $a_{n}(\mathbf{x})$  to simplify notation) can be rewritten as
$$
a _ {n} (\mathbf {x}) = \prod_ {i \in \mathcal {P} _ {n} ^ {l}} \max  \left(\min  \left(\frac {W _ {i} \mathbf {x} + b _ {i} + U _ {i}}{| | W _ {i} | |}, \frac {2 U _ {i}}{| | W _ {i} | |\}\right), 0\right) \prod_ {i \in \mathcal {P} _ {n} ^ {r}} \max  \left(\min  \left(\frac {- W _ {i} \mathbf {x} - b _ {i} + U _ {i}}{| | W _ {i} | |}, \frac {2 U _ {i}}{| | W _ {i} | |\}\right), 0\right), \tag {7}
$$
where we use the buffer width  $U_{i}$  to bound the corresponding distance term from above by  $\frac{2U_i}{||W_i||}$ .
Let  $\mathcal{N}_d$  consist of all the nodes at the  $d^{th}$  layer of the tree,  $d\in \{1,2,\dots,D\}$ . The nodes in  $\mathcal{N}_d$  altogether specify a soft partition of the input space that maps an instance  $\mathbf{x}$  to one sub-space or an intersection of overlapping sub-spaces in  $|\mathcal{N}_d|$ . Rather than merely utilizing the instance partitioning map defined by the last layer of the tree ( $2^{D - 1}$  polyhedrons), we allow PAM to leverage polyhedrons generated at all layers in  $f_{PAM}$  with  $S_{\Delta} = \cup_{d = 2}^{D}\{\Delta_n|n\in \mathcal{N}_d\}$  which gives  $2^{D} - 2$  polyhedrons.
3.2 Learning the attention and value functions