New and improved NN library system in penzai.experimental.v2

[danieldjohnson]

Release 0.1.3 includes a redesign of Penzai's neural network components to reduce boilerplate, make it easier to manipulate models, and open the door to future improvements. Key changes:

• Variable system
  • Parameters are now mutable Python objects, and are PyTree leaves instead of PyTree nodes. This means they can be easily shared between layers or models. (Models themselves remain immutable PyTrees.)
  • State variables are also mutable Python objects. Effect handlers are no longer necessary to use state.
  • Passing models through JAX transformations may require "unbinding" or "freezing" these variables first
• Parameter initialization
  • Parameters are now initialized when the model is built instead of being initialized in a separate step
  • Parameter sharing now "just works" without requiring effect handlers
  • Parameter names are explicitly passed between layers as arguments, and are no longer renamed after construction
• Layer signature and side inputs
  • Layers now accept side inputs as keyword arguments, instead of using PyTree dependency injection
• Transformer implementation
  • Llama, Mistral, and Pythia/GPT-NeoX models are now supported in addition to Gemma

For a detailed walkthrough of the changes, see this documentation page.

We are currently planning to deprecate the original penzai.nn and penzai.data_effects system and replace them with this new system for release v0.2.0. (This will be a breaking change, although we plan to keep the old system available in a deprecated submodule to simplify migration.)

Please try out the new system and give feedback!

[amifalk]

Do we need to expect tools unrelated to the new layer system to be deprecated (e.g. like auto_nmap or type aliases like LayerLike) or will those be preserved in 0.2.0?

[danieldjohnson]

Tools like auto_nmap that aren't related to the new layer system will be preserved in 0.2.0. This also applies to the core definitions like pz.pytree_dataclass, pz.Struct, pz.select, pz.Selection.

However, the LayerLike alias in particular will probably be deprecated in 0.2.0, because the interface for Layer has changed.

(Also, tools like toolshed.isolate_submodel that do make use of the layer system will be replaced with V2 versions, most of which are already in penzai.experimental.v2.toolshed.)

[amifalk]

What is the recommended api for implementing checked layer calls in v2? I noticed that the new Layer class does not have an input_structure and output_structure method.

[danieldjohnson]

Thanks for the question! There isn't currently a direct replacement for checked layer calls in the V2 system. However, for individual layers, you can still check shapes manually by using pz.chk.check_structure inside the layer __call__. This is currently done in the implementation for the V2 Linear, for instance.

For more context, the reason for removing the input_structure and output_structure methods from the base Layer class was that it didn't seem like the benefits of enforcing shape checking outweighed the costs, especially given the limitations of the V1 shapechecking:

• The shape-checking inference logic couldn't easily express requiring that two inputs be broadcastable together, which meant that shape-checking was often either too strict or too loose. Most layers ended up having annotations that were essentially just Any.
• Side inputs are now passed in as keyword arguments in V2, and layers are expected to ignore unrecognized keyword arguments, so it's harder to annotate the input structures of layers. Also, side inputs can sometimes contain variable nodes.
• It's difficult to maintain correct structure annotations for higher-order layers and combinators, because their content may be modified via model surgery, and it's usually nontrivial to try to infer structure annotations from the contents.
• When parameters are unbound from models in the V2 design, there may no longer be enough information to infer the expected input and output structures for layers, because the parameter slots do not store or restrict the parameter shapes.
• Adding shape annotations in the pretty-printer ended up making the rendered outputs somewhat verbose for most models, and was mostly redundant with the contents of the layers. Overall it seemed less useful than I had initially hoped.

Overall, the value of keeping shape checking as a core part of the Layer interface seemed fairly limited, so I switched to just using shape-checking inside __call__ for the specific layers where it seemed useful.