Package: neural-tangents
Stand with Ukraine!🇺🇦Freedom of thought is fundamental to all of science. Right now, our freedom is being suppressed with bombing of civilians in Ukraine.Don't be against the war - fight against the war!supportukrainenow.org.Neural TangentsICLR 2020 Video|Paper|Quickstart|Install guide|Reference docs|Release notesOverviewNeural Tangents is a high-level neural network API for specifying complex, hierarchical, neural networks of both finite andinfinitewidth. Neural Tangents allows researchers to define, train, and evaluate infinite networks as easily as finite ones. The library has been used in>100 papers.Infinite (in width or channel count) neural networks are Gaussian Processes (GPs) with a kernel function determined by their architecture. Seethis listingof papers written by the creators of Neural Tangents which study the infinite width limit of neural networks.Neural Tangents allows you to construct a neural network model from common building blocks like convolutions, pooling, residual connections, nonlinearities, and more, and obtain not only the finite model, but also the kernel function of the respective GP.The library is written in python usingJAXand leveragingXLAto run out-of-the-box on CPU, GPU, or TPU. Kernel computation is highly optimized for speed and memory efficiency, and can be automatically distributed over multiple accelerators with near-perfect scaling.Neural Tangents is a work in progress.
We happily welcome contributions!ContentsColab NotebooksInstallation5-Minute introPackage descriptionTechnical gotchasTraining dynamics of wide but finite networksPerformanceCitationColab NotebooksAn easy way to get started with Neural Tangents is by playing around with the following interactive notebooks in Colaboratory. They demo the major features of Neural Tangents and show how it can be used in research.Neural Tangents CookbookWeight Space LinearizationFunction Space LinearizationNeural Network Phase DiagramPerformance Benchmark: simple benchmark forMyrtle kernels. See alsoPerformance[New] Empirical NTK:Fully-connected networkFLAX ResNet18Experimental: Tensorflow ResNet50[New]Automatic NNGP/NTK of elementwise nonlinearitiesInstallationTo use GPU, first followJAX'sGPU installation instructions. Otherwise, install JAX on CPU by runningpip install jax jaxlib --upgradeOnce JAX is installed install Neural Tangents by runningpip install neural-tangentsor, to use the bleeding-edge version from GitHub source,git clone https://github.com/google/neural-tangents; cd neural-tangents
pip install -e .You can now run the examples and tests by calling:pip install .[testing]
set -e; for f in examples/*.py; do python $f; done  # Run examples
set -e; for f in tests/*.py; do python $f; done  # Run tests5-Minute introSee thisColabfor a detailed tutorial. Below is a very quick introduction.Our library closely follows JAX's API for specifying neural networks,stax. Instaxa network is defined by a pair of functions(init_fn, apply_fn)initializing the trainable parameters and computing the outputs of the network respectively. Below is an example of defining a 3-layer network and computing its outputsygiven inputsx.fromjaximportrandomfromjax.example_librariesimportstaxinit_fn,apply_fn=stax.serial(stax.Dense(512),stax.Relu,stax.Dense(512),stax.Relu,stax.Dense(1))key=random.PRNGKey(1)x=random.normal(key,(10,100))_,params=init_fn(key,input_shape=x.shape)y=apply_fn(params,x)# (10, 1) jnp.ndarray outputs of the neural networkNeural Tangents is designed to serve as a drop-in replacement forstax, extending the(init_fn, apply_fn)tuple to a triple(init_fn, apply_fn, kernel_fn), wherekernel_fnis the kernel function of the infinite network (GP) of the given architecture. Below is an example of computing the covariances of the GP between two batches of inputsx1andx2.fromjaximportrandomfromneural_tangentsimportstaxinit_fn,apply_fn,kernel_fn=stax.serial(stax.Dense(512),stax.Relu(),stax.Dense(512),stax.Relu(),stax.Dense(1))key1,key2=random.split(random.PRNGKey(1))x1=random.normal(key1,(10,100))x2=random.normal(key2,(20,100))kernel=kernel_fn(x1,x2,'nngp')Note thatkernel_fncan computetwocovariance matrices corresponding to theNeural Network Gaussian Process (NNGP)andNeural Tangent (NT)kernels respectively. The NNGP kernel corresponds to theBayesianinfinite neural network. The NTK corresponds to the(continuous) gradient descent trainedinfinite network. In the above example, we compute the NNGP kernel, but we could compute the NTK or both:# Get kernel of a single typenngp=kernel_fn(x1,x2,'nngp')# (10, 20) jnp.ndarrayntk=kernel_fn(x1,x2,'ntk')# (10, 20) jnp.ndarray# Get kernels as a namedtupleboth=kernel_fn(x1,x2,('nngp','ntk'))both.nngp==nngp# Trueboth.ntk==ntk# True# Unpack the kernels namedtuplenngp,ntk=kernel_fn(x1,x2,('nngp','ntk'))Additionally, if no third-argument is specified then thekernel_fnwill return aKernelnamedtuple that contains additional metadata. This can be useful for composing applications ofkernel_fnas follows:kernel=kernel_fn(x1,x2)kernel=kernel_fn(kernel)print(kernel.nngp)Doing inference with infinite networks trained on MSE loss reduces to classical GP inference, for which we also provide convenient tools:importneural_tangentsasntx_train,x_test=x1,x2y_train=random.uniform(key1,shape=(10,1))# training targetspredict_fn=nt.predict.gradient_descent_mse_ensemble(kernel_fn,x_train,y_train)y_test_nngp=predict_fn(x_test=x_test,get='nngp')# (20, 1) jnp.ndarray test predictions of an infinite Bayesian networky_test_ntk=predict_fn(x_test=x_test,get='ntk')# (20, 1) jnp.ndarray test predictions of an infinite continuous# gradient descent trained network at convergence (t = inf)# Get predictions as a namedtupleboth=predict_fn(x_test=x_test,get=('nngp','ntk'))both.nngp==y_test_nngp# Trueboth.ntk==y_test_ntk# True# Unpack the predictions namedtupley_test_nngp,y_test_ntk=predict_fn(x_test=x_test,get=('nngp','ntk'))Infinitely WideResnetWe can define a more complex, (infinitely)Wide Residual Networkusing the sament.staxbuilding blocks:fromneural_tangentsimportstaxdefWideResnetBlock(channels,strides=(1,1),channel_mismatch=False):Main=stax.serial(stax.Relu(),stax.Conv(channels,(3,3),strides,padding='SAME'),stax.Relu(),stax.Conv(channels,(3,3),padding='SAME'))Shortcut=stax.Identity()ifnotchannel_mismatchelsestax.Conv(channels,(3,3),strides,padding='SAME')returnstax.serial(stax.FanOut(2),stax.parallel(Main,Shortcut),stax.FanInSum())defWideResnetGroup(n,channels,strides=(1,1)):blocks=[]blocks+=[WideResnetBlock(channels,strides,channel_mismatch=True)]for_inrange(n-1):blocks+=[WideResnetBlock(channels,(1,1))]returnstax.serial(*blocks)defWideResnet(block_size,k,num_classes):returnstax.serial(stax.Conv(16,(3,3),padding='SAME'),WideResnetGroup(block_size,int(16*k)),WideResnetGroup(block_size,int(32*k),(2,2)),WideResnetGroup(block_size,int(64*k),(2,2)),stax.AvgPool((8,8)),stax.Flatten(),stax.Dense(num_classes,1.,0.))init_fn,apply_fn,kernel_fn=WideResnet(block_size=4,k=1,num_classes=10)Package descriptionTheneural_tangents(nt) package contains the following modules and functions:stax- primitives to construct neural networks likeConv,Relu,serial,paralleletc.predict- predictions with infinite networks:predict.gradient_descent_mse- inference with a single infinite width / linearized network trained on MSE loss with continuous gradient descent for an arbitrary finite or infinite (t=None) time. Computed in closed form.predict.gradient_descent- inference with a single infinite width / linearized network trained on arbitrary loss with continuous (momentum) gradient descent for an arbitrary finite time. Computed using an ODE solver.predict.gradient_descent_mse_ensemble- inference with an infinite ensemble of infinite width networks, either fully Bayesian (get='nngp') or inference with MSE loss using continuous gradient descent (get='ntk'). Finite-time Bayesian inference (e.g.t=1., get='nngp') is interpreted asgradient descent on the top layer only, since it converges to exact Gaussian process inference with NNGP (t=None, get='nngp'). Computed in closed form.predict.gp_inference- exact closed form Gaussian process inference using NNGP (get='nngp'), NTK (get='ntk'), or both (get=('nngp', 'ntk')). Equivalent topredict.gradient_descent_mse_ensemblewitht=None(infinite training time), but has a slightly different API (accepting precomputed kernel matrixk_train_traininstead ofkernel_fnandx_train).monte_carlo_kernel_fn- compute a Monte Carlo kernel estimate  ofany(init_fn, apply_fn), not necessarily specified viant.stax, enabling the kernel computation of infinite networks without closed-form expressions.Tools to investigate training dynamics ofwide but finiteneural networks, likelinearize,taylor_expand,empirical_kernel_fnand more. SeeTraining dynamics of wide but finite networksfor details.Technical gotchasnt.staxvsjax.example_libraries.staxWe remark the following differences between our library and the JAX one.Allnt.staxlayers are instantiated with a function call, i.e.nt.stax.Relu()vsjax.example_libraries.stax.Relu.All layers with trainable parameters use theNTK parameterizationby default. However,DenseandConvlayers also support thestandard parameterizationvia aparameterizationkeyword argument.nt.staxandjax.example_libraries.staxmay have different layers and options available (for examplent.staxlayers supportCIRCULARpadding, haveLayerNorm, but noBatchNorm.).CPU and TPU performanceFor CNNs w/ pooling, our CPU and TPU performance is suboptimal due to low core
utilization (10-20%, looks like an XLA:CPU issue), and excessive padding
respectively. We will look into improving performance, but recommend NVIDIA GPUs
in the meantime. SeePerformance.Training dynamics of wide but finite networksThe kernel of an infinite networkkernel_fn(x1, x2).ntkcombined withnt.predict.gradient_descent_msetogether allow to analytically track the outputs of an infinitely wide neural network trained on MSE loss throughout training. Here we discuss the implications forwide but finiteneural networks and present tools to study their evolution inweight space(trainable parameters of the network) andfunction space(outputs of the network).Weight spaceContinuous gradient descent in an infinite networkhas been shown into correspond to training alinear(in trainable parameters) model, which makes linearized neural networks an important subject of study for understanding the behavior of parameters in wide models.For this, we provide two convenient functions:nt.linearize, andnt.taylor_expand,which allow us to linearize or get an arbitrary-order Taylor expansion of any functionapply_fn(params, x)around some initial parametersparams_0asapply_fn_lin = nt.linearize(apply_fn, params_0).One can useapply_fn_lin(params, x)exactly as you would any other function
(including as an input to JAX optimizers). This makes it easy to compare the
training trajectory of neural networks with that of its linearization.
Prior theory and experiments have examined the linearization of neural
networks from inputs to logits or pre-activations, rather than from inputs to
post-activations which are substantially more nonlinear.Example:importjax.numpyasjnpimportneural_tangentsasntdefapply_fn(params,x):W,b=paramsreturnjnp.dot(x,W)+bW_0=jnp.array([[1.,0.],[0.,1.]])b_0=jnp.zeros((2,))apply_fn_lin=nt.linearize(apply_fn,(W_0,b_0))W=jnp.array([[1.5,0.2],[0.1,0.9]])b=b_0+0.2x=jnp.array([[0.3,0.2],[0.4,0.5],[1.2,0.2]])logits=apply_fn_lin((W,b),x)# (3, 2) jnp.ndarrayFunction space:Outputs of a linearized modelevolve identically to those of an infinite onebut with a different kernel - precisely, theNeural Tangent Kernelevaluated on the specificapply_fnof the finite network given specificparams_0that the network is initialized with. For this we provide thent.empirical_kernel_fnfunction that accepts anyapply_fnand returns akernel_fn(x1, x2, get, params)that allows to compute the empirical NTK and/or NNGP (based onget) kernels on specificparams.Example:importjax.randomasrandomimportjax.numpyasjnpimportneural_tangentsasntdefapply_fn(params,x):W,b=paramsreturnjnp.dot(x,W)+bW_0=jnp.array([[1.,0.],[0.,1.]])b_0=jnp.zeros((2,))params=(W_0,b_0)key1,key2=random.split(random.PRNGKey(1),2)x_train=random.normal(key1,(3,2))x_test=random.normal(key2,(4,2))y_train=random.uniform(key1,shape=(3,2))kernel_fn=nt.empirical_kernel_fn(apply_fn)ntk_train_train=kernel_fn(x_train,None,'ntk',params)ntk_test_train=kernel_fn(x_test,x_train,'ntk',params)mse_predictor=nt.predict.gradient_descent_mse(ntk_train_train,y_train)t=5.y_train_0=apply_fn(params,x_train)y_test_0=apply_fn(params,x_test)y_train_t,y_test_t=mse_predictor(t,y_train_0,y_test_0,ntk_test_train)# (3, 2) and (4, 2) jnp.ndarray train and test outputs after `t` units of time# training with continuous gradient descentWhat to ExpectThe success or failure of the linear approximation is highly architecture
dependent. However, some rules of thumb that we've observed are:Convergence as the network size increases.For fully-connected networks one generally observes very strong
agreement by the time the layer-width is 512 (RMSE of about 0.05 at the
end of training).For convolutional networks one generally observes reasonable
agreement by the time the number of channels is 512.Convergence at small learning rates.With a new model it is therefore advisable to start with large width on a small dataset using a small learning rate.PerformanceIn the table below we measure time to compute a single NTK
entry in a 21-layer CNN (3x3filters, no strides,SAMEpadding,ReLU) on inputs of shape3x32x32. Precisely:layers=[]for_inrange(21):layers+=[stax.Conv(1,(3,3),(1,1),'SAME'),stax.Relu()]CNN with poolingTop layer isstax.GlobalAvgPool():_, _, kernel_fn = stax.serial(*(layers + [stax.GlobalAvgPool()]))PlatformPrecisionMilliseconds / NTK entryMax batch size (NxN)CPU, >56 cores, >700 Gb RAM32112.90>= 128CPU, >56 cores, >700 Gb RAM64258.5595 (fastest - 72)TPU v232/163.255016TPU v332/162.302224NVIDIA P100325.943326NVIDIA P1006411.34918NVIDIA V100322.700126NVIDIA V100646.205818CNN without poolingTop layer isstax.Flatten():_, _, kernel_fn = stax.serial(*(layers + [stax.Flatten()]))PlatformPrecisionMilliseconds / NTK entryMax batch size (NxN)CPU, >56 cores, >700 Gb RAM320.120132048 <= N < 4096 (fastest - 512)CPU, >56 cores, >700 Gb RAM640.34142048 <= N < 4096 (fastest - 256)TPU v232/160.0015722512  <= N < 1024TPU v332/160.0010647512  <= N < 1024NVIDIA P100320.015171512  <= N < 1024NVIDIA P100640.019894512  <= N < 1024NVIDIA V100320.0046510512  <= N < 1024NVIDIA V100640.010822512  <= N < 1024Tested using version0.2.1. All GPU results are per single accelerator.
Note that runtime is proportional to the depth of your network.
If your performance differs significantly,
pleasefile a bug!Myrtle networkColab notebookPerformance Benchmarkdemonstrates how one would construct and benchmark kernels. To demonstrate
flexibility, we took theMyrtle architectureas an example. WithNVIDIA V10064-bit precision,nttook 316/330/508 GPU-hours on full 60k CIFAR-10 dataset for Myrtle-5/7/10 kernels.CitationIf you use the code in a publication, please cite our papers:# Infinite width NTK/NNGP:@inproceedings{neuraltangents2020,title={Neural Tangents: Fast and Easy Infinite Neural Networks in Python},author={Roman Novak and Lechao Xiao and Jiri Hron and Jaehoon Lee and Alexander A. Alemi and Jascha Sohl-Dickstein and Samuel S. Schoenholz},booktitle={International Conference on Learning Representations},year={2020},pdf={https://arxiv.org/abs/1912.02803},url={https://github.com/google/neural-tangents}}# Finite width, empirical NTK/NNGP:@inproceedings{novak2022fast,title={Fast Finite Width Neural Tangent Kernel},author={Roman Novak and Jascha Sohl-Dickstein and Samuel S. Schoenholz},booktitle={International Conference on Machine Learning},year={2022},pdf={https://arxiv.org/abs/2206.08720},url={https://github.com/google/neural-tangents}}# Attention and variable-length inputs:@inproceedings{hron2020infinite,title={Infinite attention: NNGP and NTK for deep attention networks},author={Jiri Hron and Yasaman Bahri and Jascha Sohl-Dickstein and Roman Novak},booktitle={International Conference on Machine Learning},year={2020},pdf={https://arxiv.org/abs/2006.10540},url={https://github.com/google/neural-tangents}}# Infinite-width "standard" parameterization:@misc{sohl2020on,title={On the infinite width limit of neural networks with a standard parameterization},author={Jascha Sohl-Dickstein and Roman Novak and Samuel S. Schoenholz and Jaehoon Lee},publisher={arXiv},year={2020},pdf={https://arxiv.org/abs/2001.07301},url={https://github.com/google/neural-tangents}}# Elementwise nonlinearities and sketching:@inproceedings{han2022fast,title={Fast Neural Kernel Embeddings for General Activations},author={Insu Han and Amir Zandieh and Jaehoon Lee and Roman Novak and Lechao Xiao and Amin Karbasi},booktitle={Advances in Neural Information Processing Systems},year={2022},pdf={https://arxiv.org/abs/2209.04121},url={https://github.com/google/neural-tangents}}