Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

Ryuichi Kanoh 1 2 Mahito Sugiyama 1 2

Abstract

service deployment in real-world settings. In this study, we
consider ensemble learning using soft trees as weak learners.

Axis-aligned rules are known to induce an important inductive bias in machine learning models such as typical hard decision tree ensembles.
However, theoretical understanding of the learning behavior is largely unrevealed due to the discrete nature of rules. To address this issue, we impose the axis-aligned constraint on differentiable
decision trees, or soft trees, which relax the splitting process of decision trees and are trained using
the gradient method. The differentiable property
enables us to derive their Neural Tangent Kernel(NTK) that can analytically describe the training
behavior. Two cases are realized: imposing theaxis-aligned constraint throughout the entire training process, or only at the initial state. Moreover,
we extend the NTK framework to handle various
tree architectures simultaneously, and prove thatany axis-aligned non-oblivious tree ensemble canbe transformed into an axis-aligned oblivious tree
ensemble with the same limiting NTK. By excluding non-oblivious trees from the search space,
the cost of trial-and-error procedures required formodel selection can be massively reduced.

In general, since all input features are taken into accountin each splitting process of soft trees, splitting boundaries
are oblique. The axis-aligned splitting is considered to be
an important inductive bias in typical hard decision trees.
Although a number of machine learning models have been
proposed that are aware of axis-aligned partitioning (Changet al., 2022; Humbird et al., 2019), it has not been theoretically clear what properties emerge when the axis-aligned
constraints are imposed.
Recently, there has been progress in the theoretical analysis
of soft tree ensembles (Kanoh & Sugiyama, 2022; 2023)
using the Neural Tangent Kernel (NTK) (Jacot et al., 2018).
The NTK framework provides analytical descriptions ofensemble learning with infinitely many soft trees. However,the current analysis is limited to soft trees with obliquesplitting and cannot directly incorporate the axis-aligned
constraint. In this paper, we extend the NTK concept to theaxis-aligned soft tree ensembles to uncover the theoreticalproperties of the axis-aligned constraint.

2. Preliminary
2.1. Soft Tree Ensembles
We formulate regression using soft decision trees based on
Kontschieder et al. (2015). Let x ∈ RF ×N be input dataconsisting of N samples with F features. Assume that there
are M soft decision trees and each tree has N splitting
nodes and L leaf nodes. We denote parameters of an mth soft decision tree, where m ∈ [M ] = {1, . . . , M }, aswm ∈ RF ×N and bm ∈ R1×N for splitting nodes andπ m ∈ R1×L for leaf nodes.

1. Introduction
A soft tree is a differentiable model that continuously relaxes
the splitting process in a decision tree and is trained using
the gradient method. There are various reasons why formulating trees in a differentiable manner is beneficial. Soft tree
ensemble models are recognized for their high empiricalperformance (Kontschieder et al., 2015; Popov et al., 2020;
Hazimeh et al., 2020). In addition, unlike hard decision
trees, soft tree models can be updated sequentially (Ke et al.,2019) and trained in conjunction with pre-training (Arik &Pfister, 2021), resulting in desirable traits for continuous

Unlike typical hard decision trees, the leaf nodes ℓ ∈ [L] ={1, . . . , L} in soft decision trees hold values µm,ℓ (xi , wm )that represent the probability of theQinput data reachN⊤ing that leaf: µm,ℓ (xi , wm , bm ) =n=1 σ(w m,n xi +
1ℓ↙n
⊤
1n↘ℓ
, where w =βbm,n )
(1 − σ(wm,n xi + βbm,n ))(w1 , . . . , wm ), b = (b1 , . . . , bm ), and π = (π 1 , . . . , π m )correspond to feature selection, splitting threshold, and prediction value of the leaf, respectively, in a typical decision
tree, and 1ℓ↙n (1n↘ℓ ) is a binary function that returns 1 if
the ℓ-th leaf is on the left (right) side of node n, and 0 other-1

National Institute of Informatics 2 The Graduate University
for Advanced Studies, SOKENDAI. Correspondence to: Ryuichi
Kanoh <kanoh@nii.ac.jp>.
Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

wise. Parameters w, b, and π are randomly initialized using
independent and identically distributed normal distributions
with mean 0 and variance 1, and updated using gradientdescent. The factor β ∈ R+ is a parameter that allows us totune the influence of the bias on the training.

Initial model
Trained model(a)

The internal nodes operate using a decision function that
resembles the sigmoid. In this paper, we use the scaled
2022;2023): σ(c) =error function (Kanoh & RSugiyama,
αc −t21
1
1
1 √2edt + 2 for c ∈ R. As
2 erf(αc) + 2 = 2
π 0

From Equation (4), if the NTK does not change during
training, the formulation of the gradient flow in the function
space becomes a simple ordinary differential equation, and it
becomes possible to analytically calculate how the model’s
output changes during training. When the NTK is positive
definite, it is known that the kernel does not change from its
initial value during the gradient descent with an infinitesimal
step size when considering an infinite number of soft binary
trees (Lee et al., 2019; Kanoh & Sugiyama, 2022; 2023)
under the formulation described in Section 2.1.

The function fm : RF × RF ×N × R1×N × R1×L → R
that returns the prediction of the m-th tree is given by theweighted sum of leaf-specific parameters πm,ℓ , weighted by
the probability that the input data xi reaches each leaf:
LX

πm,ℓ µm,ℓ (xi , wm , bm ). (1)ℓ=1Furthermore, the function f : RF × RM ×F ×N ×
RM ×1×N × RM ×1×L → R that returns the output of themodel using M trees is formulated as follows:

2.3. Setup on Axis-Aligned Splitting
Since wm,n is initialized randomly, the splitting is generally
oblique. In this study, we analyze the NTK when someof the initial values of wm,n are set to be zero. With thissetting, the corresponding features are eliminated from consideration of the splitting direction. This is technically notstraightforward, as Gaussian random initialization is generally assumed in the existing NTK approaches. In particular,
we consider the axis-aligned case where only one feature isused for splitting. We conduct a theoretical analysis of twocases: one where the parameters with zero initialization are
not updated during training, as illustrated in Figure 1(a), andthe other where they are updated, as shown in Figure 1(b).
These two cases are referred to as AAA (“A”lways “A”xis“A”ligned) and AAI (“A”xis-“A”ligned at “I”nitialization,
but not during training) in this paper.

M

1 X
fm (xi , wm , bm , π m ).
f (xi , w, b, π) = √
M m=1

(2)

2.2. Neural Tangent Kernels
We introduce the NTK based on the gradient flow using
training data x ∈ RF ×N , the prediction target y ∈ RN ,trainable parameters θ τ ∈ RP at time τ , and an arbitrarymodel function g(x, θ τ ) : RF ×N × RP → RN . With thelearning rate η and the mean squared error loss function L,
the gradient flow equation is given as∂θ τ∂g(x, θ τ )∂L(θ τ )= −η
(g(x, θ τ ) − y)). (3)
= −η
∂τ
∂θ τ
∂θ τ
Considering the formulation of the gradient flow in thefunction space using Equation (3), we obtain

The features to be assigned to each node need to be predetermined before training. This is different from typicaldecision trees, which search for features to be used for splitting in training. However, predetermining splitting features
is used in practice. For example, Extremely Randomized
Trees (Geurts et al., 2006) implemented in scikit-learn (Pedregosa et al., 2011) employs such predetermination.

⊤∂g(x, θ τ ) ∂θ τ∂g(x, θ τ )
=
∂τ
∂θ τ
∂τ
= −η

(c)

Figure 1. Schematics of training procedures. (a): AAA, Always
Axis-Aligned, during training, (b): AAI, Axis-Aligned at Initialization, but not during training, (c): Oblique splitting conductedby typical soft trees.

the scaling factor α ∈ R+ (Frosst & Hinton, 2017) tends
towards infinity, sigmoid-like decision functions become
step functions that correspond to (hard) Boolean operation.

fm (xi , wm , bm , π m ) =(b)⊤∂g(x, θ τ ) ∂g(x, θ τ )(g(x, θ τ ) − y)).
∂θ τ
∂θ τ
{z
}
|
Neural Tangent Kernel

3. Theoretical Results(4)Here, a matrix called the NTK, which has a shape of RN ×N ,appears. In this paper, the (i, j)-th component of the matrixb τ (xi , xj ).
at τ is denoted as Θ

3.1. The NTK Induced by Axis-Aligned Soft Trees
For an input vector xi , let xi,s ∈ R be the sth componentof xi . For both AAA and AAI conditions, at initialization,
2

Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels
2.5

AAA, Tree architecture=(A)2.5

Oblique

1.5

1.0

1.0

0.5

0.5

0.0

0.0

0.5

0.5

AAA, Tree architecture=(B)

AAA

1.5

Oblique

AAI

1.5

1.0

1.0

0.5

0.5

0.0

0.0

0.5

0.5

1.0

1.0

f xi w

2.0

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

AAI, Tree architecture=(A)

1.0
1.00
2.5

Oblique

= 4.0

1.0

1.0

0.5

0.5

0.0

0.0

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

AAI, Tree architecture=(B)Oblique

1.5

,

1.5

Kernel value,

2.0

1.5

(

,

1.0
1.00
2.5
= 0.1

, )

,

1.5

= 4.0

2.0

Kernel value,

= 0.1

2.0

0.5
1.0
1.00

0

200

400
600
(iteration)

800
Analytical

1000

1.5

M

= 16

0

200

M

400
600
(iteration)

800

1000

= 4096

0.5
0.75

0.50

0.25 0.00 0.25 0.50
Inner product of the inputs
15

30

0.75

1.00

1.0
1.00

45
Rotation angle (degree)

0.75

0.50

0.25 0.00 0.25 0.50
Inner product of the inputs

60

75

0.75

1.00

Figure 3. Output dynamics of test data points for axis-aligned soft
tree ensembles with two conditions. (Left): AAA, Always AxisAligned, during training, (Right): AAI, Axis-Aligned at Initialization, but not during training. Each data point is represented by adifferent line color. The left and right figures are created using
exactly the same training and test data.

Figure 2. The rotation angle dependency of Θ(xi , xj ). Different
training procedures, AAA and AAI, are listed vertically, and two
settings of tree architectures are listed horizontally. Dotted lines
show the limiting NTK induced by typical oblique soft tree ensembles defined in (Kanoh & Sugiyama, 2022; 2023), which isrotationally invariant.

In AAA, the impact of features that are not used for splitting
is completely ignored, while in AAI, the kernel is affectedby all features through the inner product of the inputs.

we can obtain the NTK induced by axis-aligned soft tree
ensembles in a closed-form as the number of trees M → ∞.
Theorem 3.1. Assume all M trees have the same soft treearchitecture. Let {a1 , · · · , aℓ , · · · , aL } denote the set ofdecomposed paths of the trees from the root to the leaves,and let h(aℓ ) ⊂ N be the set of feature indices used in eachsplit of the input path aℓ . For any binary tree architectures,as the number M of trees goes to infinity, the NTK for an
ensemble of axis-aligned soft trees converges in probability
to the deterministic kernel:

Figure 2 shows the visualization of Θ(xi , xj ). We set
α = 4.0 and β = 0.1. We calculated the kernel values for two rotated vectors: xi = (cos(ω), sin(ω)), xj =(cos(ω +ϕ), sin(ω +ϕ)) where ω ∈ [0, π/2], and ϕ ∈ [0, π].
The line colors show ω, and the x-axis shows ϕ. We use an
oblivious tree with depth 2, where we use the first feature atboth depths 1 and 2 for the architecture (A) (left column),
and we use the first feature at depth 1 and the second feature
at depth 2 for (B) (right column). We can see that the invariance with respect to rotational transformations for input
b 0 (xi , xj )has disappeared. This is different from the NTK induced
Θ(xi , xj ) := lim ΘM →∞ by typical soft tree ensembles, shown by the dotted lines.
L
X
X
Y
Moreover, when we compare left and right plots, we canΣ{i,j},s Ṫ{i,j},s=T{i,j},t  see that the kernel varies depending on the features used forℓ=1 s∈h(aℓ )t∈h(aℓ )\{s}{z} splitting.
|contribution from internal nodes
For both AAA and AAI, as the number of trees increases,
LXY
the trajectory obtained analytically from the limiting kernel(5)+
T{i,j},u ,and the trajectory during gradient descent training becomeℓ=1 u∈h(aℓ )more identical, as shown in Figure 3. This demonstrates{z}|contribution from leaves
the validity of using the NTK framework to analyze thetraining behavior. For our experiment, we consider an enwhere Σ{i,j},s = xi,s xj,s + β 2 if AAA is used,semble of oblivious trees with α = 4.0, β = 0.1, where the2
if AAI is used.
Byor Σ{i,j},s = x⊤i xj + β
first feature is used for splitting at depth 1 and the second=
denoting xi,s xj,s + β 2 as A{i,j},s ,  T{i,j},sfeatureat depth 2. This model is trained using full-batch
α2 A{i,j},s11√,arcsin+gradient
descent with a learning rate of 0.1, and the ini2π
4(α2 A{i,i},s +0.5)(α2 A{j,j},s +0.5)
tial
outputsare shifted to zero (Chizat et al., 2019). The2
1
Ṫ{i,j},s = απ q.
training
andtest datasets contain 10 randomly generated
2(1+2α2 A{i,i},s )(1+2α2 A{j,j},s )−4α4 A{i,j},sF = 2 dimensional points each. The prediction targets
The difference between AAA and AAI is whether partial
are also randomly generated. Based on Lee et al. (2019),
features of inputs are used or all features are used in Σ{i,j},s .
to derive analytical trajectories, we use the limiting kernel
3Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

Figure 4. Ensemble trees with different tree architectures. The
color of tree nodes indicates a feature used for splitting.
Figure 5. Conversion to oblivious trees inducing exactly the same
NTK. The color of tree nodes indicates a feature used for splitting.

(Theorem 3.1), as f (v, θ τ ) = H(v, x)H(x, x)−1 (I −′exp[−ηH(x, x)τ ])y, where H(x, x′ ) ∈ RN ×N denotethe limiting NTK matrix for two input matrices, and I represent an identity matrix. The input vector v ∈ RF is
arbitrary, and the training dataset and targets are denoted byx ∈ RF ×N and y ∈ RN , respectively. The behavior of theprediction trajectory changes depending on the configurations (AAA or AAI), even when exactly the same training
and test data are used.

makes the oblivious tree structure a popular choice in opensource libraries such as CatBoost (Prokhorenkova et al.,2018) and NODE (Popov et al., 2020). Kanoh & Sugiyama(2022; 2023) demonstrated that parameter sharing used in
oblivious trees does not affect the NTK of soft tree ensembles. However, their analysis does not give any insight if
only specific features are used for each splitting node.
With Theorem 3.1 and Proposition 3.2, we show that we
can always convert axis-aligned non-oblivious tree ensembles into axis-aligned oblivious tree ensembles that induce
exactly the same limiting NTK:

3.2. The NTK Induced by Ensembles of Various Trees
In the previous subsections, we have assumed that soft tree
ensembles are composed of weak learners with identicalstructures, as shown in the upper side of Figure 4. However,it can be more practical if the tree structure and features used
for splitting vary in an ensemble, as shown in the bottomside of Figure 4. To address this issue, we theoreticallyanalyze ensembles with various tree architectures mixed
together. Assuming the existence of an infinite number oftrees in an ensemble, the NTK can be computed analyticallyif the amount (ratio) of each structure in the ensemble is
known.

Proposition 3.3. Let D be the maximum depth of a single
tree. For any axis-aligned tree ensembles, a set of oblivious
trees that induce exactly the same limiting NTK can bealways constructed using 2D−1 copies of the same tree.
Figure 5 shows examples with D = 2 or 3. This insightsupports the validity of using oblivious trees when usingaxis-aligned soft trees, as in Chang et al. (2022). Creating
copies multiplies the NTK values by a constant factor, buttheoretically, this does not have a significant impact. As
Equation (4) shows, even if two kernels differ only up toa constant, adjusting the learning rate η can make theirtraining behavior exactly the same.

Proposition 3.2. Let a function p(xi , θ τ ) be the sum of
two model functions q(xi , θ ′τ ) and r(xi , θ ′′τ ), where θ ′τ ∈′′′RP , θ ′′τ ∈ RP are trainable parameters, and θ is theconcatenation of trainable parameters θ ′τ and θ ′′τ . For any
input pair xi and xj , the NTK induced by p is equal to the
b (p)
b (q)
sum of the NTKs of q and r: Θτ (xi , xj ) = Θτ (xi , xj ) +b (r)Θτ (xi , xj ).

Although various trial-and-error processes are necessary formodel selection to determine features used at each node, thisfinding can reduce the number of processes by excluding
non-oblivious trees from the search space.

For example, let q and r be functions that represent perfectbinary tree ensemble models with a depth of 1 and 2, respectively. In this case, the NTK induced by trees with a depthof 1 and 2 is the sum of the NTK induced by trees with adepth of 1 and the NTK induced by trees with a depth of 2.

4. Conclusion
In this paper, we have formulated the NTK induced by theaxis-aligned soft tree ensembles, and we have succeeded
in describing the analytical training trajectory. We havetheoretically analyzed two scenarios, one where the axisaligned constraint is applied throughout the training process,and the other where the initial model is axis-aligned andtraining proceeds without any constraints. We have alsopresented a theoretical framework to deal with non-identicaltree structures simultaneously and used it to provide theoretical support for the validity of using oblivious trees.

3.3. InsightThe oblivious tree architecture is a practical design where
the decision rules for tree splitting are shared across thesame depth. This approach reduces the number of required
splitting calculations from an exponential time complexityof O(2D ) to a linear time complexity of O(D), where D
represents the depth of the perfect binary tree. This property
4

Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

References

Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V.,
Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P.,
Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M., and Duchesnay, E.
Scikit-learn: Machine Learning in Python. Journal of
Machine Learning Research, 2011.

Arik, S. and Pfister, T. TabNet: Attentive Interpretable
Tabular Learning. Proceedings of the AAAI Conferenceon Artificial Intelligence, 2021.
Chang, C.-H., Caruana, R., and Goldenberg, A. NODEGAM: Neural generalized additive model for interpretable deep learning. In International Conference on
Learning Representations, 2022.

Popov, S., Morozov, S., and Babenko, A. Neural Oblivious
Decision Ensembles for Deep Learning on Tabular Data.
In International Conference on Learning Representations,2020.

Chizat, L., Oyallon, E., and Bach, F. On Lazy Training
in Differentiable Programming. In Advances in Neural
Information Processing Systems, 2019.

Prokhorenkova, L., Gusev, G., Vorobev, A., Dorogush, A. V.,
and Gulin, A. CatBoost: unbiased boosting with categorical features. In Advances in Neural Information
Processing Systems, 2018.

Frosst, N. and Hinton, G. E. Distilling a Neural Network
Into a Soft Decision Tree. CoRR, 2017.
Geurts, P., Ernst, D., and Wehenkel, L. Extremely Randomized Trees. Machine Learning, 2006.
Hazimeh, H., Ponomareva, N., Mol, P., Tan, Z., andMazumder, R. The Tree Ensemble Layer: Differentiability meets Conditional Computation. In Proceedings ofthe 37th International Conference on Machine Learning,2020.
Humbird, K. D., Peterson, J. L., and Mcclarren, R. G. Deep
Neural Network Initialization With Decision Trees. IEEE
Transactions on Neural Networks and Learning Systems,
2019.
Jacot, A., Gabriel, F., and Hongler, C. Neural Tangent Kernel: Convergence and Generalization in Neural Networks.
In Advances in Neural Information Processing Systems,
2018.
Kanoh, R. and Sugiyama, M. A Neural Tangent Kernel
Perspective of Infinite Tree Ensembles. In International
Conference on Learning Representations, 2022.
Kanoh, R. and Sugiyama, M. Analyzing Tree Architecturesin Ensembles via Neural Tangent Kernel. In International
Conference on Learning Representations, 2023.
Ke, G., Xu, Z., Zhang, J., Bian, J., and Liu, T.-Y. DeepGBM: A Deep Learning Framework Distilled by GBDTfor Online Prediction Tasks. In Proceedings of the 25th
ACM SIGKDD International Conference on Knowledge
Discovery & Data Mining, 2019.
Kontschieder, P., Fiterau, M., Criminisi, A., and Bulò, S. R.
Deep Neural Decision Forests. In IEEE International
Conference on Computer Vision, 2015.
Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., SohlDickstein, J., and Pennington, J. Wide Neural Networksof Any Depth Evolve as Linear Models Under Gradient
Descent. In Advances in Neural Information Processing
Systems, 2019.
5

Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

A. The NTK Induced by Typical Soft Trees
The NTK induced by a typical soft tree ensemble with infinitely many trees is known to be obtained in closed-form at
initialization.
Theorem A.1 ((Kanoh & Sugiyama, 2022; 2023)). Assume all M trees have the same soft tree architecture. Let Q : N →N ∪ {0} be a function that takes the depth as input and returns the number of leaves connected to internal nodes at that
depth. For any given tree architecture, as the number of trees in an ensemble of soft trees approaches infinity, the NTK
converges in probability to a deterministic kernel:b Oblique (xi , xj )
ΘOblique (xi , xj ) := lim Θ
0
M →∞

D
X

d−1
=
Q(d)  d Σ{i,j} T{i,j}Ṫ{i,j} +
{z
}
|
d=1
contribution from internal nodes

where

d
T{i,j}| {z }

T{i,j} =Ṫ{i,j} =α Σ{i,j}1
arcsin p 22π(α Σ{i,i} + 0.5)(α2 Σ{j,j} + 0.5)α2qπ
1 + 2α2 Σ

{i,i}




,

(A.1)

contribution from leaves2Σ{i,j} = x⊤i xj + β ,
2



!

(A.2)
1
+ ,
4

(A.3)

.

(A.4)

1

(1 + 2α2 Σ{j,j} )−4α4 Σ2{i,j}This kernel has rotational invariance with respect to input data.

B. Proofs
B.1. Proof of Theorem 3.1
Proof. Based on the independence of parameters at each leaf and the symmetry of the decision function, Kanoh & Sugiyama(2023) showed that the NTK induced by arbitrary soft tree ensembles can be decomposed into the sum of the NTKs inducedby the rule sets, which are constructed by paths from the tree root to leaves. This property of the independence of parametersat each leaf and the symmetry of the decision function also holds in our formulation (Section 2.1). Therefore, we formulate
the NTK induced by rule sets and use it to derive the NTK induced by axis-aligned soft tree ensembles.
For simplicity, first we consider the case where β = 0. Let D be the depth of a rule set, which is a path from the root toa leaf ℓ. We consider the contribution from internal nodes Θ(D,Rule,nodes) and the contribution from leaves Θ(D,Rule,leaves)
separately, such that
Θ(D,Rule) (xi , xj ) = Θ(D,Rule,nodes) (xi , xj ) + Θ(D,Rule,leaves) (xi , xj ) .

(B.1)

As for internal nodes, when we consider the axis-aligned case (Section 2.3), only a single parameter in wm,n is nonzero. When calculating the NTK as shown in Equation (4), the parameter derivatives in terms of trainable parameters are
considered. In the cases of AAA and AAI, they are given as follows:1f (D,Rule) (xi , w, π)(D−1,Rule)(xi , wm,−n , π m ) ,= √ xi,sn σ̇(wm,n,k xi,kn )fm∂wm,n,k
M

(AAA)

(B.2)

f (D,Rule) (xi , w, π)1(D−1,Rule)= √ xi σ̇(wm,n,k xi,kn )fm(xi , wm,−n , π m ) ,
∂wm,n
M

(AAI)

(B.3)

where xi,kn and wm,n,k are single features that are used in the n-th node and a single node parameter that corresponds toxi,kn , respectively. wm,−n denotes the internal node parameter matrix except for the parameters of the node n. Since there
are D possible locations for n, we obtain6

Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels

Θ(D,Rule,nodes) (xi , xj ) =Xs∈h(aℓ )whereΣ{i,j},s Ṫ{i,j},sYt∈h(aℓ )\{s}

T{i,j},t  ,

(B.4)

h
i
(D,Rule)(D,Rule)Em fm(xi , wm , π m ) fm
(xj , wm , π m )





⊤
2 
⊤
⊤
⊤
⊤
= Em σ(w⊤m,1 xi )σ(w m,1 xj ) σ(w m,2 xi )σ(w m,2 xj ) · · · σ(w m,D xi )σ(w m,D xj ) πm,1 
{z
}|
{z
}
{z
} |{z}||=

Yt∈h(aℓ )→T{i,j},s1→T{i,j},s2→1→T{i,j},sD

T{i,j},t(B.5)

is used. Here, the symbol “→” denotes the expected value of the corresponding term will take.
For leaves, Since
f (D,Rule) (xi , w, π)
1(D,Rule)√ fm=(xi , wm , π m ) ,∂πm,1
πm,1 M

(B.6)

we haveΘ(D,Rule,leaves) (xi , xj ) =Yu∈h(aℓ )Combining Equation (B.5) and Equation (B.7), we obtainXΣ{i,j},s Ṫ{i,j},sΘ(D,Rule,nodes) (xi , xj ) =s∈h(aℓ )Y

T{i,j},u .

t∈h(aℓ )\{s}

(B.7)



T{i,j},t  +Yu∈h(aℓ )T{i,j},u .

(B.8)

When we sum up this NTK over multiple rule sets constructed by multiple leaves, it becomes the NTK of the axis-aligned
soft tree ensembles:


L
X
Y
X
Y

Θ(xi , xj ) =Σ{i,j},s Ṫ{i,j},sT{i,j},t +T{i,j},u  .
(B.9)
ℓ=1

s∈h(aℓ )t∈h(aℓ )\{s}u∈h(aℓ )Up until this point, we have been considering the case where β = 0. It is straightforward to take the case β ̸= 0 into accountbecause, in the case of soft tree ensemble, the bias term can be represented by using an extra feature that takes a constantvalue β as input. This allows us to generally express the bias term by adding β 2 to Σ{i,j},s .

B.2. Proof of Proposition 3.2
Proof. The NTK induced by this model can be decomposed into the sum of the NTKs of each tree architecture as follows:b τ(p) (xi , xj ) = ∂(q(xi , θ 1 ) + r(xi , θ 2 )) , ∂(q(xj , θ 1 ) + r(xj , θ 2 ))
Θ
∂θ
∂θ

 

∂r(xi , θ 2 ) ∂r(xj , θ 2 )∂q(xi , θ 1 ) ∂q(xj , θ 1 )+.
(B.10)
=
,
,
∂θ 1
∂θ 1
∂θ 2
∂θ 2
{z
} |
{z
}
|
b (q)
Θ
τ (xi ,xj )b (r)Θτ (xi ,xj )By repeatedly using this property, the proposition can treat any number of sub-models.
7

Kernel value, = 4.0, = 0.1

Kernel value, = 4.0, = 0.1

Investigating Axis-Aligned Differentiable Trees through Neural Tangent Kernels
2.5
2.0
1.5
1.0
0.5
0.0
0.5
1.0
2.5
2.0
1.5
1.0
0.5
0.0
0.5
1.0

xi

1.00

0.75

0.50

xi

1.00

0.75

, Tree architecture=(A)= (1, 0)0.25 0.00 0.25 0.50
Inner product of the inputs= (1, 0)0.75

0.25 0.00 0.25 0.50 0.75
Inner product of the inputs
AAA: Always axis-aligned during training
= 16
= 64

1.00

, Tree architecture=(B)

0.50
M
M

2.5
2.0
1.5
1.0
0.5
0.0
0.5
1.0

M
M

= 256
= 1024

M
M

2.5
2.0
1.5
1.0
0.5
0.0
0.5
1.0

1.00

xi

1.00

0.75

xi

1.00

= (1/

0.50

= (1/

0.75

=

2)

, Tree architecture=(A)0.25 0.00 0.25 0.50
Inner product of the inputs

2 , 1/

2)

0.75

1.00

, Tree architecture=(B)0.50

0.25 0.00 0.25 0.50 0.75 1.00
Inner product of the inputs
AAI: Axis-aligned at initialization, but not during trainingMM= 40962 , 1/

= 16
= 64

M
M

= 256
= 1024

MM= 4096=b 0 (xi , xj ) to the fixed limit Θ(xi , xj ) as M increases. Two cases using
Figure A.1. An empirical demonstration of convergence of Θdifferent features are listed vertically, and two settings of vectors for computing the kernel are listed horizontally. In the upper left figure,the two cases overlap.

B.3. Proof of Proposition 3.3
Proof. The number of leaves of a perfect binary tree with the depth d is always 2d . Also, since we are considering binary
splitting at any node, the number of leaves at each node is 2. Therefore, if we multiply this by 2d−1 , we get 2d leaves at thesame depth required for a perfect binary tree. Therefore, for the maximum depth D of a tree, by having 2D−1 copies of thesame tree, we can generate all oblivious trees corresponding to root-to-leaf paths.

C. Convergence of the Kernels
Figure A.1 shows the convergence of the kernels as the number M of trees increases. We set α = 4.0 and β = 0.1. The
kernels induced by finite trees M = {16, 64, 256, 1024, 4096} are computed numerically by re-initializing the parameters 10
times. We plot two cases: xi = (1, 0), xj = (cos(ω), sin(ω)) with ω ∈ [0, π], and xi = ( √12 , √12 ), xj = (cos(ω), sin(ω))
with ω ∈ [ π4 , 5π
4 ]. This visualization confirms that as the number of trees increases, the kernel asymptotically approaches theformula defined in Theorem 3.1.

Acknowledgement
This work was supported by JST, CREST Grant Number JPMJCR22D3, Japan, and JSPS KAKENHI Grant Number
JP21H03503.

8