Freezing filtered parameter collections

[Jacob-Worrell]

I'm trying to work out how to do transfer learning with nnx. Below is my attempt to freeze the kernel of an nnx.Linear instance and optimize only the bias. I think maybe I'm not correctly setting up the 'wrt' argument to my optimizer.

from jax import numpy as jnp
from jax import random
from flax import nnx
import optax
from matplotlib import pyplot as plt

def f(x,m=2.234,b=-1.123):
    return m*x+b

def compute_loss(model, inputs, obs):
    prediction = model(inputs)
    error = obs - prediction
    loss = jnp.mean(error ** 2)
    mae = jnp.mean(jnp.abs(error ) )
    return loss, mae

if __name__ == '__main__':
    shape = (2,55,1)
    epochs = 123

    rngs = nnx.Rngs(123)
    model = nnx.Linear( 1, 1, rngs=rngs )

    model.kernel.value = jnp.array([[2.0]]) #load pretrained kernel  

    skey = rngs.params()
    xx = random.uniform( skey, shape, minval=-10, maxval=10 ) 
    obs1,obs2 = f(xx)
    x1,x2 = xx
    
    loss_grad = nnx.value_and_grad(compute_loss, has_aux = True)
    @nnx.scan(
        in_axes=(nnx.Carry,None,None,),
        out_axes=(nnx.Carry,0),
        length=epochs
    )
    def optimizer_scan( optimizer, x, obs ):
        (loss,mae), grads = loss_grad( optimizer.model, x, obs )        
        optimizer.update( grads )
        return optimizer, (loss,mae)

    transfer_params = nnx.All(nnx.PathContains("bias"))
    optimizer_transfer = nnx.Optimizer(model, optax.adam(learning_rate=1e-3), wrt = transfer_params)

    optimizer, (losses,maes) = optimizer_scan( optimizer_transfer, x1, obs1 )

    print( ' AFTER TRAINING' )
    print( 'training loss:', losses[-1] )

    y1,y2 = optimizer.model(xx)
    error = obs2-y2
    loss = jnp.mean( error*error )
    print( 'test loss:',loss )
    print( 'm approximation:', optimizer.model.kernel.value )
    print( 'b approximation:', optimizer.model.bias.value )

Possibly of note, the parameter filtering here is inspired by the response to issue #4167.

However, this results in the following error:

ValueError: Mismatch custom node data: ('bias', 'kernel') != ('bias',); value: State({ 'bias': VariableState( type=Param, value=Traced<ShapedArray(float32[1])>with<DynamicJaxprTrace(level=1/0)> ) }).

[Jacob-Worrell]

The missing link for me was nnx.DiffState. This is described in documentation for nnx.grad() on the nnx "transforms" page: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html

Anyway, effectively the only changes that need be made for the code to work as intended are:

1. Move the declaration of transfer_params to before the value_and_grad call,
2. Create an nnx.DiffState object,
3. Give the nnx.DiffState as the argnums keyword for nnx.value_and_grad.

And that does it!

I realized the precise nature of my problem upon returning to issue #4167 -

The following is a corrected version of the example code:

from jax import numpy as jnp
from jax import random
from flax import nnx
import optax
from matplotlib import pyplot as plt

def f(x,m=2.234,b=-1.123):
    return m*x+b

def compute_loss(model, inputs, obs):
    prediction = model(inputs)
    error = obs - prediction
    loss = jnp.mean(error ** 2)
    mae = jnp.mean(jnp.abs(error ) )
    return loss, mae

if __name__ == '__main__':
    shape = (2,55,1)
    epochs = 123

    rngs = nnx.Rngs(123)
    model = nnx.Linear( 1, 1, rngs=rngs )

    model.kernel.value = jnp.array([[2.0]]) #load pretrained kernel

    skey = rngs.params()
    xx = random.uniform( skey, shape, minval=-10, maxval=10 ) 
    obs1,obs2 = f(xx)
    x1,x2 = xx

    transfer_params = nnx.All(nnx.PathContains("bias"))
    diff_state = nnx.DiffState(0,transfer_params)
    
    loss_grad = nnx.value_and_grad(compute_loss, argnums = diff_state, has_aux = True)
    @nnx.scan(
        in_axes=(nnx.Carry,None,None,),
        out_axes=(nnx.Carry,0),
        length=epochs
    )
    def optimizer_scan( optimizer, x, obs ):
        (loss,mae), grads = loss_grad( optimizer.model, x, obs )        
        optimizer.update( grads )
        return optimizer, (loss,mae)

    optimizer_transfer = nnx.Optimizer(model, optax.adamw(learning_rate = 1e-3), wrt = transfer_params)

    optimizer, (losses,maes) = optimizer_scan( optimizer_transfer, x1, obs1 )

    print( ' AFTER TRAINING' )
    print( 'training loss:', losses[-1] )

    y1,y2 = optimizer.model(xx)
    error = obs2-y2
    loss = jnp.mean( error*error )
    print( 'test loss:',loss )
    print( 'm approximation:', optimizer.model.kernel.value )
    print( 'b approximation:', optimizer.model.bias.value )