Source-linked AI summary
Adaptive Neural Trees
Ryutaro Tanno, Kai Arulkumaran, Daniel C. Alexander, Antonio Criminisi, Aditya Nori
TL;DR
Neural networks and decision trees offer complementary strengths but typically separate representation learning from data-driven architecture and hierarchical computation. ANTs unite these paradigms through neural-network paths and routers with progressive architecture growth, achieving strong results across regression and classification while supporting conditional inference and structured feature separation.
Problem
Neural networks typically use pre-specified architectures, while decision trees use pre-specified features and can lack expressive representation learning and flexible routing.
Method
ANTs combine neural-network representation learning in tree paths and routing functions with a backpropagation-based mechanism that progressively grows the architecture.
Results
ANTs achieve competitive performance across SARCOS regression, MNIST classification, and CIFAR-10 classification, including over 99% accuracy on MNIST and over 90% on CIFAR-10.
Takeaways & Limitations
ANTs provide conditional single-path inference, hierarchical feature separation, and architectures adapted to the size and complexity of the training data.
Takeaways & Limitations
The best ANT has a comparable error rate and half the parameter count of ResNet-110, while shortcut connections are left for future work.
Abstract
from arXiv · showhide
Deep neural networks and decision trees operate on largely separate paradigms; typically, the former performs representation learning with pre-specified architectures, while the latter is characterised by learning hierarchies over pre-specified features with data-driven architectures. We unite the two via adaptive neural trees (ANTs) that incorporates representation learning into edges, routing functions and leaf nodes of a decision tree, along with a backpropagation-based training algorithm that adaptively grows the architecture from primitive modules (e.g., convolutional layers). We demonstrate that, whilst achieving competitive performance on classification and regression datasets, ANTs benefit from (i) lightweight inference via conditional computation, (ii) hierarchical separation of features useful to the task e.g. learning meaningful class associations, such as separating natural vs. man-made objects, and (iii) a mechanism to adapt the architecture to the size and complexity of the training dataset.
1. Introduction
ANTs combine neural-network representation learning with decision-tree structure, using progressive growth to adapt architectures while enabling conditional, lightweight inference. Experiments across regression and classification show competitive performance, hierarchical feature separation, and benefits from combining feature learning with data partitioning.
- Neural networks learn representations but typically require hand-designed architectures and engage every network component during inference.
- Decision trees learn data-driven hierarchies and use lightweight single-path inference, but often depend on hand-engineered features and simplistic routing functions.
- ANTs represent routing decisions and root-to-leaf paths as neural networks, combining hierarchical representation learning with tree-structured partitioning.
- Progressive backpropagation-based growth chooses between deepening the network and partitioning data, adapting architecture to data availability and complexity.
- ANTs perform conditional computation by selecting one root-to-leaf path per sample, activating only a subset of model parameters during inference.
- Across SARCOS, MNIST, and CIFAR-10, ANTs achieve strong regression and classification performance, while ablations show that feature learning and data partitioning are jointly required for best prediction.
2. Related work
Prior tree-structured neural models combine neural transformations, learned routing, or conditional computation, but commonly use fixed architectures or lack representation learning along tree paths. ANTs provide a framework that combines path-wise representation learning, feature-learning routers, and architecture growth.
- Earlier soft decision trees and hierarchical mixtures of experts used fixed tree structures with simple or linear routers and identity transformations.
- More recent models added nonlinear transformations, convolutional or multilayer-perceptron routers, and conditional computation, but their architectures remained pre-specified and fixed.
- ANTs provide tree-structured models with neural-network paths, feature-learning routing functions, and a mechanism for learning the architecture.
- Table 1 compares methods by whether tree paths are neural networks, routers learn features, and architectures grow or remain pre-specified.
3. Adaptive Neural Trees
Adaptive Neural Trees combine tree-structured partitioning with learned neural representations by assigning differentiable routers, transformers, and solvers throughout a binary tree. Each input follows a routed root-to-leaf path whose neural expert produces the prediction, while single-path inference reduces computation.
- Model topology and operations: An ANT specifies a tree topology and operations, with primitive modules selected from task-appropriate choices such as CNN operations for image inputs.The model is defined over supervised data to estimate the conditional distribution p(y|x).
- Model topology and operations: ANTs are binary tree-structured models whose nodes and edges carry differentiable operations for transforming, routing, and predicting from input data.The topology consists of internal router nodes, transformer-bearing edges, and solver-bearing leaves.
- Model topology and operations: Routers send representations to child branches, transformers apply nonlinear operations that can grow deeper, and solvers output conditional-distribution estimates at leaves.Examples include CNN-based routers, convolutional transformer layers followed by ReLU, and linear classifiers for classification.
- Probabilistic model and inference: Each root-to-leaf path acts as a neural-network expert within a hierarchical mixture, with representations shared hierarchically between related experts.Standard hierarchical mixtures of experts arise as the special case in which transformers are identity functions.
- Probabilistic model and inference: Multi-path inference averages predictions across all leaves, whereas single-path inference greedily selects one router path and computes only its prediction.The single-path scheme trades the full predictive distribution for more memory- and time-efficient inference.
4. Optimisation
ANT optimisation separates architecture growth from parameter refinement. The growth phase locally tests whether to split data or deepen transformations, and a final global phase jointly tunes the fixed architecture.
- Optimisation phases: ANT training has a growth phase that learns the architecture and a refinement phase that globally tunes its parameters.Both phases use negative log-likelihood as the common objective.
- Parameter optimisation: With fixed topology, differentiable routers, transformers, and solvers are trained by backpropagation and gradient descent to minimise negative log-likelihood.The parameters jointly represented are Θ = (θ, ψ, φ).
- Architecture growth: At each leaf, growth evaluates three choices: split data with a router, deepen the incoming transform with a transformer, or keep the current model.New modules are locally optimised while earlier graph parameters remain fixed.
- Architecture growth: The model retains an added split or transformation only when it improves validation negative log-likelihood, repeating breadth-first until no candidate growth operation passes validation.This selection process controls architecture expansion using validation performance.
- Parameter refinement: After topology selection, global optimisation jointly updates all modules to refine hierarchical data grouping and the associated expert neural networks.This final stage can correct suboptimal local decisions made during growth.
5. Experiments
Experiments across regression and image classification show that ANTs combine competitive predictive performance with conditional computation, adaptive complexity, interpretable hierarchies, and refinement-based pruning. Ablations and dataset-size studies support the value of combining representation learning with hierarchical partitioning.
- Conditional computation: Single-path inference reduces active parameters while changing error by at most 0.06% for classification and 0.158 for regression.The resulting approximation remains accurate and reduces FLOPS because routers make confident decisions.
- Ablation study: Disabling routers or transformer modules consistently increases error across all three datasets, supporting the combination of feature learning and hierarchical partitioning.The ablations correspond respectively to CNN-like and SDT/HME-like alternatives.
- Model performance: ANT-CIFAR-A achieves higher accuracy than CNNs without shortcut connections using fewer parameters in single-path inference, while simpler modules produce more compact models with marginal accuracy compromises.Pretraining transformer and router parameters further reduced ANT-MNIST-A's error rate by 20%.
- Interpretability: ANTs discover task-useful hierarchical structures, including class-specialised paths such as man-made versus natural objects and road vehicles versus other vehicles.The authors caution that human-interpretable hierarchies are not necessarily optimal representations.
- Refinement: Global refinement consistently converges to higher test accuracy than the best value reached during growth, although it can initially reduce performance while changing the optimisation landscape.Refinement also polarises router probabilities and can effectively prune branches; one branch had validation visitation probability 0.09%.
- Adaptive complexity: As labelled data increases, discovered ANT complexity generally increases, while fixed-size All-CNN models overfit small datasets and linear classifiers underfit.The adaptive procedure constructs models of adequate complexity and yields better generalisation across the tested CIFAR-10 data sizes.
6. Conclusion
ANTs combine decision-tree architecture learning, conditional computation, and hierarchical clustering with neural-network representation learning and gradient-based optimization. Their progressive-growth training adapts architecture to dataset size and complexity while maintaining high performance across regression and classification.
- ANTs unite decision-tree architecture learning, conditional computation, and hierarchical clustering with deep neural-network representation learning and gradient descent.
- Progressive growth jointly tunes ANT parameters and architecture to the size and complexity of the training dataset.
- ANTs achieve high performance on SARCOS regression and MNIST and CIFAR-10 classification tasks.
A. Training algorithm
ANT training grows a tree-structured neural architecture through local decisions, then refines the fixed topology globally. The procedure uses primitive neural modules and validation-based stopping while relating ANTs to conditional computation and neural architecture search.
- Training algorithm: ANT optimisation initializes a root solver and transformer, trains parameters by gradient descent, and repeatedly expands suboptimal leaves.
- Training algorithm: Each growth decision adds either a router to partition data or a transformer to deepen representation learning.
- Connections: ANT trees implement conditional computation, while their growing procedure is related to neural architecture search over simple tree-structured networks.
- Training protocol: Patience level 5 was selected because very low or high values can cause local underfitting or overfitting and limit further growth.
- Primitive modules: Primitive modules comprise convolutional, global-average-pooling, and fully connected layers, with linear solvers and sigmoid binary-classifier routers.
D. Training times
ANT growth uses local optimization restricted to newly added graph components, enabling candidate models to be grown in under three hours on one GPU. Training-time comparisons include growth, refinement, and the All-CNN baseline.
- Under 3 hours on one GPU, local optimization grows a good candidate model by constraining gradients to the newly added graph component.
- Table 5 compares growth-phase and refinement-phase time and epochs with the training time of the All-CNN baseline.
E. Effect of training steps in the growth phase
Growth-phase patience controls how long new modules are optimized and strongly affects validation accuracy. Too little patience can terminate growth prematurely, while too much can cause local overfitting; routing quality is also reflected by poor performance from the least likely expert.
- Effect of patience: At patience 1, architecture growth terminates prematurely and validation accuracy plateaus at 80%, whereas patience 15 locally overfits at 87%.
- Effect of patience: Figure 4 plots CIFAR-10 validation-accuracy trajectories for different patience levels during training.
- Expert specialization: Using the least likely expert causes a substantial classification-accuracy drop, reaching near-random or worse performance for large trees.
G. Visualisation of discovered architectures
Discovered ANT architectures show high-confidence class separation through learned routers, while single-path inference reduces computation across ANT models.
- G. Visualisation of discovered architectures: Learned routers often separate examples by class with high confidence, revealing hierarchical clustering in discovered ANT architectures.This property resembles traditional decision trees but is not characteristic of neural tree models with pre-specified architectures.
- G. Visualisation of discovered architectures: Single-path inference reduces FLOPS for every ANT model, although the reduction varies across architectures.The comparison includes two inference schemes and reports FLOPS for ANT models alongside ResNet110 and DenseNet references.
I. Ensembling
Ensembling eight ANTs improves prediction performance under both multi-path and single-path inference. On MNIST, the ensemble approaches state-of-the-art performance while using substantially fewer parameters.
- I. Ensembling: Ensembling eight ANTs noticeably improves performance for both multi-path and single-path inference across the evaluated cases.The ensembles use ANT-A configurations for classification, with each model trained using a randomly chosen training-validation split.
- I. Ensembling: MNIST error reaches 0.29% versus 0.25% for the state-of-the-art reference, while using 851k versus 8.2M parameters.These figures compare the ANT ensemble with the cited state-of-the-art result and parameter count.