reference/code/ELF/src/train_step.py

"""Per-device pmap'd training step for the ELF diffusion language model."""

from typing import Dict, Tuple

import jax
import jax.numpy as jnp

from utils.train_utils import TrainState
from utils.encoder_utils import encode_text
from utils.sampling_utils import (
    sample_cfg_scale, add_noise, sample_timesteps,
    net_out_to_v_x, restore_cond,
)


Array = jnp.ndarray


def train_step(
    state: TrainState,
    encoder_params: Dict,
    encoder_apply_fn,
    batch: Dict[str, Array],
    config,
) -> Tuple[TrainState, Dict[str, float]]:
    """Perform a single training step."""
    t_eps = config.t_eps
    self_cond_prob = config.self_cond_prob
    latent_mean, latent_std = config.latent_mean, config.latent_std

    decoder_prob = config.decoder_prob
    decoder_noise_scale = config.decoder_noise_scale

    new_dropout_rng, current_step_rng = jax.random.split(state.dropout_rng, 2)
    current_step_rng = jax.random.fold_in(current_step_rng, jax.lax.axis_index(axis_name="batch"))
    (
        t_rng, noise_rng, self_cond_mask_rng, self_cond_cfg_rng,
        model_dropout_rng, decoder_step_rng, decoder_rng,
        decoder_lambda_rng, decoder_noise_rng,
    ) = jax.random.split(current_step_rng, 9)

    # encoder_attention_mask: cond sees cond, x sees all
    encoder_attention_mask = batch["encoder_attention_mask"]

    # Label drop before encoding: prevent target tokens from attending to
    # condition tokens so x0 is truly unconditional for dropped samples.
    if config.label_drop_prob > 0:
        drop = batch["label_drop_mask"][:, None, None]  # (B, 1, 1)
        cond_mask = batch["cond_seq_mask"]  # (B, S)
        # block_mask is 1 only at (non-cond row, cond col) — leaves cond↔cond unchanged
        block_mask = (1 - cond_mask)[:, :, None] * cond_mask[:, None, :]
        encoder_attention_mask = encoder_attention_mask * (1 - drop * block_mask)

    x0 = encode_text(
        input_ids=batch["input_ids"],
        attention_mask=encoder_attention_mask,
        encoder_apply_fn=encoder_apply_fn,
        encoder_params=encoder_params,
        latent_mean=latent_mean,
        latent_std=latent_std,
    )

    batch_size, seq_length = x0.shape[0], x0.shape[1]

    t = sample_timesteps(
        t_rng, batch_size,
        P_mean=config.denoiser_p_mean, P_std=config.denoiser_p_std,
        time_schedule=config.time_schedule,
    )

    noise = jax.random.normal(noise_rng, x0.shape, dtype=x0.dtype)

    cond_seq_mask = batch["cond_seq_mask"][:, :, None]
    attention_mask = batch["attention_mask"]
    if config.pad_token == "pad":
        loss_mask = attention_mask
    else:
        loss_mask = jnp.ones_like(attention_mask)
    loss_mask = loss_mask * (1 - batch["cond_seq_mask"])

    denoiser_z = add_noise(x0, noise, t, config, cond_seq_mask=cond_seq_mask)

    drop = batch["label_drop_mask"][:, None]
    if config.label_drop_prob > 0:
        denoiser_z = jnp.where(drop[:, :, None] & (cond_seq_mask > 0), jnp.zeros_like(denoiser_z), denoiser_z)
        x0 = jnp.where(drop[:, :, None] & (cond_seq_mask > 0), jnp.zeros_like(x0), x0)

    decoder_targets = batch["input_ids"]  # (B, S)
    decoder_step_active = jax.random.bernoulli(decoder_step_rng, decoder_prob)

    # Decoder-branch input: logit-normal-noised latent (decoder_z) at t=1
    decoder_lambda_rng, decoder_noise_rng = jax.random.split(decoder_rng)
    decoder_z_vals = (
        jax.random.normal(decoder_lambda_rng, (batch_size * seq_length,))
        * config.decoder_p_std + config.decoder_p_mean
    )
    decoder_lambda_t = jax.nn.sigmoid(decoder_z_vals).reshape(batch_size, seq_length, 1)
    decoder_noise = jax.random.normal(decoder_noise_rng, x0.shape, dtype=x0.dtype) * decoder_noise_scale
    decoder_z = decoder_lambda_t * x0 + (1 - decoder_lambda_t) * decoder_noise

    t_expanded = t.reshape(-1, 1, 1)
    v_target = (x0 - denoiser_z) / jnp.maximum(1 - t_expanded, t_eps)

    if self_cond_prob > 0:
        use_self_cond_mask = (
            (jax.random.uniform(self_cond_mask_rng, (batch_size,)) < self_cond_prob)
            .reshape(-1, 1, 1).astype(x0.dtype)
        )
    else:
        use_self_cond_mask = None

    if config.num_self_cond_cfg_tokens > 0:
        self_cond_cfg_scale = sample_cfg_scale(
            self_cond_cfg_rng, batch_size,
            cfg_min=config.self_cond_cfg_min, cfg_max=config.self_cond_cfg_max,
        )
    else:
        self_cond_cfg_scale = None

    def get_z_input(params, z, t_input, self_cond_cfg_input, x_tokens):
        # Self-conditioning: with probability self_cond_prob, compute initial estimate
        if self_cond_prob == 0:
            return z
        z_uncond = restore_cond(jnp.zeros_like(z), x_tokens, cond_seq_mask)
        z_with_zeros = jnp.concatenate([z, z_uncond], axis=-1)
        net_out_init = state.apply_fn(
            {"params": params}, z_with_zeros, t_input,
            deterministic=True,
            self_cond_cfg_scale=self_cond_cfg_input,
        )
        net_out_init = jax.lax.stop_gradient(net_out_init)
        _, x_pred_init = net_out_to_v_x(net_out_init, z, t_input, t_eps)
        x_pred_init = restore_cond(x_pred_init, x_tokens, cond_seq_mask)
        x_pred_cond = x_pred_init * use_self_cond_mask.astype(z.dtype)
        x_pred_cond = restore_cond(x_pred_cond, x_tokens, cond_seq_mask)
        return jnp.concatenate([z, x_pred_cond], axis=-1)

    def reduce_token_loss(per_token_loss, loss_mask):
        loss_mask = loss_mask.astype(per_token_loss.dtype)
        safe_loss = jnp.where(loss_mask > 0, per_token_loss, jnp.zeros_like(per_token_loss))
        return (safe_loss * loss_mask).sum() / jnp.maximum(loss_mask.sum(), 1.0)

    def get_sc_cond_and_uncond(params, z, t, cond_mask, x_tokens):
        kwargs = {
            "self_cond_cfg_scale": self_cond_cfg_scale,
            "deterministic": True,
        }
        if config.self_cond_prob == 0:
            net_out_uncod = state.apply_fn({"params": params}, z, t, **kwargs)
            v_uncond, _ = net_out_to_v_x(net_out_uncod, z, t, t_eps)
            return v_uncond, v_uncond

        z_uncond = restore_cond(jnp.zeros_like(z), x_tokens, cond_mask)
        z_input_uncond = jnp.concatenate([z, z_uncond], axis=-1)
        net_out_uncond = state.apply_fn({"params": params}, z_input_uncond, t, **kwargs)
        v_uncond, x_uncond = net_out_to_v_x(net_out_uncond, z, t, t_eps)
        x_uncond = restore_cond(x_uncond, x_tokens, cond_mask)

        z_input_cond = jnp.concatenate([z, x_uncond], axis=-1)
        net_out_cond = state.apply_fn({"params": params}, z_input_cond, t, **kwargs)
        v_cond, _ = net_out_to_v_x(net_out_cond, z, t, t_eps)
        return v_cond, v_uncond

    def get_sc_guided_v(params, z, t, base_v_target, x_tokens):
        """v target with self-conditioning guidance."""
        v_cond, v_uncond = get_sc_cond_and_uncond(
            params, z, t, cond_mask=cond_seq_mask, x_tokens=x_tokens
        )
        sc_w = self_cond_cfg_scale.reshape(batch_size, 1, 1)
        sc_guidance = (1 - 1 / sc_w) * (v_cond - v_uncond)
        sc_guidance = jnp.where(use_self_cond_mask, sc_guidance, jnp.zeros_like(sc_guidance))
        return jax.lax.stop_gradient(base_v_target + sc_guidance)

    def get_v_target(params, z, t, base_v_target, x_tokens):
        """Compute final v target with self-conditioning guidance."""
        if config.num_self_cond_cfg_tokens > 0 and config.self_cond_prob > 0:
            return get_sc_guided_v(params, z, t, base_v_target=base_v_target, x_tokens=x_tokens)
        return base_v_target

    def loss_fn(params):

        def _decoder_branch(_):
            # Decoder mode: encoder-noised latent (decoder_z) at t=1, CE loss on tokens.
            decoder_t = jnp.ones_like(t)
            decoder_input = (
                jnp.concatenate([decoder_z, jnp.zeros_like(decoder_z)], axis=-1)
                if config.self_cond_prob > 0 else decoder_z
            )
            _, decoder_logits = state.apply_fn(
                {"params": params}, decoder_input, decoder_t,
                deterministic=False,
                rngs={"dropout": model_dropout_rng},
                self_cond_cfg_scale=self_cond_cfg_scale,
                decoder_step_active=jnp.array(True),
            )
            log_probs = jax.nn.log_softmax(decoder_logits.astype(jnp.float32), axis=-1)
            ce = -jnp.take_along_axis(log_probs, decoder_targets[..., None], axis=-1).squeeze(-1)
            ce_loss = (ce * loss_mask).sum() / jnp.maximum(loss_mask.sum(), 1.0)
            return ce_loss, ce_loss, jnp.zeros(())

        def _denoiser_branch(_):
            # Denoiser mode: x0-noised latent (denoiser_z) at random t, L2 loss on velocity.
            denoiser_t = t
            denoiser_input = get_z_input(
                params, denoiser_z, denoiser_t,
                self_cond_cfg_input=self_cond_cfg_scale,
                x_tokens=x0,
            )
            net_out, _ = state.apply_fn(
                {"params": params}, denoiser_input, denoiser_t,
                deterministic=False,
                rngs={"dropout": model_dropout_rng},
                self_cond_cfg_scale=self_cond_cfg_scale,
                decoder_step_active=jnp.array(False),
            )
            v_pred, _ = net_out_to_v_x(net_out, denoiser_z, denoiser_t, t_eps)
            v_final_target = get_v_target(
                params, denoiser_z, denoiser_t, base_v_target=v_target, x_tokens=x0,
            )
            per_dim_loss = (v_pred - v_final_target) ** 2
            l2_loss = reduce_token_loss(jnp.mean(per_dim_loss, axis=-1), loss_mask)
            return l2_loss, jnp.zeros(()), l2_loss

        loss, ce_loss, l2_loss = jax.lax.cond(
            decoder_step_active, _decoder_branch, _denoiser_branch, None,
        )
        return loss, (l2_loss, ce_loss)

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, (l2_loss_val, ce_loss_val)), grads = grad_fn(state.params)

    grads = jax.lax.pmean(grads, axis_name="batch")
    loss = jax.lax.pmean(loss, axis_name="batch")
    l2_loss_val = jax.lax.pmean(l2_loss_val, axis_name="batch")
    ce_loss_val = jax.lax.pmean(ce_loss_val, axis_name="batch")

    new_state = state.apply_gradients(grads=grads, dropout_rng=new_dropout_rng)

    # Update EMA only on actual optimizer steps, not on gradient accumulation steps.
    # With optax.MultiSteps, params only change every grad_accum_steps mini-batches; updating
    # EMA every mini-batch would make effective decay decay^grad_accum_steps instead of decay.
    def ema_update(ema_params, params, decay):
        return jax.tree_util.tree_map(lambda e, p: e * decay + p * (1 - decay), ema_params, params)

    is_optimizer_step = (new_state.step % config.grad_accum_steps) == 0
    new_ema_params1 = jax.lax.cond(
        is_optimizer_step,
        lambda: ema_update(state.ema_params1, new_state.params, config.ema_decay1),
        lambda: state.ema_params1,
    )
    new_state = new_state.replace(ema_params1=new_ema_params1, dropout_rng=new_dropout_rng)

    # Rescale per-branch losses by their sampling probability so they reflect the
    # per-branch loss rather than the expected loss conditioned on the branch firing.
    decoder_prob_arr = jnp.asarray(decoder_prob, dtype=jnp.float32)
    denoiser_prob_arr = jnp.asarray(1.0 - decoder_prob, dtype=jnp.float32)
    active_ce_loss_val = jnp.where(
        decoder_prob_arr > 0.0, ce_loss_val / decoder_prob_arr, jnp.zeros_like(ce_loss_val),
    )
    active_l2_loss_val = jnp.where(
        denoiser_prob_arr > 0.0, l2_loss_val / denoiser_prob_arr, jnp.zeros_like(l2_loss_val),
    )
    metrics = {
        "loss": loss,
        "l2_loss": active_l2_loss_val,
        "ce_loss": active_ce_loss_val,
    }
    return new_state, metrics