CNO

Convolutional Neural Operator (CNO)

neojax.models.cno.CNO

Convolutional Neural Operator (CNO).

A U-Net-style convolutional architecture, in which every convolution is paired with an alias-free activation \(\Sigma_{s, s'}\) that upsamples, applies the leaky ReLU and downsamples again with low-pass filters, such that the network maps between function spaces instead of between fixed grids. Implements the classic variant of the reference (CNO2d_classic) for any number of spatial dimensions.

The input is lifted to a latent space and passed through an encoder of n_layers levels. Each level consists of n_res residual blocks, whose output is kept for the skip connection, and a downsampling block (D-block). The bottleneck applies n_res_neck residual blocks. The decoder has n_layers levels, each consisting of an upsampling block (U-block) and optionally an invariant block (I-block) applied after the skip connection of the corresponding encoder level was concatenated (along the channel dimension). The channels double with every encoder level, starting at channel_multiplier. The output lives on the same grid as the input.

All spatial sizes of the input must be divisible by 2 ** n_layers. The model is otherwise resolution-agnostic, since the filters only depend on the ratio of the input and output resolution of the blocks.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of input channels (e.g. coordinates + initial conditions).	required
out_channels	int	Number of output channels (e.g. solution field).	required
ndim	int	Number of spatial dimensions of the input (e.g. 1, 2 or 3).	required
n_layers	int	Number of encoder (and decoder) levels. Must be at least 1.	required
n_res	int	Number of residual blocks per encoder level. Default is 1.	1
n_res_neck	int	Number of residual blocks in the bottleneck. Default is 6.	6
channel_multiplier	int	Number of channels after the first encoder level. The lifted input has channel_multiplier // 2 channels and each further level doubles the channel count. Default is 32.	32
latent_lift_proj_channels	int	Number of latent channels inside the lifting and projection blocks. Default is 64.	64
conv_kernel	int	Odd size of the convolution kernels along each dimension. Default is 3.	3
cutoff_den	float	Denominator of the cutoff frequency of the filters, see CNOConvBlock. Default is 2.0001.	2.0001
half_width_mult	float	Multiplier of the half-width of the transition band of the filters, see CNOConvBlock. Default is 0.8.	0.8
filter_size	int	Number of filter taps per phase of the resampling, see CNOConvBlock. Default is 6.	6
add_inv	bool	Whether to add an invariant block after the concatenation of the skip connection in each decoder level. Default is True.	True
positional_embedding	Union[GridEmbeddingNd, Literal['grid'], NoneType]	Positional embedding to apply to last channels of raw input before passing through the CNO. Default is None. Pass "grid" for a regular GridEmbeddingNd.	None
batch_norm	bool	Whether to apply batch normalization after the convolutions of the blocks (but not in the lifting and projection), as in the reference. If False, there is no normalization. Default is True. Batch normalization makes the model stateful, see below.	True
batch_axis_name	str	Name of the jax.vmap axis the batch statistics are computed over if batch_norm=True. Default is "batch".	'batch'
batch_norm_momentum	float	Weight of the running statistics (not of the batch statistics) in their update if batch_norm=True. Default is 0.9, which corresponds to PyTorch's momentum=0.1.	0.9
pad_mode	Literal['wrap', 'reflect', 'constant']	Padding mode of the convolutions and activations. Use "wrap" for a periodic domain, otherwise "reflect" or "constant" (zero padding, as in the reference). Default is "constant".	'constant'
negative_slope	float	Negative slope of the leaky ReLU. Default is 0.2.	0.2

Raises:

Type	Description
ValueError	If n_layers is smaller than 1.
ValueError	If an invalid positional embedding is passed.

Stateful models (batch normalization)

With batch_norm=True (the default, as in the reference), the model has running statistics and is stateful (is_stateful() returns True), following the convention of Equinox:

• The model and its initial state are created with model, state = eqx.nn.make_with_state(CNO)(...).
• The model is called as out, state = model(x, state), and has to be called inside of a jax.vmap over the batch with axis_name=batch_axis_name, in which the state is not mapped: in_axes=(0, None), out_axes=(0, None). The batch statistics are computed over this axis (see example).
• inference=True uses the running statistics instead of the batch statistics and leaves the state unchanged.

CNO's with batch_norm=False are stateless and called as out = model(x).

Differences to the reference

• The batch normalization uses the "batch" mode of eqx.nn.BatchNorm: Like PyTorch, the batch statistics are used for training. For inference, an (initialization-corrected) exponential moving average of the batch statistics is used.
• The reference supports different input and output resolutions (out_size) and enlarging the input to a multiple of 2 ** n_layers (expand_input). Neither are supported. The output has the resolution of the input.
• The reference builds an unused invariant block in the last decoder level, which is omitted here.

Internal Attributes

These fields store the internal layers state (and weights).

• positional_embedding (GridEmbeddingNd | None): Positional embedding.
• lift (CNOLiftProjectBlock): Maps inputs to the latent space.
• encoder (tuple[CNOConvBlock, ...]): Downsampling blocks of the encoder.
• res_nets (tuple[tuple[CNOResidualBlock, ...], ...]): Residual blocks of each encoder level.
• res_net_neck (tuple[CNOResidualBlock, ...]): Residual blocks of the bottleneck.
• ed_expansion (tuple[CNOConvBlock, ...]): Invariant blocks applied to the features passed from the encoder to the decoder.
• decoder_inv (tuple[CNOConvBlock, ...] | None): Invariant blocks of the decoder or None.
• decoder (tuple[CNOConvBlock, ...]): Upsampling blocks of the decoder.
• project (CNOLiftProjectBlock): Maps the latent features to the output channels.
• n_layers (int): Number of encoder (and decoder) levels.
• ndim (int): Number of spatial dimensions.

Examples:

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from neojax.models import CNO

key = jr.key(0)
model = CNO(
    key=key,
    in_channels=3,
    out_channels=1,
    ndim=2,
    n_layers=3,
    n_res=1,
    n_res_neck=2,
    channel_multiplier=8,
    batch_norm=False,
)
x = jnp.ones((3, 32, 32))
# out.shape = (1, 32, 32)
out = model(x)

# Batch normalization (default): stateful model
model, state = eqx.nn.make_with_state(CNO)(
    key=key,
    in_channels=3,
    out_channels=1,
    ndim=2,
    n_layers=3,
    n_res=1,
    n_res_neck=2,
    channel_multiplier=8,
    batch_norm=True,
)
xs = jnp.ones((4, 3, 32, 32))
batched_model = jax.vmap(
    model, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
)
# outs.shape = (4, 1, 32, 32)
outs, state = batched_model(xs, state)

Cite

Convolutional Neural Operators for robust and accurate learning of PDEs

@inproceedings{NEURIPS2023_f3c1951b,
    author = {Raonic, Bogdan and Molinaro,
    Roberto and De Ryck, Tim and Rohner, Tobias and Bartolucci,
    Francesca and Alaifari, Rima and Mishra,
    Siddhartha and de B\'{e}zenac, Emmanuel},
    booktitle = {Advances in Neural Information Processing Systems},
    doi = {10.52202/075280-3376},
    editor = {A. Oh and T. Naumann and A. Globerson
    and K. Saenko and M. Hardt and S. Levine},
    pages = {77187--77200},
    title = {Convolutional Neural Operators for
    robust and accurate learning of PDEs},
    volume = {36},
    year = {2023}
}

is_stateful() -> bool

Whether the model is stateful, i.e. uses batch normalization.

Returns:

Type	Description
bool	True if the model has to be called as model(x, state) and returns (out, state), False if it is called as model(x).

__call__(x: Inexact[Array, 'in_c ...'], state: equinox.nn._stateful.State | None = None, *, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'out_c ...'] | tuple[Inexact[Array, 'out_c ...'], equinox.nn._stateful.State]

Forward pass of the CNO model.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'in_c ...']	Input array of shape (in_channels, d1, ..., dN), where every d_i is divisible by 2 ** n_layers.	required
state	State | None	The state holding the running statistics of the batch normalizations. Required if (and only if) is_stateful().	None
key	PRNGKeyArray | None	Unused. Accepted for API compatibility with other neojax models.	None
inference	bool	Whether to use the running statistics of the batch normalizations instead of the batch statistics. Has no effect if the model is stateless. Default is False.	False

Returns:

Type	Description
Inexact[Array, 'out_c ...'] | tuple[Inexact[Array, 'out_c ...'], State]	The predicted field of shape (out_channels, d1, ..., dN). If is_stateful(), the prediction and the updated state.

Raises:

Type	Description
ValueError	If the number of spatial dimensions of x does not match ndim.
ValueError	If a spatial dimension of x is not divisible by 2 ** n_layers.
ValueError	If a state is required, but not given or vice versa.