Dealing with non-trainable parameters

[joaogui1]

Do we have something similar to Haiku's merge as shown in this guide? Would be useful for freezing part of the model for fine-tuning or similar ideas

[andsteing]

Just use optax.masked().

Note that the Haiku guide predates google-deepmind/optax#60

[joaogui1]

Got it, thanks!

[andsteing]

Update (from @pharringtonp19's answer below, and @matthias-wright's answer in #1706; this was also recently asked by @kohjingyu and answered by @levskaya)

Optax now recommends to use optax.multi_transform.

For Flax it can be very handy to use flax.traverse_util.ModelParamTraversal to create the second parameter:

def flattened_traversal(fn):
  """Returns function that is called with `(path, param)` instead of pytree."""
  def mask(tree):
    flat = flax.traverse_util.flatten_dict(tree)
    return flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()})
  return mask

# Freezes all but the last layer.
label_fn = flattened_traversal(
    lambda path, _: 'sgd' if path[0] == 'head' else 'none')

tx = optax.multi_transform(
    {'sgd': optax.sgd(0.1), 'none': optax.set_to_zero()}, label_fn)

Why the update? The problem with optax.masked() is that if one forgets to specify a optax.set_to_zero() then those gradients are passed through without any modification. A behaviour that makes sense from the Optax "composable gradient transformations" point of view, but might be surprising if one simply wants to update a part of a model. Using optax.multi_transform protects from this accidental mistake.

Full example

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax


class Model(nn.Module):
  num_layers: int
  depth: int
  @nn.compact
  def __call__(self, x):
    for i in range(self.num_layers - 1):
      x = nn.relu(nn.Dense(self.depth)(x))
    return nn.Dense(self.depth, name='head')(x)


model = Model(num_layers=3, depth=10)
x = jnp.zeros([1, 7])
params = model.init(jax.random.PRNGKey(0), x)['params']

jax.tree_map(jnp.shape, params)

FrozenDict({
    Dense_0: {
        bias: (10,),
        kernel: (7, 10),
    },
    Dense_1: {
        bias: (10,),
        kernel: (10, 10),
    },
    head: {
        bias: (10,),
        kernel: (10, 10),
    },
})

def flattened_traversal(fn):
  """Returns function that is called with `(path, param)` instead of pytree."""
  def mask(tree):
    flat = flax.traverse_util.flatten_dict(tree)
    return flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()})
  return mask

# Freezes all but the last layer.
label_fn = flattened_traversal(
    lambda path, _: 'sgd' if path[0] == 'head' else 'none')

tx = optax.multi_transform(
    {'sgd': optax.sgd(0.1), 'none': optax.set_to_zero()}, label_fn)

fake_grads = jax.tree_map(jnp.ones_like, params.unfreeze())
opt_state = tx.init(params.unfreeze())
updates, opt_state = tx.update(fake_grads, opt_state)

jax.tree_map(lambda x: jnp.sum(jnp.abs(x)), updates)

{'Dense_0': {'bias': DeviceArray(0., dtype=float32),
  'kernel': DeviceArray(0., dtype=float32)},
 'Dense_1': {'bias': DeviceArray(0., dtype=float32),
  'kernel': DeviceArray(0., dtype=float32)},
 'head': {'bias': DeviceArray(1., dtype=float32),
  'kernel': DeviceArray(9.999999, dtype=float32)}}

[pharringtonp19]

I am not sure whether this is a Flax issue/ Optax issue/ Or my incorrect implementation, but I haven't been able to make the above suggestion work.

from typing import Sequence  
import flax.linen as nn     
import optax                 
import jax 
import jax.numpy as jnp

class MLP(nn.Module):
  features: Sequence[int]

  @nn.compact 
  def __call__(self, x):
    for feat in self.features[:-1]:
      x = nn.relu(nn.Dense(feat)(x))
    x = nn.Dense(self.features[-1])(x)
    return x 

model = MLP([2, 2, 1])
optimizer = optax.masked(optax.sgd(0.01), {"Dense_0": False, "Dense_1": False, "Dense_2": True})

# Initialize Model
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 1)))

# Initialize Optimizer
opt_state = optimizer.init(params)

Which produces the following error

ValueError: Expected dict, got FrozenDict({
    params: {
        Dense_0: {
            kernel: DeviceArray([[-0.13036619, -0.83167976]], dtype=float32),
            bias: DeviceArray([0., 0.], dtype=float32),
        },
        Dense_1: {
            kernel: DeviceArray([[ 0.31147462,  0.8027802 ],
                         [ 0.59190094, -0.5613291 ]], dtype=float32),
            bias: DeviceArray([0., 0.], dtype=float32),
        },
        Dense_2: {
            kernel: DeviceArray([[0.32363892],
                         [0.53026754]], dtype=float32),
            bias: DeviceArray([0.], dtype=float32),
        },
    },
}).

[jheek]

Note that init returns all variables and not just params so it should be: {"params": {"Dense_0": False, "Dense_1": False, "Dense_2": True} or you should pass only the params to optimizer.init.
Also the variables are in a flax.core.FrozenDict not a normal dict to avoid in-place mutation.

[pharringtonp19]

@jheek Thanks for the advice. The code now runs! But, the output is somewhat surprising.

from typing import Sequence  
import flax.linen as nn     
import optax                 
import jax 
import jax.numpy as jnp
from flax.core import freeze, unfreeze


### Define Model
class MLP(nn.Module):
  features: Sequence[int]

  @nn.compact 
  def __call__(self, x):
    for feat in self.features[:-1]:
      x = nn.relu(nn.Dense(feat)(x))
    x = nn.Dense(self.features[-1])(x)
    return x 

### Instantiate Model + Optimizer
model = MLP([2, 2, 1])
optimizer = optax.masked(optax.sgd(0.1), {"params": {"Dense_0": False, "Dense_1": False, "Dense_2": True}})

### Initialize Model Parameters
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 1)))

### Construct Gradients
grad = jax.tree_map(jnp.ones_like, params)

### Initialize Optimizer State
opt_state = optimizer.init(unfreeze(params))

### Compute Updates
masked_updates, new_state = optimizer.update(unfreeze(grad), opt_state, unfreeze(params))
print(masked_updates)

The masked updates are not what I expected -- similar to issues raised in Optax

{'params': {'Dense_0': {'bias': DeviceArray([1., 1.], dtype=float32), 'kernel': DeviceArray([[1., 1.]], dtype=float32)}, 'Dense_1': {'bias': DeviceArray([1., 1.], dtype=float32), 'kernel': DeviceArray([[1., 1.],
             [1., 1.]], dtype=float32)}, 'Dense_2': {'bias': DeviceArray([-0.1], dtype=float32), 'kernel': DeviceArray([[-0.1],
             [-0.1]], dtype=float32)}}}

[pharringtonp19]

Based on the recent discussion in the optax repository, it seems that the solution to this is actually to use optax.multi_transform instead of optax.mask.