Update only the last layer of the pretrained model during finetuning

[snat1505027]

I want to finetune only the last layer (proj) of a pretrained model. Initially the parameters is loaded from the checkpoint and converted to a frozenDict. I tried to update only the proj layer values with this in the finetune_step:

new_params_proj = optax.apply_updates(state.params['proj'], updates['proj'])
new_params = state.params.copy({'proj':new_params_proj})

But I am getting worst validation loss with this. I am not sure whether this is the right approach.

Here is my code:

params = load_checkpoint(args.checkpoint)['params'] 
for k in ['head', 'span_encoder']:
    params.pop(k, None)
layer_shape = params['joint_transformer']['final_ln']['bias'].shape[0]
params['proj'] = {'kernel': np.random.randn(layer_shape, 1).astype(np.float32)}
params = freeze(params)
state, tx_fns = initiate_finetuning_state(opt_config=config['optimizer'], model=model, params=params)

def initiate_finetuning_state(opt_config, model, params, only_state=False):
    def mask_fn(p):
        return jax.tree_map(lambda x: (x.ndim > 1) and (x.size > 4096), p)

    tx_fns = [
        scale_by_bfloat16_adam(b1=opt_config.get('beta_1', 0.9),
                               b2=opt_config.get('beta_2', 0.98),
                               eps=opt_config.get('eps', 1e-6),
                               use_bfloat16=opt_config.get('use_bfloat16_adam', True),
                               do_bias_correction=opt_config.get('do_bias_correction', True),
                               ),
        subtract_old_weights(weight_decay=opt_config['weight_decay_rate'], mask=mask_fn),
        optax.add_decayed_weights(weight_decay=opt_config['weight_decay_rate'], mask=mask_fn),
        optax.scale_by_schedule(lr_scale_linearwarmup_lineardecay(num_warmup_steps=opt_config['num_warmup_steps'],
                                                                  num_train_steps=opt_config['num_train_steps'])),
        optax.scale(-opt_config['learning_rate']),
    ]
    tx = optax.chain(*tx_fns)

    state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
    if only_state:
        return state

    # move state to device
    state = state.replace(step=jax_utils.replicate(state.step))
    state = state.replace(opt_state=jax.tree_map(_shard_opt, state.opt_state))
    state = state.replace(params=jax_utils.replicate(state.params))
    return state, tx_fns

def finetune_step(state: train_state.TrainState, batch,
                        loss_fn: Callable[[train_state.TrainState, FrozenDict, Dict], Tuple[float, Dict]], tx_fns,
                        scan_minibatch=False):
 
    def _loss_fn(params):
        return loss_fn(state, params, batch)

    grad_fn = jax.value_and_grad(_loss_fn, has_aux=True)
    params = f32_to_bf16(state.params)

    if scan_minibatch:
        def _microbatch(old_grads, microbatch):
            def _loss_fn(params):
                return loss_fn(state, params, {k: v[None] for k, v in microbatch.items()})
            grad_fn = jax.value_and_grad(_loss_fn, has_aux=True)
            (loss, loss_info), grads = grad_fn(params)
            grads = jax.tree_multimap(lambda a, b: a + b, old_grads, grads)
            return grads, (loss, loss_info)

        grads, (loss, loss_info) = jax.lax.scan(_microbatch,
                                                init=jax.tree_map(lambda x: jnp.zeros_like(x, dtype=jnp.bfloat16), params),
                                                xs=batch)
        loss_info = jax.tree_map(lambda x: x.mean(), loss_info)
    else:
        (loss, loss_info), grads = grad_fn(params)

    grads = jax.tree_map(lambda x: jnp.nan_to_num(x, copy=False), grads)
    grads = jax.lax.pmean(grads, axis_name='batch')

    def _idx_grad(x):
        if x.shape[0] % 8 != 0:
            return x
        else:
            x = jnp.reshape(x, [8, x.shape[0] // 8] + list(x.shape[1:]))
            idx = jax.lax.axis_index('batch') % 8
            return x[idx]
    updates = jax.tree_map(_idx_grad, grads)
    updates = bf16_to_f32(updates)

    assert len(state.opt_state) == 5
    updates, nos0 = tx_fns[0].update(updates, state.opt_state[0], None)
    updates, nos1 = tx_fns[1].update(updates, state.opt_state[1], None)

    def _fix_grad(update, param):
        if update.shape == param.shape:
            return update
        else:
            aig = [[(j * 8 + i) for i in range(8)] for j in range(jax.device_count() // 8)]
            update = jax.lax.all_gather(update, axis_name='batch', axis_index_groups=aig)
            return jnp.reshape(update, param.shape)
    updates = jax.tree_multimap(_fix_grad, updates, state.params)

    new_opt_state = [nos0, nos1]
    for i in range(3):
        updates, nos = tx_fns[i + 2].update(updates, state.opt_state[i + 2], state.params)
        new_opt_state.append(nos)

    new_params = optax.apply_updates(state.params, updates)
    
    loss_info = jax.lax.pmean(loss_info, axis_name='batch')
    loss_info = bf16_to_f32(loss_info)

    new_state = state.replace(
        step=state.step + 1,
        params=new_params,
        opt_state=tuple(new_opt_state),
    )
    return new_state, loss_info

[cgarciae]

Hey @snat1505027 , check out the Transfer Learning guide's Optimization section. It shows how to use optax.multi_transform to freeze layers.