Simple reverse mode automatic differentiation of scalar values in javascript
npm install simplegradThis library implements efficient automatic differentiation of scalar valued functions.
It is lightweight and fast.
Similar to existing educational libraries like micrograd,
this library operate on computational graphs. Unlike existing educational libraries,
this library compiles computational graphs into efficient javascript code at runtime,
which results in extremely fast forward/backward passes, making this library quite
useful for quick prototyping and even real world applications.
First, you need to install it from npm.
```
npm install --save simplegrad
Now let's create a simple function f(a, b) = (a + b) * a, and compute its valuea = 2, b = 3
at point :
` js
import {Variable, ValueStorage} from 'simplegrad';
// First, let's construct computational graph of our f(a, b) function:
let valueStorage = new ValueStorage();
let a = new Variable(vs);
let b = new Variable(vs);
let f = a.add(b).mul(a);
// The graph is ready, let's produce efficient code for
// function computation:
f.compile();
// At this point we are ready to compute f(2, 3), but how do we
// assign concrete values to variables? Use setValue() method:
a.setValue(2);
b.setValue(3);
// The values are set, let's propagate them forward in the computational graph:
f.forward();
// If we want to know the value of f(2, 3), we can call getValue():getValue()
// Note: performance of is O(1), as all values are already computedf.forward()
// in call.
assert(f.getValue() === (2 + 3) * 2);
// To learn value of the df_da and df_db at point (2, 3)1
// let's propagate gradient backward.
// First, we need to tell what the value df_df is, as it could be
// coming from external source. Most of the time however it should be just :
f.setGradient(1);
// Propagate gradient backward:
f.backward();
// Remember our f(a, b) = (a + b) a, so df_da = 2 a + b, df_db = a:
assert(a.getGradient() === 2 * a.getValue() + b.getValue());
assert(b.getGradient() === a.getValue());
`
If you open the hood of this library you'll notice a beautiful engine.
Each Variable instance upon creation reserves a slot in the ValueStorage.ValueStorage is a simple object that stores values of variables in typed array:
`v
// somewhere inside ValueStorage
// variable stores values of all variables for the forward pass:gv
v = new Float64Array(variablesCount);
// variable stores gradient values for the backward pass:`
gv = new Float64Array(variablesCount);
When you call compile(), the computational graph is traversed in topologicalMaximum call stack size exceeded
order only once. The topological traversal implementation
is non-recursive, which allows to traverse huge graphs without error.
Once traversal is done, each variable has an opportunity to produce code for
forward/backward passes, and ValueStorage creates and returns dynamic functions
that can be reused for different values without re-compilation, graph traversal, or
memory allocations.
For example, function f(a, b) = (a + b) * a will be compiled by f.compile() call, whichValueStorage
in turn will ask to produce f.forward() and f.backward() functions:
` jsValueStorage
// This code is dynamically compiled by :
var v = new Float64Array(4);
var gv = new Float64Array(4);
function forward() {
v[2] = v[0] + v[1];
v[3] = v[2] * v[0];
}
function backward() {
gv[2] += v[0] gv[3]; gv[0] += v[2] gv[3]
gv[0] += gv[2]; gv[1] += gv[2]
}
return {
forward: forward,
backward: backward,
v: v,
gv: gv
}
`
As you can see the produced code is extremely efficient, each variable is
computed only once, no new objects are allocated, and the code can be optimized
by javascript engines.
Each operation that can be performed on Variable` instance is implemented in
Variable.js file. You can easily extend it with your own custom
functions.
Please let me know how it goes.
You can also sponsor my projects here - your funds will
be dedicated to more libraries and data visualizations.