Source-linked AI summary
TNLearn: An Open Source Python Package for Task-based Neurons
Meng Wang, Tieyun Li, Juntong Fan, Hanyu Pei, Jing-Xiao Liao, Yaodong Yang, Jianwei Ma, Fenglei Fan
TL;DR
Most successful artificial networks use a common inner-product-and-activation neuron, unlike the brain’s diverse neurons. TNLearn discovers task-based neurons and integrates them into networks; evaluations report consistent gains over linear neurons and superior or competitive performance against classical quadratic neurons.
Problem
Most successful networks use the same inner-product-and-activation neuron type, unlike the brain’s diverse neurons; the paper asks what neuronal diversity can add to artificial networks.
Source: §1 · pp. 1–2Method
TNLearn is an open-source Python package that discovers task-based neuron formulas with symbolic-regression methods and integrates them into existing architectures through modular layers.
Source: §1 · p. 2, §2.1 · pp. 3–4, §2.3 · p. 5Results
Across evaluated datasets and architectures, task-based neurons consistently surpass linear neurons in regression and classification, and achieve superior or competitive performance against classical quadratic neurons.
Source: §E.1 · pp. 17, 19Takeaways & Limitations
TNLearn provides a convenient platform for benchmarking different neuron designs.
Source: §E.1 · p. 19Potential advantages of task-based over generic neurons depend on the specific task and neural-network design.
Source: §3 · p. 6
Abstract
from arXiv · showhide
The brain does not rely on a single type of neuron to perform all kinds of tasks; instead, it designs different neurons for different tasks. The concept of task-based neurons represents a paradigm shift compared to task-based architectures. It argues that solving a specific problem requires customized neurons, as task-based neurons capture useful prior knowledge from task-related data. To facilitate the use of task-based neurons in scientific research and industrial applications, we introduce TNLearn, an open-source Python package that provides automated construction of task-based neurons and networks, enabling smooth training of task-based networks. Comprehensive documentation, including technical exposition, API reference, and representative examples, is available online. TNLearn is open-sourced at https://github.com/NewT123-WM/tnlearn and has become a PyTorch ecosystem project.
1. Introduction
The paper argues that artificial networks' reliance on a common neuron type contrasts with biological neuronal diversity, motivating task-specific neurons alongside architecture design. TNLearn addresses limited software support for task-based neurons with an open-source Python library for constructing them and integrating them into networks.
1. Introduction
- Artificial networks largely use neurons based on inner products and nonlinear activation, contrasting with the varied neurons coordinating intelligent behavior in the brain.The paper frames introducing diverse neurons into artificial networks as a way to examine what neuronal diversity can contribute.Source: §1 · pp. 1–2
- Prior work has explored new and quadratic neurons and applied them across scientific and industrial domains, supporting the feasibility of developing deep learning with new neuron types.The cited applications include medical imaging, industrial informatics, fault diagnosis, clinical imaging, and civil engineering.Source: §1 · p. 2
- Task-based neurons are customized for specific problems and use task-related data to incorporate task-specific priors, unlike task-based architectures.The authors argue that no single neuron type should be expected to perform best across all tasks.Source: §1 · p. 2
- TNLearn is an open-source Python library for creating task-based neurons with different search methods and integrating them into networks.The authors present it as a user-friendly tool for beginners and a benchmark for researchers.Source: §1 · p. 2
2. Design and Features of TNLearn
TNLearn provides a modular workflow for discovering task-based neurons, constructing them, and integrating them into networks. Its discovery engine offers multiple formula-search approaches, while customizable layers extend symbolic neuron expressions across network architectures.
2. Design and Features of TNLearn
- TNLearn uses Python's machine-learning ecosystem and automatic differentiation, with a modular design intended to support user-friendly instantiation and training with minimal code.The authors report that implementations undergo automatic unit testing.Source: §2 · p. 3
2.1 Neuron Discovery Engine
- Vectorized symbolic regression learns homogeneous formulas by applying the same operations to each input variable, with coefficients allowed to differ.TNLearn offers four discovery methods, and the paper states that no principle currently determines which is superior under which conditions.Source: §2.1 · p. 3
- GPSymRegressor evolves expressions through genetic programming, while LLMSymRegressor uses a language model to generate equation skeletons and fits their coefficients with optimizers such as BFGS.The GPSymRegressor example learns a homogeneous formula from synthetic regression data.Source: §2.1 · p. 4
- RLSymRegressor searches candidate terms with risk-seeking REINFORCE and Ridge regression, while PolyTensorRegressor combines polynomial, tensor-decomposed interaction, and periodic streams.PolyTensorRegressor jointly optimizes structural gates and projection matrices through gradient descent.Source: §2.1 · p. 4
2.2 Task-based Neuron Construction
- Task-based neuron construction parameterizes an acquired formula for learning or initialization, while retaining or allowing customization of activation functions such as ReLU.Learned parameters let neurons adapt and interact within a network.Source: §2.2 · pp. 4–5
2.3 Task-based Network Construction
- TNLearn integrates task-based neurons into specified networks through a scikit-learn estimator API and offers modular fully connected, convolutional, recurrent, and Transformer layers with user-defined symbolic expressions.These layers substitute for corresponding PyTorch modules without requiring changes to the surrounding model; recurrent cells and Transformer encoders accept the same expressions.Source: §2.3 · p. 5
- TNLearn provides documentation covering technical exposition, task-based neuron guidance, and representative examples.The paper summarizes its quick-start resources in Table 1.Source: §2.3 · p. 5, Table 1 · p. 6
3. What Can You Do with TNLearn?
The paper presents TNLearn as a tool for both designing task-based models and benchmarking neural-network research. It describes potential advantages over generic neurons while noting that these depend on the task and network design.
Model Design
- Task-based neurons may improve computational efficiency and time for specific tasks by incorporating more prior knowledge, though the advantages depend on task and network design.The paper presents these as potential advantages rather than universal guarantees.Source: §3 · p. 6
- Task-based neurons may be easier to interpret when their functions relate directly to their designed tasks, an advantage the paper characterizes as mild.Source: §3 · p. 6
Benchmark
- TNLearn can serve as a baseline reference for comparing neural-network models in machine-learning and data-science research.The paper describes this as a way to evaluate the effectiveness of different models.Source: §3 · p. 6
4. Conclusion and Outlook
TNLearn is presented as a user-friendly, versatile package for task-based neurons, with planned additions to its architectures and utilities. The paper also describes its package organization and installation options.
4. Conclusion and Outlook
- TNLearn aims to be easy to use, versatile across data, and performant across tasks, and the authors welcome community contributions.They plan to add architectures such as Mamba and more utilities while maintaining the package.Source: §4 · p. 6
- The authors plan to equip GPSymRegressor with GPU parallelization to scale to larger datasets while preserving its lightweight and extensible design.They cite EvoGP's tensorized tree encoding and adaptive parallelism as a recent demonstration of efficient GPU acceleration for tree-based genetic programming.Source: §4 · p. 6
A.1 Project Structure
- TNLearn's package directory contains regressors, custom operators, network-layer construction, and module integration, while separate directories hold tests, benchmarks, and examples.Source: §A.1 · p. 7
A.2 Installation
- TNLearn is available through pip, requires Python 3.8 or later, and resolves its PyTorch, scikit-learn, and NumPy dependencies during installation.For GPU support, users should first install a PyTorch version compatible with their CUDA version.Source: §A.2 · p. 8
- The documentation, including API reference and tutorial notebooks, is available online.Source: §A.2 · p. 8
Appendix B. Architecture and Implementation Guide
TNLearn’s appendix describes a modular path from data preparation through task-specific neuron discovery, network construction, and training, alongside integration and customization options.
B.1 General Workflow
- The general workflow loads and optionally preprocesses data, discovers or accepts a neuron formula, builds a standard or customized network, then trains and evaluates it.Users can also define a computational graph and optimization loop directly for networks assembled from tnlearn.modules.Source: §B.1 · pp. 8–9
B.2 Core Usage Pattern
- The core usage pattern applies GPSymRegressor to training data, then supplies its discovered formula as the neuron in an MLP regressor.The same pattern supports other symbolic regressors or a manually defined formula.Source: §B.2 · p. 9
B.3 Integration with Scikit-learn Workflows
- TNLearn follows the scikit-learn estimator API, so its estimators can enter pipelines, cross-validation, and GridSearchCV workflows.A task-based classifier can be passed to GridSearchCV with minimal additional code.Source: §B.3 · p. 9
B.4 Extending TNLearn
- Users can supply mathematical expressions as neuron formulas, subclass CustomNeuronLayer to change aggregation, and extend activation or loss-function registries.Developer documentation guides these customization paths.Source: §B.4 · p. 10
Appendix C. Vectorized Symbolic Regressors in TNLearn
- Appendix C frames symbolic regression as learning a neuron aggregation formula from data, treating good raw-data fit as task-aligned inductive bias and providing four complementary search strategies.The regressors can produce homogeneous formulas with built-in feature interactions.Source: Appendix C · p. 10
C.1 GPSymRegressor
The three symbolic regressors discover task-specific neuron formulas through genetic programming, LLM-guided proposals, or reinforcement-learning selection of basis terms.
C.1 GPSymRegressor
- GPSymRegressor evolves parse-tree formulas using arithmetic and inner-product operators, scoring fit by MSE while excluding expressions with fewer than two occurrences of x.Its inner-product operator supports feature interactions in vectorized expressions.Source: §C.1 · pp. 10–11
- Crossover, mutation, reproduction, and immigration modify the population; elite trees survive, and the best-fit final expression is returned, as shown in Figure 2.Subtree selection favors deeper nodes to preserve larger structures while allowing local refinement.Source: §C.1 · p. 11, Figure 2 · p. 11
C.2 LLMSymRegressor
- LLMSymRegressor prompts iterative formula proposals using an experience buffer of score-clustered islands and stage-wise hints that vary favored functional forms over the search.The buffer sampling supplies diverse, high-quality expressions as in-context examples.Source: §C.2 · pp. 12–13
- It fits proposal coefficients by multi-restart BFGS on training MSE, registers candidates in a closed loop, and returns the best optimized homogeneous expression.The output can be used as a neuron in a task-based network.Source: Figure 3 · p. 12, §C.2 · p. 13
C.3 RLSymRegressor
- RLSymRegressor samples power and interaction basis terms from a policy, vectorizes them into features, and fits their coefficients with Ridge regression.The predefined library’s maximum power and interaction order are user-specified.Source: §C.3 · p. 13
- Validation R^2 rewards update the policy via REINFORCE; the best selected structure is returned in inner-product form, with coefficients later re-optimized during network training.The policy learns which candidate terms to select, while Ridge regression fits the selected expression’s coefficients.Source: §C.3 · pp. 14–15, Figure 4 · p. 14
C.4 PolyTensorRegressor
PolyTensorRegressor makes neuron discovery a continuous optimization problem by restricting aggregation functions to a tensor-decomposed representation. Its differentiable streams and soft gates produce compact, trainable neurons that can be embedded in neural network backbones.
C.4 PolyTensorRegressor
- PolyTensorRegressor restricts candidate aggregation functions to a tensor-decomposed form, turning neuron discovery into a continuous optimization problem.Given an input matrix X ∈ R^(d×N), it represents the candidate function through three complementary streams.Source: §C.4 · p. 15
- The polynomial stream models individual-feature transformations, the CP-decomposed interaction stream captures high-order dependencies, and the periodic component adds a sinusoidal term for cyclic patterns.Directly parameterizing the interaction stream reduces parameter complexity from O(d^m) to O(mRd) when d_out is treated as constant.Source: §C.4 · p. 15
- Differentiable soft gates are jointly optimized with projection matrices, then components are retained or pruned according to a threshold after convergence.Training uses gradient descent and may include a sparsity-inducing regularizer such as an ℓ1 penalty.Source: §C.4 · p. 15
- The differentiable search is described as more stable and reproducible than discrete genetic programming, especially for high-dimensional inputs or high-order interactions.The discovered neuron can be returned as a trainable layer and embedded into neural network backbones without architectural redesign.Source: §C.4 · p. 15
Appendix D. Discovered Neuron Formulas on Benchmark Datasets
The appendix presents GPSymRegressor-discovered formulas on six benchmark datasets and describes how formula discovery is used to construct task-based networks. The benchmarks span varied domains, sample sizes, and feature dimensionalities.
Appendix D. Discovered Neuron Formulas on Benchmark Datasets
- Table 2 presents formulas discovered by GPSymRegressor on six benchmark datasets.The table identifies each dataset’s instance count, feature count, task type, and discovered formula.Source: Table 2 · p. 16, Appendix D · p. 17
- The OpenML datasets span varied application domains, sample sizes, and feature dimensionalities, and are identified by their OpenML IDs.The named set includes satimage, phoneme, adult, pol, fried, Bike Sharing Demand, and sarcos.Source: Appendix D · pp. 16–17
- Formula discovery uses up to 1,000 training instances for computational efficiency, and displayed numeric constants become learnable parameters when building task-based neurons.The resulting formulas are supplied to MLPRegressor or MLPClassifier when constructing a task-based network.Source: Appendix D · p. 17
E.1 Evaluation on Classification and Regression Tasks
The evaluation compares task-based networks with linear neurons, random formulas, and quadratic neuron variants. Across the reported comparisons, task-based neurons perform better than linear neurons and are superior or competitive against quadratic variants.
E.1 Evaluation on Classification and Regression Tasks
- The evaluation constructs task-based networks for regression and classification using MLPRegressor or MLPClassifier with formulas discovered by TNLearn’s regressors.The comparison covers datasets and network architectures, with Table 3 reporting the linear-versus-task-based results.Source: §E.1 · p. 17, Table 3 · p. 17
- Across evaluated datasets and architectures, task-based neurons surpass linear neurons in both regression and classification tasks.The comparison uses the same network architectures; the authors attribute the result to enhanced feature representation from task-based inductive bias.Source: §E.1 · p. 17, Table 3 · p. 17
- Task-based networks consistently outperform random-formula networks across nearly all datasets, including with smaller architectures and fewer parameters.The reported outcomes are higher regression R2 or higher classification accuracy; Table 4 identifies the discovered-formula and random-formula rows.Source: Table 4 · p. 18, §E.1 · p. 19
- Against four representative quadratic neuron variants, task-based neurons achieve superior or competitive performance, and the authors report overall superiority to classical quadratic neurons.The comparison constructs the quadratic-neuron networks in TNLearn using their corresponding expressions.Source: §E.1 · p. 19, Table 5 · p. 19
E.2 Evaluation on Images
On MNIST, TNConv2d achieves competitive performance with comparable final training loss and faster early convergence; its image-neuron discovery pipeline first compresses image patches with an autoencoder, then applies vectorized symbolic regression.
E.2 Evaluation on Images
- The controlled MNIST comparison uses 60,000 training and 10,000 test 28 × 28 grayscale images across 10 classes, with identical lightweight backbones and training protocols.Each experiment is repeated five times, and average results are reported.Source: §E.2 · pp. 19–20
- An autoencoder compresses image patches into a low-dimensional latent representation before vectorized symbolic regression discovers the neuron aggregation function.The compressed representation is described as preserving task-relevant semantic information while reducing discovery cost and improving convergence.Source: §E.2 · p. 20
- TNConv2d achieves competitive performance while maintaining comparable final training loss and faster early convergence on MNIST.Table 6 summarizes training loss and test accuracy across epochs.Source: Table 6 · p. 20, §E.2 · p. 20
- Task-specific symbolic neurons capture more intricate pixel interactions, while training-loss trends indicate convergence speed comparable to standard Conv2d.The paper presents this as a balance between representational power and trainability.Source: §E.2 · p. 21