Jacobian-Enhanced Neural Network (JENN)

Jacobian-Enhanced Neural Networks (JENN) are fully connected multi-layer perceptrons, whose training process is modified to predict partial derivatives more accurately. This is accomplished by minimizing a modified version of the Least Squares Estimator (LSE) that accounts for Jacobian prediction error (see paper). The main benefit of jacobian-enhancement is better accuracy with fewer training points compared to standard fully connected neural nets, as illustrated below.

Example #1	Example #2

Example #3

Citation

If you use JENN in a scientific publication, please consider citing it:

@misc{berguin2024jacobianenhanced,
      title={Jacobian-Enhanced Neural Networks}, 
      author={Steven H. Berguin},
      year={2024},
      eprint={2406.09132},
      archivePrefix={arXiv},
      primaryClass={id='cs.LG' full_name='Machine Learning' is_active=True alt_name=None in_archive='cs' is_general=False description='Papers on all aspects of machine learning research (supervised, unsupervised, reinforcement learning, bandit problems, and so on) including also robustness, explanation, fairness, and methodology. cs.LG is also an appropriate primary category for applications of machine learning methods.'}
}

Main Features

• Multi-Task Learning : predict more than one output with same model Y = f(X) where Y = [y1, y2, ...]
• Jacobian prediction : analytically compute the Jacobian (i.e. forward propagation of dY/dX)
• Gradient-Enhancement: minimize prediction error of partials (i.e. back-prop accounts for dY/dX)

Installation

pip install jenn 

For the optional agent integration (MCP server), which requires Python >= 3.10:

pip install "jenn[mcp]"

Example Usage

See demo notebooks for more details

Import library:

import jenn

Generate example training and test data:

x_train, y_train, dydx_train = jenn.utilities.sample(
    f=jenn.synthetic_data.sinusoid.compute, 
    f_prime=jenn.synthetic_data.sinusoid.compute_partials, 
    m_random=0, 
    m_levels=4, 
    lb=-3.14, 
    ub=3.14,
)
x_test, y_test, dydx_test = jenn.utilities.sample(
    f=jenn.synthetic_data.sinusoid.compute, 
    f_prime=jenn.synthetic_data.sinusoid.compute_partials, 
    m_random=30, 
    m_levels=0, 
    lb=-3.14, 
    ub=3.14,
)

Train a model:

nn = jenn.NeuralNet(
    layer_sizes=[1, 12, 1],
).fit(
    x=x_train,  
    y=y_train, 
    dydx=dydx_train,
    lambd=0.1,  # regularization parameter 
    is_normalize=True,  # normalize data before fitting it
)

Make predictions:

y, dydx = nn(x) 

# OR 

y = nn.predict(x)
dydx = nn.predict_partials(x)

Save model (parameters) for later use:

nn.save('parameters.json')  

Reload saved parameters into new model:

reloaded = jenn.NeuralNet.load('parameters.json')

Check goodness of fit:

jenn.plot_goodness_of_fit(
    y_true=y_test, 
    y_pred=nn.predict(x_test), 
    title="y (JENN)"
)

Check goodness of fit of partials:

jenn.plot_goodness_of_fit(
    y_true=dydx_test, 
    y_pred=nn.predict_partials(x_test), 
    title="dy/dx (JENN)"
)

Show sensitivity profiles:

jenn.plot_sensitivity_profiles(
    func=[jenn.synthetic_data.sinusoid.compute, nn.predict], 
    x_min=x_train.min(), 
    x_max=x_train.max(), 
    x_true=x_train, 
    y_true=y_train, 
    resolution=100, 
    legend_label=['true', 'pred'], 
    xlabels=['x'], 
    ylabels=['y'],
)

MCP Server

JENN ships an optional Model Context Protocol server so an AI agent (e.g. Claude Code) can build and validate a surrogate model end-to-end — from data (with optional partials) to a portable artifact — without writing a training script by hand.

pip install "jenn[mcp]"
claude mcp add --transport stdio jenn -- jenn-mcp

It exposes tools for the full lifecycle (ingest, train, evaluate, export, load_model, predict, and listing helpers), jenn://files resources for discovering local data/model files — the whole folder, plus each file individually so an agent's @ menu can browse them — and a surrogate_workflow prompt. See the MCP Server section of the documentation for the full tool reference, a hands-on tutorial with practice data, and JENN_DIR setup.

Use Case

JENN is intended for the field of computer aided design, where there is often a need to replace computationally expensive, physics-based models with so-called surrogate models in order to save time down the line. Since the surrogate model emulates the original model accurately in real time, it offers a speed benefit that can be used to carry out orders of magnitude more function calls quickly, opening the door to Monte Carlo simulation of expensive functions for example.

In general, the value proposition of a surrogate is that the computational expense of generating training data to fit the model is much less than the computational expense of performing the analysis with the original physics-based model itself. However, in the special case of gradient-enhanced methods, there is the additional value proposition that partials are accurate which is a critical property for one important use-case: surrogate-based optimization. The field of aerospace engineering is rich in applications of such a use-case.

Limitations

Gradient-enhanced methods require responses to be continuous and smooth, but they are only beneficial if the cost of obtaining partials is not excessive in the first place (e.g. adjoint methods), or if the need for accuracy outweighs the cost of computing the partials. Users should therefore carefully weigh the benefit of gradient-enhanced methods relative to the needs of their application.

A Note on AI-Assisted Development

JENN began in 2018, well before today's generation of AI coding tools, and its foundation is hand-derived neural-network mathematics rather than generated code. We recognize how capable AI has since become, and we leverage it judiciously and under human oversight — applying it where it genuinely improves the project, not as a substitute for understanding it. Every change, whatever its source, is held to the same standard of review.

License

Distributed under the terms of the MIT License.

Acknowledgement

This code used the code by Prof. Andrew Ng in the Coursera Deep Learning Specialization as a starting point. It then built upon it to include additional features such as line search and plotting but, most of all, it fundamentally changed the formulation to include gradient-enhancement and made sure all arrays were updated in place (data is never copied). The author would like to thank Andrew Ng for offering the fundamentals of deep learning on Coursera, which took a complicated subject and explained it in simple terms that even an aerospace engineer could understand.