reference/code/flm/algo.py

import os
import collections
import copy
import pickle

import fsspec
import numpy as np
import torch
import torch.nn.functional as F
import wandb
import trainer_base
import utils
import math
import models
from torch.func import functional_call
from models.dit import modulate_fused
import functools
from entmax import entmax_bisect

class AR(trainer_base.TrainerBase):
    def __init__(self, config, tokenizer):
        vocab_size = tokenizer.vocab_size
        if (not hasattr(tokenizer, 'mask_token')
                or tokenizer.mask_token is None):
            self.mask_index = vocab_size
            vocab_size += 1
        else:
            self.mask_index = tokenizer.mask_token_id
        super().__init__(config, tokenizer,
                         vocab_size=vocab_size)
        self.save_hyperparameters()
        self._validate_configuration()

    def _validate_configuration(self):
        super()._validate_configuration()
        assert not self.config.algo.time_conditioning
        assert self.config.prior.type == 'none'

    def _process_model_input(self, x0, valid_tokens):
        input_tokens = x0[:, :-1]
        output_tokens = x0[:, 1:]
        valid_tokens = valid_tokens[:, 1:]
        return input_tokens, output_tokens, valid_tokens

    def nll(self, input_tokens, output_tokens,
            current_accumulation_step):
        del current_accumulation_step
        output = self.backbone(input_tokens, None)
        output[:, :, self.mask_index] = self.neg_infinity
        output = output.log_softmax(-1)
        return - output.gather(
            -1, output_tokens[:, :, None])[:, :, 0]

    def generate_samples(self, num_samples, **kwargs):
        # precompute token buffer
        num_pred_tokens = self.num_tokens - 1
        x = torch.zeros(
            (num_samples, num_pred_tokens + 1),
            dtype=torch.long,
            device=self.device)
        x[:, 0] = self.tokenizer.bos_token_id
        # precompute noise
        noise = (torch.distributions.Gumbel(0, 1)
                 .sample((num_samples, num_pred_tokens, self.vocab_size))
                 .to(self.device))
        if self.config.sampling.use_float64:
            noise = noise.to(torch.float64)
        for i in range(num_pred_tokens):
            output = self.backbone(x[:, :i + 1], None)
            output[:, :, self.mask_index] = self.neg_infinity
            output = output.log_softmax(-1)
            y = (output[:, -1, :] + noise[:, i, :]).argmax(-1)
            x[:, i + 1] = y
        return x

    def _process_sigma(self, sigma):
        del sigma
        return None


class MDLM(trainer_base.AbsorbingState):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self._validate_configuration()

    def _validate_configuration(self):
        # ancestral sampling isn't desirable because it's slow
        assert self.sampler == 'ancestral_cache'

    def _process_model_output(self, model_output, xt, sigma):
        del sigma
        model_output[:, :, self.mask_index] += self.neg_infinity

        # Normalize the model_output such that x.exp() is
        # a probability distribution over vocab_size.
        model_output = model_output - torch.logsumexp(
            model_output, dim=-1, keepdim=True)
        # Apply updates directly in the logits matrix.
        # For the logits of the unmasked tokens, set all values
        # to -infinity except for the indices corresponding to
        # the unmasked tokens.
        unmasked_indices = (xt != self.mask_index)
        model_output[unmasked_indices] = self.neg_infinity
        model_output[unmasked_indices, xt[unmasked_indices]] = 0
        return model_output

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                      dalpha_t, low_var=False):
        del xt
        log_p_theta = torch.gather(
            input=log_x_theta,
            dim=-1,
            index=x0[:, :, None]).squeeze(-1)
        return log_p_theta * dalpha_t / (1 - alpha_t)

    def _get_score(self, x, sigma):
        model_output = self.forward(x, sigma)
        # score(x, t) = p_t(y) / p_t(x)
        # => log score(x, t) = log p_t(y) - log p_t(x)

        # case 1: x = masked
        #   (i) y = unmasked
        #     log score(x, t) = log p_\theta(x)|_y + log k
        #     where k = exp(- sigma) / (1 - exp(- sigma))
        #   (ii) y = masked
        #     log score(x, t) = 0

        # case 2: x = unmasked
        #   (i) y != masked, y != x
        #     log score(x_i, t) = - inf
        #   (ii) y = x
        #     log score(x_i, t) = 0
        #   (iii) y = masked token
        #     log score(x_i, t) = - log k
        #     where k = exp(- sigma) / (1 - exp(- sigma))

        log_k = - torch.log(torch.expm1(sigma)).squeeze(-1)
        assert log_k.ndim == 1

        masked_score = model_output + log_k[:, None, None]
        masked_score[:, :, self.mask_index] = 0

        unmasked_score = self.neg_infinity * torch.ones_like(
            model_output)
        unmasked_score = torch.scatter(
            unmasked_score,
            -1,
            x[..., None],
            torch.zeros_like(unmasked_score[..., :1]))
        unmasked_score[:, :, self.mask_index] = - (
            log_k[:, None] * torch.ones_like(x))

        masked_indices = (x == self.mask_index).to(
            model_output.dtype)[:, :, None]
        model_output = (
            masked_score * masked_indices
            + unmasked_score * (1 - masked_indices))
        return model_output.exp()


class D3PMAbsorb(trainer_base.AbsorbingState):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self._validate_configuration()

    def _validate_configuration(self):
        super()._validate_configuration()
        assert self.noise.type == 'log-linear'
        assert self.parameterization == 'mean'

    def _process_model_output(self, model_output, xt, sigma):
        del xt
        del sigma
        if self.subs_masking:
            model_output[:, :, self.mask_index] += self.neg_infinity
        return model_output.log_softmax(dim=-1)

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                      dalpha_t, low_var=False):
        del dalpha_t
        assert not low_var
        dt = 1 / self.T
        t = 1 - alpha_t  # Only valid for log-linear schedule.
        t = t.clamp(0., 1.0 - 1e-4)
        alpha_t = alpha_t + torch.zeros_like(xt)
        alpha_s = t - dt + torch.zeros_like(xt)
        assert alpha_s.shape == xt.shape
        assert alpha_t.shape == xt.shape
        log_x_theta_at_x0 = torch.gather(
            log_x_theta, -1, x0[:, :, None]).squeeze(-1)
        log_x_theta_at_m = log_x_theta[:, :, self.mask_index]
        x_theta_at_m = log_x_theta_at_m.exp()

        term_1_coef = dt / t
        term_1_log_nr = torch.log(alpha_t * x_theta_at_m / t + 1)
        term_1_log_dr = log_x_theta_at_x0

        term_2_coef = 1 - dt / t
        term_2_log_nr = term_1_log_nr
        term_2_log_dr = torch.log(
            alpha_s * x_theta_at_m / (t - dt) + 1)
        L_vb_masked = (
            term_1_coef * (term_1_log_nr - term_1_log_dr)
            + term_2_coef * (term_2_log_nr - term_2_log_dr))

        diffusion_loss = self.T * L_vb_masked * (xt == self.mask_index)
        return self._reconstruction_loss(x0) + diffusion_loss


class SEDDAbsorb(trainer_base.AbsorbingState):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self._validate_configuration()

    def _validate_configuration(self):
        super()._validate_configuration()
        assert self.config.sampling.predictor == 'analytic'

    def _get_score(self, x, sigma):
        return self.forward(x, sigma).exp()

    def _process_model_output(self, model_output, xt, sigma):
        esigm1_log = torch.where(
            sigma < 0.5,
            torch.expm1(sigma),
            sigma.exp() - 1).log().to(model_output.dtype)
        # logits shape
        # (batch_size, context_length, vocab_size)
        model_output = (model_output
                        - esigm1_log[:, None, None]
                        - np.log(model_output.shape[-1] - 1))
        # The below scatter operation sets the log score
        # for the input word to 0.
        model_output = torch.scatter(
            model_output, -1, xt[..., None],
            torch.zeros_like(model_output[..., :1]))
        return model_output

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                      dalpha_t, low_var=False):
        """Computes the SEDD loss for the Absorbing State Diffusion.

        Args:
          log_x_theta: float torch.Tensor with shape (batch_size,
              context_length, vocab_size),
              log score, output of the denoising network.
          xt: int torch.Tensor with shape (batch_size,
              context_length), input.
          x0: int torch.Tensor with shape (batch_size,
              context_length), input.
          alpha_t: float torch.Tensor with shape (batch_size, 1),
              signal level.
          alpha_t: float torch.Tensor with shape (batch_size, 1),
              signal level.
          dalpha_t: float or float torch.Tensor with shape (batch_size, 1),
              time derivative of signal level.
          low_var: bool, low variance loss during training.

        Returns:
          loss with shape (batch_size, context_length).
        """
        assert not low_var
        masked_indices = xt == self.mask_index
        sigma = self._sigma_from_alphat(alpha_t)
        dsigma = - dalpha_t / alpha_t

        expsig_minus_1 = torch.expm1(sigma).expand_as(xt)
        q_ratio = 1 / expsig_minus_1[masked_indices]

        words_that_were_masked = x0[masked_indices]

        neg_term = q_ratio * torch.gather(
            log_x_theta[masked_indices],
            -1,
            words_that_were_masked[..., None]).squeeze(-1)
        score = log_x_theta[masked_indices].exp()
        if self.mask_index == self.vocab_size - 1:
            pos_term = score[:, :-1].sum(dim=-1)
        else:
            pos_term = score[:, : self.mask_index].sum(
                dim=-1) + score[:, self.mask_index + 1:].sum(dim=-1)
        const = q_ratio * (q_ratio.log() - 1)

        entropy = torch.zeros(* xt.shape, device=xt.device)
        entropy[masked_indices] += pos_term - neg_term + const
        return dsigma * entropy


def stopgrad(x):
    """Stop gradient for x."""
    return x.detach()


def adaptive_l2_loss(error, gamma=0.5, c=1e-3):
    """
    Adaptive L2 loss: sg(w) * ||Δ||_2^2, where w = 1 / (||Δ||^2 + c)^p, p = 1 - γ
    """
    delta_sq = torch.mean(error ** 2, dim=(1, 2), keepdim=False)  # (B,)
    p = 1.0 - gamma
    w = 1.0 / (delta_sq + c).pow(p)
    loss = delta_sq  # ||Δ||^2
    return (stopgrad(w) * loss).mean()


def mse_loss(error):
    per_sample = (error ** 2).mean(dim=(1, 2))  # [B]
    return per_sample.mean()


class DUO_BASE(trainer_base.UniformState):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self._validate_configuration()

    def on_save_checkpoint(self, checkpoint):
        checkpoint['state_dict'] = collections.OrderedDict(
            (k, v) for k, v in checkpoint['state_dict'].items()
            if not k.startswith('teacher'))
        super().on_save_checkpoint(checkpoint)

    def on_load_checkpoint(self, checkpoint):
        # Handle _orig_mod prefix from torch.compile and filter teacher keys
        new_state_dict = collections.OrderedDict()
        for k, v in checkpoint['state_dict'].items():
            # Filter out teacher keys
            if k.startswith('teacher'):
                continue
            # Strip _orig_mod prefix from torch.compile
            new_key = k.replace('._orig_mod.', '.')
            new_state_dict[new_key] = v
        checkpoint['state_dict'] = new_state_dict
        super().on_load_checkpoint(checkpoint)

    def _process_model_output(self, model_output, xt, sigma):
        del xt, sigma
        return model_output.log_softmax(dim=-1)

    def _compute_posterior(self, x, xt, alpha_s, alpha_t):
        """Computes the posterior / approximate posterior.

        Args:
          x: Either clean input `x0` (one-hot),
            or model's predicted `x_theta` of shape (B, L, V).
          xt: The noisy latent (as indices) of shape (B, L).
          alpha_s: Noise level at s of shape (B, [L | 1], 1).
          alpha_t: Noise level at t of shape (B, [L | 1], 1).

        Returns:
          Posterior / approximate posterior of shape (B, L, V).
        """
        if self.config.sampling.use_float64:
            x = x.to(torch.float64)
        if alpha_s.ndim == 2:
            alpha_s = alpha_s.unsqueeze(-1)
        if alpha_t.ndim == 2:
            alpha_t = alpha_t.unsqueeze(-1)
        alpha_ts = alpha_t / alpha_s
        d_alpha = alpha_s - alpha_t
        xt_one_hot = F.one_hot(xt, self.vocab_size).to(
            self.dtype).to(self.device)
        return (
            (alpha_t * self.vocab_size * x * xt_one_hot + (
                alpha_ts - alpha_t) * xt_one_hot + d_alpha * x + (
                1 - alpha_ts) * (1 - alpha_s) / self.vocab_size) / (
                alpha_t * self.vocab_size * torch.gather(
                    x, -1, xt[..., None]) + (1 - alpha_t)))

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                      dalpha_t, low_var=False):  # Computes Eq 5.
        assert alpha_t.ndim == 2
        assert x0.ndim == 2
        assert xt.ndim == 2
        if torch.is_tensor(dalpha_t) and dalpha_t.ndim == 1:
            dalpha_t = dalpha_t.unsqueeze(-1)
        assert not torch.is_tensor(dalpha_t) or dalpha_t.ndim == 2
        x_reconst = log_x_theta.exp()  # convert logits to probabilities
        x_bar_theta = self.vocab_size * alpha_t[
            :, :, None] * x_reconst + 1 - alpha_t[:, :, None]
        coeff = dalpha_t / (self.vocab_size * alpha_t)
        x_eq_xt = (x0 == xt).float()
        x_neq_xt = 1 - x_eq_xt
        xbar_xt = (1 - alpha_t) + self.vocab_size * alpha_t * x_eq_xt
        xbar_theta_xt = torch.gather(
            x_bar_theta, -1, xt.unsqueeze(-1)).squeeze(-1)
        xbar_theta_x = torch.gather(
            x_bar_theta, -1, x0.unsqueeze(-1)).squeeze(-1)
        term1 = self.vocab_size * (1 / xbar_xt
                                   - 1 / xbar_theta_xt)  # Eq 5. term 1

        const = (1 - alpha_t) / (self.vocab_size * alpha_t
                                 + 1 - alpha_t)
        term2_coefs = x_eq_xt * const + x_neq_xt
        term2_offset = ((self.vocab_size - 1) * const * x_eq_xt
                        - (1 / const) * x_neq_xt) * const.log()
        term2_theta = - term2_coefs * (
            x_bar_theta.log().sum(-1)
            - self.vocab_size * xbar_theta_xt.log())
        term2_theta = (
            term2_theta
            - self.vocab_size * alpha_t / (1 - alpha_t) * (
                xbar_theta_x.log() - xbar_theta_xt.log()) * x_neq_xt)
        term2 = term2_theta + term2_offset
        diffusion_loss = coeff * (term1 - term2)
        assert diffusion_loss.ndim == 2
        return diffusion_loss

    def _ancestral_update(self, x, t, dt, p_x0=None,
                          noise_removal_step=False, step_index=None):
        del p_x0
        _, alpha_t = self.noise(t)
        if noise_removal_step:
            alpha_s = torch.ones_like(alpha_t)
        else:
            _, alpha_s = self.noise(t - dt)
        sigma_t = self._sigma_from_alphat(alpha_t)

        assert alpha_t.ndim == 2

        q_xs = self._compute_posterior(
            x=self.forward(x, sigma_t).exp(),
            xt=x,
            alpha_s=alpha_s,
            alpha_t=alpha_t)
        if self.p_nucleus < 1:
            q_xs = utils.top_k_top_p_filtering(
                q_xs.log(), top_p=self.p_nucleus)
        return None, trainer_base.sample_categorical(q_xs, self.config.sampling.temperature)


class Integral(torch.autograd.Function):
    """
    torch module calculating UDLM's p_t 
    """

    @staticmethod
    def forward(ctx, gamma_t, data):
        gamma_max = data['gamma_max']
        gamma_min = data['gamma_min']
        if (gamma_t.max() > gamma_max) or (
                gamma_t.min() < gamma_min):
            # print('max:{} {}'.format(gamma_t.max(), gamma_max))
            # print('min:{} {}'.format(gamma_t.min(), gamma_min))
            gamma_t = torch.clip(gamma_t, gamma_min, gamma_max)
        indices = torch.round(
            (data['num_points'] - 1) * (gamma_t - gamma_min) / (
                gamma_max - gamma_min)).long()
        grad_pt = data['grad_pt']
        ctx.grad_pt = grad_pt[indices]

        pt = data['pt'][indices]
        assert pt.shape == gamma_t.shape
        return pt

    @staticmethod
    def backward(ctx, grad_output):
        return ctx.grad_pt * grad_output, None


class DUO(DUO_BASE):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        with fsspec.open(self.config.algo.integral_cache_path,
                                         'rb') as f:
            self.integral_cache = pickle.load(f)
        self.integral_cache['pt'] = torch.from_numpy(
            self.integral_cache['pt'])
        self.integral_cache['grad_pt'] = torch.from_numpy(
            self.integral_cache['grad_pt'])
        self.gamma_min = self.config.algo.gamma_min
        self.gamma_max = self.config.algo.gamma_max
        self.gumbel_tau_log10_start = self.config.algo.gumbel_tau_log10_start
        self.gumbel_tau_log10_end = self.config.algo.gumbel_tau_log10_end
        self.curriculum_start = self.config.algo.curriculum_start
        self.curriculum_end = self.config.algo.curriculum_end
        self.loss_type = self.config.algo.loss_type
        self._validate_configuration()

    def to(self, *args, **kwargs):
        self = super().to(*args, **kwargs)
        self.integral_cache['pt'] = self.integral_cache[
            'pt'].to(*args, **kwargs)
        self.integral_cache['grad_pt'] = self.integral_cache[
            'grad_pt'].to(*args, **kwargs)
        return self

    def _compute_gumbel_tau_inverse(self):
        start = self.gumbel_tau_log10_start
        end = self.gumbel_tau_log10_end
        delta = end - start
        if self.global_step < self.curriculum_start:
            tau = start
        elif self.global_step < self.curriculum_end:
            frac = (self.global_step - self.curriculum_start) / (
                self.curriculum_end - self.curriculum_start)
            tau = start + frac * delta
        else:
            tau = -10
        return 10 ** (-tau)

    def training_step(self, batch, batch_idx):
        self.log(name='gumbel_tau_log10',
                         value=1 / self._compute_gumbel_tau_inverse(),
                         on_step=True,
                         on_epoch=False,
                         sync_dist=True)
        return super().training_step(batch, batch_idx)

    def _gamma_to_alphat(self, gamma_t): # eq 10.
        integral = Integral.apply(gamma_t, self.integral_cache)
        return (self.vocab_size * integral - 1) / (
            self.vocab_size - 1)

    def _prior_loss(self):
        alpha_1 = self._gamma_to_alphat(
            torch.tensor(self.gamma_max))
        loss = ((alpha_1 + (1 - alpha_1) / self.vocab_size) * torch.log(
            (self.vocab_size - 1) * alpha_1 + 1) + (
                1 - 1 / self.vocab_size) * (1 - alpha_1) * torch.log(1 - alpha_1))
        return loss.item()

    def _q_xt_gaussian(self, x, gamma_t):
        """Computes the noisy sample xt."""
        assert gamma_t.ndim == 1
        assert x.ndim == 3
        gamma_t = gamma_t.unsqueeze(-1).unsqueeze(-1)
        alpha_t = torch.sigmoid(-gamma_t).sqrt()
        sigma_t = torch.sigmoid(gamma_t).sqrt()
        epsilon = torch.randn(x.shape, dtype=torch.float32,
                                                    device=self.device)
        return alpha_t * x + sigma_t * epsilon

    def nll(self, x0, output_tokens,
                    current_accumulation_step=None, train_mode=False, xT=None, **kwargs):
        # TODO: use xT
        use_true_nll = (self.global_step > self.curriculum_end
                                        or not train_mode)
        if use_true_nll:
            return super().nll(x0, output_tokens,
                                                 current_accumulation_step)
        del output_tokens
        t = self._sample_t(x0.shape[0], current_accumulation_step)
        gamma_t = self.gamma_min + t * (self.gamma_max - self.gamma_min)    
        gamma_t_prime = self.gamma_max - self.gamma_min
        alpha_t = self._gamma_to_alphat(gamma_t)
        T = 1000
        dalpha_t = gamma_t_prime * T * (
            self._gamma_to_alphat(gamma_t + 1 / T) - alpha_t)
        alpha_t = alpha_t.unsqueeze(-1)
        dalpha_t = dalpha_t.unsqueeze(-1)
        assert alpha_t.ndim == 2
        sigma = self._sigma_from_alphat(alpha_t)

        x0_one_hot = F.one_hot(x0, self.vocab_size)
        xt = self._q_xt_gaussian(x0_one_hot, gamma_t)
        xt = xt * self._compute_gumbel_tau_inverse()
        xt_usdm = xt.argmax(-1)
        log_x_theta = self.forward(xt, sigma=sigma)

        return self.nll_per_token(log_x_theta=log_x_theta,
                                                            xt=xt_usdm,
                                                            x0=x0,
                                                            alpha_t=alpha_t,
                                                            dalpha_t=dalpha_t,
                                                            low_var=False)




class Distillation(DUO):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self.update_teacher_every = config.algo.update_teacher_every
        self.save_hyperparameters()
        self.teacher = None
        self.teacher_ema = config.algo.teacher_ema
        self.linear_growth_dt = config.algo.linear_growth_dt
        self.linear_growth_min = config.algo.linear_growth_min
        self.linear_growth_max = config.algo.linear_growth_max

    def _validate_configuration(self):
        assert os.path.exists(
            self.config.algo.integral_cache_path), (
            'The integral cache (Eq. 10 in the paper) for '
            f'the {self.config.data.tokenizer_name_or_path} '
            ' tokenizer doesnt exist at '
            f'{self.config.algo.integral_cache_path}. '
            'Please generate it by running the utils.py script, '
            'and ensure the correct path is specified using the '
            'algo.integral_cache_path flag.')
        assert self.loss_type in {
            'kl-fwd', 'kl-bwd', 'posterior', 'kl-posterior'}

    def _maybe_update_teacher_weights(self):
        if self.global_step % self.update_teacher_every != 0:
            return
        if self.teacher_ema:
            self.ema.copy_to(self.teacher.parameters())
        else:
            for better_param, current_param in zip(
                    self.backbone.parameters(), self.teacher.parameters()):
                if current_param.requires_grad:
                    current_param.data.copy_(better_param.data)

    @torch.no_grad()
    def _teacher_logits(self, xt, sigma):
        if self.teacher is None:
            self.teacher = copy.deepcopy(self.backbone)
        self._maybe_update_teacher_weights()

        sigma = self._process_sigma(sigma)
        with torch.cuda.amp.autocast(dtype=torch.float32):
            model_output = self.teacher(xt, sigma)
        logits = self._process_model_output(
            model_output=model_output, xt=xt, sigma=sigma)
        return logits.detach()

    def _sample_trajectory(self, x0, gamma_t, gamma_s):
        """Computes the noisy sample xt."""
        assert gamma_t.ndim == 1
        assert gamma_s.ndim == 1
        assert x0.ndim == 2
        x0 = F.one_hot(x0, self.vocab_size).to(
            self.dtype).to(self.device)
        gamma_t = gamma_t.unsqueeze(-1).unsqueeze(-1)
        alpha_t = torch.sigmoid(-gamma_t).sqrt()
        sigma_t = torch.sigmoid(gamma_t).sqrt()

        gamma_s = gamma_s.unsqueeze(-1).unsqueeze(-1)
        alpha_s = torch.sigmoid(-gamma_s).sqrt()
        sigma_s = torch.sigmoid(gamma_s).sqrt()

        epsilon = torch.randn(x0.shape, dtype=torch.float32,
                              device=self.device)
        xt = alpha_t * x0 + sigma_t * epsilon
        xs = alpha_s * x0 + sigma_s * epsilon
        return xt, xs

    def _compute_dt(self):
        if self.linear_growth_dt:
            scale = self.global_step / self.trainer.max_steps
            return self.linear_growth_min + scale * (
                self.linear_growth_max - self.linear_growth_min)
        n = self.global_step // self.update_teacher_every
        return 2 ** n / self.T

    def nll(self, x0, output_tokens,
            current_accumulation_step=None, train_mode=None, xT=None):
        # TODO: use xT
        del output_tokens, train_mode
        t = self._sample_t(x0.shape[0], current_accumulation_step)
        dt = self._compute_dt()
        t = torch.clip(t + dt, 0, 1)

        gamma_t = self.gamma_min + t * (self.gamma_max
                                        - self.gamma_min)
        gamma_s = self.gamma_min + (
            t - dt) * (self.gamma_max - self.gamma_min)

        alpha_t = self._gamma_to_alphat(gamma_t)
        alpha_t = alpha_t.unsqueeze(-1)
        assert alpha_t.ndim == 2
        usdm_alpha_s = self._gamma_to_alphat(gamma_s)
        usdm_alpha_s = usdm_alpha_s.unsqueeze(-1)
        assert usdm_alpha_s.ndim == 2

        xt, xs = self._sample_trajectory(x0, gamma_t, gamma_s)
        xt_discrete = xt.argmax(-1)
        xs_discrete = xs.argmax(-1)
        log_x_theta_student = self.forward(
            xt_discrete, sigma=self._sigma_from_alphat(alpha_t))
        log_x_theta_teacher = self._teacher_logits(
            xs_discrete, sigma=self._sigma_from_alphat(usdm_alpha_s))
        if self.config.training.loss_precision == 'float64':
            log_x_theta_student = log_x_theta_student.to(torch.float64)
            log_x_theta_teacher = log_x_theta_teacher.to(torch.float64)
        if self.loss_type == 'kl-fwd':
            return (log_x_theta_teacher.exp() * (
                log_x_theta_teacher - log_x_theta_student)).sum(-1)
        elif self.loss_type == 'kl-bwd':
            return (log_x_theta_student.exp() * (
                log_x_theta_student - log_x_theta_teacher)).sum(-1)

    def training_step(self, batch, batch_idx):
        self.log(name='dt',
                 value=self._compute_dt(),
                 on_step=True,
                 on_epoch=False,
                 sync_dist=True)
        return super().training_step(batch, batch_idx)


class Rectification(DUO):  # Training as duo, without curriculum
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self.save_hyperparameters()
        self.use_linear_schedule = config.algo.use_linear_schedule
        self.use_simple_loss = config.algo.use_simple_loss
        self.onestep_mode = config.algo.onestep_mode
        self.debug = getattr(config.algo, 'debug', False)

    def _compute_gumbel_tau_inverse(self):
        return 1e-10

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                      dalpha_t, low_var=False, simple_loss=False):
        if simple_loss:
            loss = F.cross_entropy(
                log_x_theta.view(-1, self.vocab_size),
                x0.view(-1),
                reduction='none')
            loss = loss.view(xt.shape)
            return loss
        else:
            return super().nll_per_token(
                log_x_theta=log_x_theta,
                xt=xt,
                x0=x0,
                alpha_t=alpha_t,
                dalpha_t=dalpha_t,
                low_var=low_var
            )

    def nll(self, x0, output_tokens,
            current_accumulation_step=None, train_mode=False, xT=None, given_t=None, not_sampling_t=False):
        del output_tokens
        if given_t is not None:
            if not_sampling_t:
                assert torch.is_tensor(given_t)
                t = 1-given_t
            else:
                t = self._sample_t(
                    x0.shape[0], current_accumulation_step, given_t=1-given_t)
        else:
            t = self._sample_t(x0.shape[0], current_accumulation_step)
        assert t.shape[0] == x0.shape[0]
        if self.T > 0:
            assert 0

        dalpha_t, alpha_t = self.noise(t)

        alpha_t = alpha_t.unsqueeze(-1)
        dalpha_t = dalpha_t.unsqueeze(-1)
        assert alpha_t.ndim == 2
        sigma = self._sigma_from_alphat(alpha_t)

        if given_t is not None and xT is not None:
            # x0 with alpha_t, xT with (1-alpha_t)
            random = torch.rand_like(x0, dtype=torch.float32)
            given_t = given_t.unsqueeze(1)
            random = given_t + random * (1 - given_t)
            if self.onestep_mode:
                # always larger than alpha_t
                random = torch.ones_like(random) + 1
            xt = torch.where(random <= alpha_t, x0, xT)
        elif xT is None or self.debug:
            if not self.debug:
                assert not self.training, 'xT should be provided during training'
            xT = self.prior_sample(x0.shape[0], x0.shape[1])
            random = torch.rand_like(x0, dtype=torch.float32)
            if self.onestep_mode:
                # always larger than alpha_t
                random = torch.ones_like(random) + 1
            xt = torch.where(random <= alpha_t, x0, xT)
        else:
            # x0 with alpha_t, xT with (1-alpha_t)
            random = torch.rand_like(x0, dtype=torch.float32)
            if self.onestep_mode:
                # always larger than alpha_t
                random = torch.ones_like(random) + 1
            xt = torch.where(random <= alpha_t, x0, xT)

        log_x_theta = self.forward(xt, sigma=sigma)

        return self.nll_per_token(log_x_theta=log_x_theta,
                                  xt=xt,
                                  x0=x0,
                                  alpha_t=alpha_t,
                                  dalpha_t=dalpha_t,
                                  low_var=train_mode and self.loss_type == 'low_var',
                                  simple_loss=self.use_simple_loss,
                                  )

class FLMBase(trainer_base.TrainerBase):
    """Base class for FLM/FMLM.
    """

    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self.t_min = config.algo.t_min
        self.t_max = config.algo.t_max
        self.lut_a2g, self.lut_g2a = utils.build_luts(K=self.vocab_size)
        self._is_resuming = (
            config.checkpointing.resume_from_ckpt
            and config.checkpointing.resume_ckpt_path is not None
            and utils.fsspec_exists(config.checkpointing.resume_ckpt_path)
        )

    def _validate_configuration(self):
        pass

    def training_step(self, batch, batch_idx):
        return super().training_step(batch, batch_idx)

    def _process_sigma(self, sigma):
        if sigma.ndim == 1:
            sigma = sigma.unsqueeze(-1)
        assert sigma.ndim == 2
        sigma = sigma.mean(-1).squeeze()
        if sigma.ndim == 0:
            sigma = sigma.unsqueeze(0)
        if not self.config.algo.time_conditioning:
            sigma = torch.zeros_like(sigma)
        assert sigma.ndim == 1, sigma.shape
        return sigma

    def _process_model_output(self, model_output, xt, sigma, cap_value = 30.0):
        del xt, sigma
        model_output = cap_value * torch.tanh(model_output / cap_value)
        return model_output.log_softmax(dim=-1)

    def _process_model_input(self, x0, valid_tokens):
        return x0, None, valid_tokens

    def _loss(self, x0, valid_tokens,
              current_accumulation_step=None,
              train_mode=False,
              xT=None, given_t=None, not_sampling_t=False):
        """Override to always dispatch to self.loss() for all FLM classes."""
        (input_tokens, output_tokens,
         valid_tokens) = self._process_model_input(x0, valid_tokens)
        loss = self.loss(input_tokens, output_tokens,
                         current_accumulation_step, train_mode,
                         xT=xT, given_t=given_t,
                         not_sampling_t=not_sampling_t)
        assert loss.ndim == 2
        if self.ignore_bos:
            loss[:, 1:] = loss[:, 1:]
            valid_tokens[:, 1:] = valid_tokens[:, 1:]

        nlls = (loss * valid_tokens).sum()
        num_tokens = valid_tokens.sum()
        token_nll = nlls / num_tokens
        return trainer_base.Loss(loss=token_nll,
                                 nlls=nlls,
                                 prior_loss=0.0,
                                 num_tokens=num_tokens)

    def loss(self, x0, output_tokens,
             current_accumulation_step=None, train_mode=False,
             xT=None, given_t=None, not_sampling_t=False):
        raise NotImplementedError

    def nll(self, input_tokens, output_tokens,
            current_accumulation_step=None, train_mode=False):
        raise NotImplementedError

    def _sample_t_interval(self, n, accum_step, t_min=None, t_max=None):
        if t_min is None:
            t_min = self.t_min
        if t_max is None:
            t_max = self.t_max
        if accum_step is not None:
            batch_dim = n
            n = self.config.loader.global_batch_size
        _eps_t = torch.rand(n, device=self.device)
        if self.antithetic_sampling:
            offset = torch.arange(n, device=self.device) / n
            _eps_t = (_eps_t / n + offset) % 1
            perm = torch.randperm(n, device=self.device)
            _eps_t = _eps_t[perm]
        t = (t_max - t_min) * _eps_t + t_min
        if accum_step is not None:
            t = t.chunk(self.trainer.num_nodes)[self.trainer.node_rank]
            t = t.chunk(self.trainer.num_devices)[self.trainer.local_rank]
            t = t.chunk(self.trainer.accumulate_grad_batches)[accum_step]
            t = t[:batch_dim]
        return t

    def _tau_to_t(self, tau):
        """Convert t to reparameterized time tau."""
        return utils.alpha_to_gamma(tau, self.lut_a2g)

    def _t_to_tau(self, t):
        """Convert t to reparameterized time tau."""
        return utils.gamma_to_alpha(t, self.lut_g2a)

    def corrupt_continuous(self, x0, t):
        """Corrupt data x0 at time t using linear interpolation with Gaussian noise."""
        t = t.unsqueeze(-1).unsqueeze(-1)
        target_data = F.one_hot(x0, self.vocab_size).float()
        noise = torch.randn_like(target_data, dtype=torch.float32)
        x_t = (1 - t) * noise + t * target_data
        return x_t, target_data

    def load_state_dict(self, state_dict, strict=True):
        return super().load_state_dict(state_dict, strict=False)

    def on_load_checkpoint(self, checkpoint):
        print("Resuming training from checkpoint...")
        self._is_resuming = True
        if 'state_dict' in checkpoint:
            checkpoint['state_dict'] = self._filter_checkpoint_state_dict(
                checkpoint['state_dict'])
        if self.config.mode == 'sample_eval':
            if getattr(self.backbone, 'learnable_loss_weighting', None) is not None:
                if not any(k.startswith('backbone.learnable_loss_weighting')
                           for k in checkpoint['state_dict'].keys()):
                    print("Learnable_loss_weighting not found in checkpoint. "
                          "Initializing from scratch for eval mode.")
                    for name, param in self.backbone.learnable_loss_weighting.named_parameters():
                        param_key = f'backbone.learnable_loss_weighting.{name}'
                        checkpoint['state_dict'][param_key] = param.data.clone()
        super().on_load_checkpoint(checkpoint)

    def on_save_checkpoint(self, checkpoint):
        checkpoint['state_dict'] = collections.OrderedDict(
            (k, v) for k, v in checkpoint['state_dict'].items()
            if not k.startswith('teacher'))
        super().on_save_checkpoint(checkpoint)

    def _filter_checkpoint_state_dict(self, state_dict):
        """Filter teacher keys and strip _orig_mod from checkpoint state_dict."""
        new_state_dict = collections.OrderedDict()
        for k, v in state_dict.items():
            if k.startswith('teacher'):
                continue
            new_key = k.replace('._orig_mod.', '.')
            new_state_dict[new_key] = v
        return new_state_dict

    def forward_no_softmax(self, xt, tau, tau_prime=None, **kwargs):
        tau = self._process_sigma(tau)
        if tau_prime is not None:
            tau_prime = self._process_sigma(tau_prime)
        with torch.amp.autocast(device_type=self.device.type, dtype=torch.float32):
            model_output = self.backbone(xt, tau, tau_prime, **kwargs)
        return model_output

    def _extract_ema_state_dict(self, model, checkpoint):
        """Extract EMA parameters from checkpoint into a state_dict for model."""
        ema_state = checkpoint.get('ema', None)
        if not ema_state:
            print("Warning: No EMA found, using regular state_dict")
            return {k.replace('backbone.', '').replace('._orig_mod.', ''): v
                    for k, v in checkpoint['state_dict'].items()
                    if k.startswith('backbone.')}

        new_sd = collections.OrderedDict()
        shadow_params = ema_state['shadow_params']
        param_names = [n for n, p in model.named_parameters()
                       if p.requires_grad]
        print(f"EMA shadow_params: {len(shadow_params)}, "
              f"Model param_names: {len(param_names)}")
        min_len = min(len(shadow_params), len(param_names))
        for name, val in zip(param_names[:min_len],
                             shadow_params[:min_len]):
            new_sd[name] = val
        for k, v in checkpoint['state_dict'].items():
            clean_k = k.replace('backbone.', '').replace('._orig_mod.', '')
            if (clean_k not in new_sd
                    and clean_k in [n for n, _ in model.named_parameters()]):
                new_sd[clean_k] = v
                print(f"Loaded missing param from state_dict: {clean_k}")
        if len(shadow_params) != len(param_names):
            print(f"Warning: EMA param count mismatch. "
                  f"Loaded {min_len}/{len(param_names)} from EMA, "
                  f"rest from state_dict")
        return new_sd

    def _load_teacher_model(self, path, use_plain_config=True):
        """Load a frozen teacher model from checkpoint.

        Args:
            path: Path to checkpoint file.
            use_plain_config: If True, temporarily disable double_temb and
                learnable_loss_weighting when building the teacher
                (to match EMA parameter shapes from a base model).
        """
        print(f"Loading teacher model from: {path}")
        if use_plain_config:
            saved = (self.config.algo.double_temb,
                     self.config.algo.learnable_loss_weighting)
            self.config.algo.double_temb = False
            self.config.algo.learnable_loss_weighting = False

        assert self.config.algo.backbone == 'dit', \
            "Only DIT backbone supported for teacher model"
        model = models.dit.DIT(self.config, vocab_size=self.vocab_size)

        if use_plain_config:
            (self.config.algo.double_temb,
             self.config.algo.learnable_loss_weighting) = saved

        checkpoint = torch.load(path, map_location='cpu', weights_only=False)
        state_dict = self._extract_ema_state_dict(model, checkpoint)
        model.load_state_dict(state_dict, strict=False)
        model = model.to(self.device).eval()
        for param in model.parameters():
            param.requires_grad = False
        return model

    def _copy_teacher_weights_to_student(self, teacher_dict):
        """Copy teacher weights to student backbone and zero-init sigma_map_prime."""
        with torch.no_grad():
            student_dict = self.backbone.state_dict()
            for name, param in teacher_dict.items():
                print(f"Copying parameter: {name}")
                if name in student_dict:
                    student_dict[name].copy_(param)
            if (hasattr(self.backbone, 'sigma_map_prime')
                    and self.backbone.sigma_map_prime is not None):
                for name, param in self.backbone.sigma_map_prime.named_parameters():
                    if 'mlp.2' in name:
                        param.zero_()
                        print(f"Zero initialized student sigma_map_prime: {name}")

    @staticmethod
    def _zero_init_module(module):
        for m in module.modules():
            if isinstance(m, torch.nn.Linear):
                m.weight.data.zero_()
                if m.bias is not None:
                    m.bias.data.zero_()

    @staticmethod
    def _random_init_module(module, std=0.02):
        for m in module.modules():
            if isinstance(m, torch.nn.Linear):
                m.weight.data.normal_(mean=0.0, std=std)
                if m.bias is not None:
                    m.bias.data.zero_()


class FLM(FLMBase):
    def loss(self, x0, output_tokens,
             current_accumulation_step=None, train_mode=False,
             xT=None, given_t=None, not_sampling_t=False):
        del given_t, not_sampling_t, output_tokens
        B = x0.shape[0]
        tau_t = self._sample_t_interval(B, current_accumulation_step,
                                    t_min=self.t_min, t_max=self.t_max)
        t = self._tau_to_t(tau_t)
        x_t, target_data = self.corrupt_continuous(x0, t)
        f = self.forward(x_t, tau_t) #condition on tau_t
        loss = -(target_data * f).sum(dim=-1)
        self.log('loss', loss.mean(), prog_bar=True)
        if self.config.algo.learnable_loss_weighting is True:
            loss_weight = self.backbone.learnable_loss_weighting(tau_t)
            loss_weight = loss_weight.unsqueeze(-1)
            loss = torch.exp(-loss_weight) * loss + loss_weight
            self.log('loss_weighted', loss.mean(), prog_bar=True)
        return loss

    @torch.no_grad()
    def generate_samples(self, num_samples, num_steps=None, eps=1e-5):
        """Generate samples using Euler ODE solver."""
        if num_steps is None:
            num_steps = self.config.sampling.steps
        B = num_samples
        V = self.vocab_size
        L = self.num_tokens
        device = self.device

        tau_vals = torch.linspace(0.0, 1.0, num_steps + 1, device=device)
        z = torch.randn((num_samples, L, V), device=device, dtype=self.dtype)

        for i in range(num_steps):
            tau_t_curr = tau_vals[i]
            tau_t_next = tau_vals[i + 1]
            tau_t_in = tau_t_curr.expand(B)
            t_in = self._tau_to_t(tau_t_in)
            dt = self._tau_to_t(tau_t_next.expand(B)) - t_in
            x_1_pred = self.forward(z, tau_t_in)
            x_1_pred_probs = x_1_pred.exp()

            if i == num_steps - 1:
                z = x_1_pred_probs
                break

            v = (x_1_pred_probs - z) / (1.0 - t_in.view(-1, 1, 1) + 1e-5)
            z = z + dt.view(-1, 1, 1) * v

        return z.argmax(dim=-1)
    
class FMLM(FLMBase):
    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self._validate_configuration()
        self.teacher_model = None

    def setup(self, stage: str):
        if self.teacher_model is None:
            self._initialize_teacher()
        if stage == 'fit' and not self._is_resuming:
            print(">>> Initializing student from teacher...")
            self._initialize_student_from_teacher()
        elif self._is_resuming:
            print(">>> Skipping student initialization (resuming from checkpoint).")

    def _initialize_teacher(self):
        path = self.config.algo.teacher_path
        if path is None or path == "":
            print("No teacher model specified, skipping teacher initialization")
            return
        self.teacher_model = self._load_teacher_model(path, use_plain_config=True)

    def _initialize_student_from_teacher(self):
        if self.teacher_model is None or not self.config.algo.initialize_student_from_teacher:
            return
        self._copy_teacher_weights_to_student(self.teacher_model.state_dict())

    def _validate_configuration(self):
        assert self.config.algo.double_temb == True, \
            "FMLM denoiser requires double time-emb conditioning to be True"

        assert type(self.config.algo.diagonal_fraction) == float, \
            "diagonal_fraction must be a float"
        assert 0 <= self.config.algo.diagonal_fraction <= 1, \
            "diagonal_fraction must be between 0 and 1"
        assert self.config.algo.distillation_method in ["PSD", "LSD", "ESD"], \
            "FMLM denoiser must distill using PSD, LSD, or ESD"
        if self.config.algo.distillation_method == "LSD":
            assert 2 >= self.config.algo.entmax_temp_lsd >= 1, \
                "entmax_temp_lsd must be in [1,2] for LSD distillation"
            assert not self.config.algo.backprop_entmax_temp_lsd, \
                "backprop_entmax_temp_lsd is not coded for LSD distillation"

    def forward(self, xt, sigma, sigma_prime=None, use_jvp_attn=False, **kwargs):
        return super().forward(xt, sigma, sigma_prime, use_jvp_attn=use_jvp_attn, **kwargs)

    def forward_with_ema(self, *args, **kwargs):
        ema_to_use = self.ema
        assert ema_to_use is not None, "EMA must be available"
        ema_to_use.store(self._get_parameters())
        ema_to_use.copy_to(self._get_parameters())
        try:
            with torch.no_grad():
                self.backbone.eval()
                out = self.forward(*args, **kwargs)
            return out
        finally:
            ema_to_use.restore(self._get_parameters())
            self.backbone.train()

    def teacher_forward(self, xt, tau=None, d=None, use_jvp_attn=False):
        del d, use_jvp_attn
        sigma = tau.unsqueeze(-1) if tau.ndim == 1 else tau
        sigma = self._process_sigma(sigma)
        with torch.no_grad():
            with torch.amp.autocast(device_type=self.device.type, dtype=torch.float32):
                model_output = self.teacher_model(xt, sigma)
        return self._process_model_output(model_output=model_output, xt=xt, sigma=sigma)

    def _d_tau_by_d_t(self, t):
        return utils.d_alpha_by_d_gamma(t, self.lut_g2a)

    def _sample_multi_t_interval(self, n, accum_step, num_per_sample, t_min=None, t_max=None):
        """Sample num_per_sample times from [t_min, t_max], returned sorted."""
        if t_min is None:
            t_min = self.t_min
        if t_max is None:
            t_max = self.t_max

        if accum_step is not None:
            batch_dim = n
            total_n = self.config.loader.global_batch_size * num_per_sample
        else:
            batch_dim = n
            total_n = n * num_per_sample

        _eps_t = torch.rand(total_n, device=self.device)
        if self.antithetic_sampling:
            offset = torch.arange(total_n, device=self.device) / total_n
            _eps_t = (_eps_t / total_n + offset) % 1
            perm = torch.randperm(total_n, device=self.device)
            _eps_t = _eps_t[perm]

        t_all = (t_max - t_min) * _eps_t + t_min

        if accum_step is not None:
            t_all = t_all.view(-1, num_per_sample)
            t_all = t_all.chunk(self.trainer.num_nodes)[self.trainer.node_rank]
            t_all = t_all.chunk(self.trainer.num_devices)[self.trainer.local_rank]
            t_all = t_all.chunk(self.trainer.accumulate_grad_batches)[accum_step]
            t_all = t_all[:batch_dim]
        else:
            t_all = t_all.view(batch_dim, num_per_sample)

        t_sorted, _ = torch.sort(t_all, dim=-1)
        return [t_sorted[:, i] for i in range(num_per_sample)]

    def _get_split_indices(self, n, accum_step, ratios):
        """Split batch indices into category sets"""
        assert sum(ratios) < 0.999, "Sum of ratios must be less than 1.0, leave one ratio out"

        batch_dim = n
        if accum_step is not None:
            n_total = self.config.loader.global_batch_size
            num_nodes = self.trainer.num_nodes
            num_devices = self.trainer.num_devices
            accum_batches = self.trainer.accumulate_grad_batches
            node_rank = self.trainer.node_rank
            local_rank = self.trainer.local_rank
            current_accum = accum_step
        else:
            n_total = n
            num_nodes = num_devices = accum_batches = 1
            node_rank = local_rank = current_accum = 0

        total_chunks = num_nodes * num_devices * accum_batches
        chunk_idx = (node_rank * num_devices * accum_batches
                     + local_rank * accum_batches
                     + current_accum)

        num_categories = len(ratios) + 1
        global_counts = [0] * num_categories
        net_ratio, prev_num = 0, 0
        for i, ratio in enumerate(ratios):
            net_ratio += ratio
            num_a = int(n_total * net_ratio)
            global_counts[i] = num_a - prev_num
            prev_num = num_a
        global_counts[-1] = n_total - prev_num

        local_counts = []
        for cnt in global_counts:
            base, rem = divmod(cnt, total_chunks)
            local_counts.append(base + (1 if chunk_idx < rem else 0))

        local_size = sum(local_counts)
        local_assignments = torch.empty(local_size, device=self.device, dtype=torch.long)
        offset = 0
        for i, cnt in enumerate(local_counts):
            local_assignments[offset:offset + cnt] = i
            offset += cnt

        seed = self.global_step * total_chunks + chunk_idx
        g = torch.Generator(device=self.device)
        g.manual_seed(seed)
        perm = torch.randperm(local_size, device=self.device, generator=g)
        local_assignments = local_assignments[perm][:batch_dim]

        return [(local_assignments == i).nonzero(as_tuple=True)[0]
                for i in range(num_categories)]

    def loss(self, x1, output_tokens,
             current_accumulation_step=None, train_mode=False, xT=None,
             given_t=None, not_sampling_t=False):
        del given_t, not_sampling_t
        del output_tokens, train_mode, xT
        B, L = x1.shape[0], x1.shape[1]

        tau_diag = self._sample_t_interval(B, current_accumulation_step,
                                           t_min=self.t_min, t_max=self.t_max)
        set_midpoint = getattr(self.config.algo, 'set_midpoint', 'midpoint')
        if self.config.algo.offdiagonal_sampling == "uniform_st":
            tau_s_offdiag, tau_t_offdiag = self._sample_multi_t_interval(
                B, current_accumulation_step, 2,
                t_min=self.t_min, t_max=self.t_max)
        else:  # uniform_diff
            tau_d_offdiag = self._sample_t_interval(
                B, current_accumulation_step, t_min=self.t_min, t_max=self.t_max)
            tau_s_offdiag = self._sample_t_interval(
                B, current_accumulation_step, t_min=self.t_min, t_max=self.t_max)
            tau_s_offdiag = tau_s_offdiag * (1 - tau_d_offdiag)
            tau_t_offdiag = tau_s_offdiag + tau_d_offdiag

        idx_diag, idx_offdiag_bndry, idx_offdiag = self._get_split_indices(
            B, current_accumulation_step,
            ratios=(self.config.algo.diagonal_fraction,
                    (1.0 - self.config.algo.diagonal_fraction)
                    * (1.0 / self.config.algo.boundary_prob)))
        tau_s = torch.zeros(B, device=self.device)
        tau_t = torch.zeros(B, device=self.device)
        tau_s[idx_diag] = tau_diag[idx_diag]
        tau_s[idx_offdiag_bndry] = 0.0
        tau_s[idx_offdiag] = tau_s_offdiag[idx_offdiag]
        tau_t[idx_diag] = tau_diag[idx_diag]
        tau_t[idx_offdiag_bndry] = 1.0
        tau_t[idx_offdiag] = tau_t_offdiag[idx_offdiag]

        tau_s = torch.clamp(tau_s, 0.0, 1.0)
        tau_t = torch.clamp(tau_t, 0.0, 1.0)

        idx_offdiag = torch.cat([idx_offdiag, idx_offdiag_bndry])
        has_diag = idx_diag.numel() > 0
        has_offdiag = idx_offdiag.numel() > 0

        if set_midpoint == 'midpoint':
            tau_u = 0.5 * (tau_s + tau_t)
        else:  # random
            tau_u = tau_s + torch.rand_like(tau_s) * (tau_t - tau_s)
        s = self._tau_to_t(tau_s)
        u = self._tau_to_t(tau_u)
        t = self._tau_to_t(tau_t)

        x_s, target_data = self.corrupt_continuous(x1, s)

        if self.teacher_model is None or not has_diag:
            on_diagonal_target = target_data[idx_diag]
        else:
            on_diagonal_target = self.teacher_forward(x_s[idx_diag], tau_s[idx_diag]).exp()
        on_diagonal_target = stopgrad(on_diagonal_target)

        if self.config.algo.distillation_method == "PSD": # Progressive Distillation. (semigroup)
            log_D_st = self.forward(x_s, tau_s, tau_t)
            _fwd = (self.forward_with_ema
                    if getattr(self.config.algo, 'use_ema_for_psd_target', False)
                    else self.forward)
            if has_offdiag:
                with torch.no_grad():
                    x_s_od = x_s[idx_offdiag]
                    s_od = s[idx_offdiag].view(-1, 1, 1)
                    u_od = u[idx_offdiag].view(-1, 1, 1)
                    t_od = t[idx_offdiag].view(-1, 1, 1)
                    tau_s_od = tau_s[idx_offdiag]
                    tau_u_od = tau_u[idx_offdiag]
                    tau_t_od = tau_t[idx_offdiag]
                    D_su_offdiag = _fwd(x_s_od, tau_s_od, tau_u_od).exp()
                    X_su = ((1 - u_od) / (1 - s_od + 1e-8)) * x_s_od \
                           + ((u_od - s_od) / (1 - s_od + 1e-8)) * D_su_offdiag
                    D_ut_offdiag = _fwd(X_su, tau_u_od, tau_t_od).exp()
                    lambda_sut = ((1 - t_od) * (u_od - s_od)
                                  / ((1 - u_od) * (t_od - s_od) + 1e-8))
                    offdiag_target = stopgrad(
                        lambda_sut * D_su_offdiag + (1 - lambda_sut) * D_ut_offdiag)

                if not self.config.algo.use_mse_loss_psd:
                    offdiag_loss = -(offdiag_target * log_D_st[idx_offdiag]).sum(dim=-1)
                else:
                    offdiag_loss = F.mse_loss(
                        log_D_st[idx_offdiag].exp(), offdiag_target,
                        reduction='none').sum(dim=-1)
            else:
                offdiag_loss = x_s.new_empty((0, L))

            if has_diag:
                if not self.config.algo.use_mse_loss_psd:
                    diag_loss = -(on_diagonal_target * log_D_st[idx_diag]).sum(dim=-1)
                else:
                    diag_loss = F.mse_loss(
                        log_D_st[idx_diag].exp(), on_diagonal_target,
                        reduction='none').sum(dim=-1)
            else:
                diag_loss = x_s.new_empty((0, L))

            if self.config.algo.rescale_offdiag_loss_psd is True and has_offdiag:
                offdiag_loss = offdiag_loss * (
                    (t_od - s_od) / (1 - s_od + 1e-8)).view(-1, 1).pow(2)

        elif self.config.algo.distillation_method == "ESD": # Eulerian Distillation
            use_jvp_attn = True
            
            if has_diag:
                log_D_st_diag = self.forward(
                    x_s[idx_diag], s[idx_diag], t[idx_diag], use_jvp_attn=False
                )
                diag_loss = -(on_diagonal_target * log_D_st_diag).sum(dim=-1)
            else:
                log_D_st_diag = x_s.new_empty((0, L, self.vocab_size))
                diag_loss = x_s.new_empty((0, L))
            
            if has_offdiag:
                x_s_od = x_s[idx_offdiag]
                s_od = s[idx_offdiag]
                t_od = t[idx_offdiag]
                tau_s_od = tau_s[idx_offdiag]
                tau_t_od = tau_t[idx_offdiag]
                
                with torch.no_grad():
                    use_teacher = (
                        self.teacher_model is not None
                        and getattr(self.config.algo, 'use_teacher_for_D_s_esd', True)
                    )
                    if use_teacher:
                        D_s = self.teacher_forward(x_s_od, tau_s_od).exp()
                    else:
                        D_s = self.forward(
                            x_s_od, tau_s_od, tau_t_od, use_jvp_attn=False
                        ).exp()
                
                d_tau_s_by_d_s = self._d_tau_by_d_t(
                    s_od.view(-1, 1, 1)
                ).squeeze()
                
                with torch.enable_grad():
                    def forward_s_x(tau_s_val, x_s_val):
                        return self.forward(
                            x_s_val, tau_s_val, tau_t_od,
                            use_jvp_attn=use_jvp_attn
                        )
                    
                    tangent_tau_s = d_tau_s_by_d_s * torch.ones_like(tau_s_od)
                    tangent_x_s = (D_s - x_s_od) / (1 - s_od.view(-1, 1, 1) + 1e-8)
                    
                    log_D_st_offdiag, d_ds_log_D_st = torch.func.jvp(
                        forward_s_x,
                        (tau_s_od, x_s_od),
                        (tangent_tau_s, tangent_x_s),
                    )
                    d_ds_log_D_st = stopgrad(d_ds_log_D_st)
                
                with torch.no_grad():
                    s_g = s_od.view(-1, 1, 1)
                    t_g = t_od.view(-1, 1, 1)
                    
                    D_st = log_D_st_offdiag.exp()
                    d_ds_D_st = D_st * d_ds_log_D_st
                    
                    coeff = (1 - s_g) * (t_g - s_g) / (1 - t_g + 1e-8)
                    offdiag_target = stopgrad(D_s + coeff * d_ds_D_st)
                
                offdiag_loss = F.mse_loss(log_D_st_offdiag.exp(), offdiag_target, reduction='none').sum(dim=-1)
            else:
                log_D_st_offdiag = x_s.new_empty((0, L, self.vocab_size))
                offdiag_loss = x_s.new_empty((0, L))
            
            log_D_st = torch.zeros(B, L, self.vocab_size, device=self.device)
            log_D_st[idx_offdiag] = log_D_st_offdiag
            log_D_st[idx_diag] = log_D_st_diag
            
        else: # Langrangian distillation
            use_jvp_attn = True

            if has_diag:
                log_D_st_diag = self.forward(
                    x_s[idx_diag], tau_s[idx_diag], tau_t[idx_diag], use_jvp_attn=False)
                diag_loss = -(on_diagonal_target * log_D_st_diag).sum(dim=-1)
            else:
                log_D_st_diag = x_s.new_empty((0, L, self.vocab_size))
                diag_loss = x_s.new_empty((0, L))

            if has_offdiag:
                with torch.no_grad():
                    x_s_od = x_s[idx_offdiag]
                    s_od = s[idx_offdiag].view(-1, 1, 1)
                    t_od = t[idx_offdiag].view(-1, 1, 1)
                    tau_s_od = tau_s[idx_offdiag]
                    tau_t_od = tau_t[idx_offdiag]

                with torch.enable_grad():
                    def forward_t(tau_t_val):
                        return self.forward(
                            x_s_od, tau_s_od, tau_t_val,
                            use_jvp_attn=use_jvp_attn)
                    tangent_t = torch.ones_like(tau_t_od)
                    log_D_st_offdiag, grad_tau_t_log_D_st_offdiag = torch.func.jvp(
                        forward_t, (tau_t_od,), (tangent_t,))
                    grad_tau_t_log_D_st_offdiag = stopgrad(grad_tau_t_log_D_st_offdiag)

                d_tau_by_d_t = self._d_tau_by_d_t(t_od)
                grad_t_log_D_st_offdiag = grad_tau_t_log_D_st_offdiag * d_tau_by_d_t

                with torch.no_grad():
                    D_st = log_D_st_offdiag.exp()
                    partial_t_D_st = D_st * grad_t_log_D_st_offdiag
                    X_st = ((1 - t_od) / (1 - s_od + 1e-8)) * x_s_od \
                           + ((t_od - s_od) / (1 - s_od + 1e-8)) * D_st
                    use_teacher_for_D_t = (
                        self.teacher_model is not None
                        and self.config.algo.use_teacher_for_D_t_lsd)
                    if use_teacher_for_D_t:
                        D_t__X_st = self.teacher_forward(X_st, tau_t_od).exp()
                    else:
                        D_t__X_st = self.forward(
                            X_st, tau_t_od, tau_t_od,
                            use_jvp_attn=False).exp()
                    offdiag_target = stopgrad(
                        D_t__X_st
                        - (t_od - s_od) * ((1 - t_od) / (1 - s_od + 1e-8))
                        * partial_t_D_st)
                    offdiag_target = entmax_bisect(
                        offdiag_target,
                        torch.tensor(
                            self.config.algo.entmax_temp_lsd,
                            dtype=torch.float32,
                            requires_grad=self.config.algo.backprop_entmax_temp_lsd,
                        ).to(self.device),
                        dim=-1,
                    )

                offdiag_loss = -(offdiag_target * log_D_st_offdiag).sum(dim=-1)
            else:
                log_D_st_offdiag = x_s.new_empty((0, L, self.vocab_size))
                offdiag_loss = x_s.new_empty((0, L))

            log_D_st = torch.zeros(B, L, self.vocab_size, device=self.device)
            log_D_st[idx_offdiag] = log_D_st_offdiag
            log_D_st[idx_diag] = log_D_st_diag

        loss = torch.zeros(B, L, device=self.device)
        
        if has_diag:
            loss[idx_diag] = diag_loss
            diag_loss_to_log = diag_loss.mean()
        else:
            diag_loss_to_log = loss.new_tensor(0.0)
        self.log('diag_loss', diag_loss_to_log, prog_bar=True, sync_dist=True)
        
        if has_offdiag:
            loss[idx_offdiag] = offdiag_loss
            offdiag_loss_to_log = offdiag_loss.mean()
        else:
            offdiag_loss_to_log = loss.new_tensor(0.0)
        self.log('offdiag_loss', offdiag_loss_to_log, prog_bar=True, sync_dist=True)
        self.log('loss', loss.mean(), prog_bar=True, sync_dist=True)

        if self.config.algo.learnable_loss_weighting is True:
            loss_weight = self.backbone.learnable_loss_weighting(tau_s, tau_t)
            loss_weight = loss_weight.unsqueeze(-1)
            loss = torch.exp(-loss_weight) * loss + loss_weight
            self.log('loss_weighted', loss.mean(), prog_bar=True, sync_dist=True)

        return loss

    @torch.no_grad()
    def generate_samples(self, num_samples, num_steps=None,
                         eps=1e-5):
        if num_steps is None:
            num_steps = self.config.sampling.steps
        gamma = getattr(self.config.sampling, 'gamma', 0.0)
        print(f"Sampling with {num_steps} steps")
        
        B = num_samples
        V = self.vocab_size
        L = self.num_tokens
        device = self.device

        tau_vals = torch.linspace(0.0, 1.0, num_steps + 1, device=device)


        z = torch.randn((num_samples, L, V), device=device, dtype=self.dtype)

        for i in range(num_steps):
            tau_curr = tau_vals[i]
            tau_next = tau_vals[i + 1]

            t_curr = self._tau_to_t(tau_curr.expand(B))

            t_next = self._tau_to_t(tau_next.expand(B))
            sigma_target = 1.0 - t_next

            sigma_tilde = sigma_target * torch.sqrt(torch.tensor(1.0 - gamma**2))
            t_tilde = 1.0 - sigma_tilde
            tau_tilde = self._t_to_tau(t_tilde)

            log_D_st_pred = self.forward(z, tau_curr.expand(B), tau_tilde)
            D_st_pred = log_D_st_pred.exp()

            if i == num_steps - 1:
                z = D_st_pred
                break

            weight_z = (1.0 - t_tilde.view(-1, 1, 1)) / (1.0 - t_curr.view(-1, 1, 1))
            weight_D = ((t_tilde.view(-1, 1, 1) - t_curr.view(-1, 1, 1))
                        / (1.0 - t_curr.view(-1, 1, 1)))
            z_tilde = weight_z * z + weight_D * D_st_pred

            if gamma > 0:
                noise_std = gamma * sigma_target.view(-1, 1, 1)
                mean_adjustment = sigma_tilde.view(-1, 1, 1) - sigma_target.view(-1, 1, 1)
                z = z_tilde + mean_adjustment * D_st_pred + noise_std * torch.randn_like(z)
            else:
                z = z_tilde
                
        return z.argmax(dim=-1)


class FMLM_TwoModel(FLMBase):
    """FMLM two-model parameterization (appendix: semigroup loss, first stage of two-stage MSE distillation)."""

    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self.log_flag = False
        self.teacher_model = None
        self._is_resuming = (
            config.checkpointing.resume_from_ckpt
            and config.checkpointing.resume_ckpt_path is not None
        )

    def setup(self, stage: str):
        if self.teacher_model is None:
            self._initialize_teacher()
        if stage == 'fit' and not self._is_resuming:
            print(">>> Initializing student from teacher...")
            self._initialize_student_from_teacher()
        elif self._is_resuming:
            print(">>> Skipping student initialization (resuming from checkpoint).")

    def _initialize_teacher(self):
        self.teacher_model = self._load_teacher_model(
            self.config.algo.teacher_path, use_plain_config=True)

    def _initialize_student_from_teacher(self):
        self._copy_teacher_weights_to_student(self.teacher_model.state_dict())
        if hasattr(self.backbone, 'output_layer'):
            self._zero_init_module(self.backbone.output_layer)
            print("Zero initialized student output_layer")

    def _on_load_checkpoint_extra(self, checkpoint):
        if self.config.mode == 'sample_eval':
            self._initialize_teacher()

    def on_train_start(self):
        super().on_train_start()
        if self.teacher_model is None:
            print("Initializing teacher model...")
            self._initialize_teacher()
            print("Initializing student from teacher")
            self._initialize_student_from_teacher()

    def teacher_forward(self, xt, tau):
        tau = self._process_sigma(tau)
        with torch.no_grad():
            with torch.amp.autocast(device_type=self.device.type, dtype=torch.float32):
                model_output = self.teacher_model(xt, tau)
        return self._process_model_output(model_output, xt, tau)

    def loss(self, x0, output_tokens,
             current_accumulation_step=None, train_mode=False,
             xT=None, given_t=None, not_sampling_t=False):
        del given_t, not_sampling_t, output_tokens
        B = x0.shape[0]

        d_tau = self._sample_t_interval(
            B, current_accumulation_step, t_min=0.0, t_max=1.0).clamp(min=1e-6)
        tau_s = torch.rand(B, device=self.device) * (1.0 - d_tau)
        if self.config.algo.add_boundary:
            p_boundary = 1.0 / self.config.algo.boundary_prob
            is_boundary = torch.rand(B, device=self.device) < p_boundary
            tau_s = torch.where(is_boundary, torch.tensor(0.0, device=self.device), tau_s)
            d_tau = torch.where(is_boundary, torch.tensor(1.0, device=self.device), d_tau)

        t_s = self._tau_to_t(tau_s)
        x_t_s, _ = self.corrupt_continuous(x0, t_s)
        f_final_sc = self.forward_no_softmax(x_t_s, tau_s, tau_s + d_tau)

        with torch.no_grad():
            d_tau_half = d_tau / 2.0
            t_s = t_s.view(-1, 1, 1)

            t_mid = self._tau_to_t(tau_s + d_tau_half).view(-1, 1, 1)
            t_e = self._tau_to_t(tau_s + d_tau).view(-1, 1, 1)
            dt_half_1 = t_mid - t_s
            dt_half_2 = t_e - t_mid
            dt = t_e - t_s

            f_theta_s = self.teacher_forward(x_t_s, tau_s).exp()
            v_s_hat = (f_theta_s - x_t_s) / (1.0 - t_s + 1e-5)
            g_theta_s_u = self.forward_no_softmax(x_t_s, tau_s, tau_s + d_tau_half)

            # v_su_hat = v_s(x_s) + 1/2(u-s)g_theta(x_s,u)
            v_s_u_hat = ((f_theta_s - x_t_s) / (1.0 - t_s + 1e-5)
                         + dt_half_1 / 2.0 * g_theta_s_u)

            # F_s,u(x_s) = x_s + (u-s)v_s(x_s) + 1/2(u-s)^2 g_theta(x_s,u)
            large_f_s_u = (x_t_s + dt_half_1 * v_s_hat
                           + 0.5 * dt_half_1 ** 2 * g_theta_s_u)

            f_theta_u = self.teacher_forward(large_f_s_u, tau_s + d_tau_half).exp()
            v_u_hat = (f_theta_u - large_f_s_u) / (1.0 - t_mid + 1e-5)
            g_theta_u_t = self.forward_no_softmax(
                large_f_s_u, tau_s + d_tau_half, tau_s + d_tau)

            # v_u,t_hat(x_u') = v_u(x_u') + 1/2*(t-u)*g_theta(x_u',u, t)
            v_u_t_hat = v_u_hat + dt_half_2 / 2.0 * g_theta_u_t

            v_hat = (dt_half_1 * v_s_u_hat + dt_half_2 * v_u_t_hat) / dt

            # x_1_hat = stopgrad(x_s + (1-s)*v_hat)
            x_boot = x_t_s + (1.0 - t_s) * v_hat
            x_boot = x_boot.detach()

        f_final_fm = f_theta_s
        weight = 0.5 * dt * (1.0 - t_s)
        f_final = f_final_fm + weight * f_final_sc
        error = x_boot - f_final  # (B, L, V)
        loss = (error ** 2).mean(dim=-1) * self.vocab_size  # (B, L)

        if self.config.algo.learnable_loss_weighting is True:
            loss_weight = self.backbone.learnable_loss_weighting(tau_s, tau_s + d_tau)
            loss_weight = loss_weight.unsqueeze(-1)
            loss = torch.exp(-loss_weight) * loss + loss_weight

        return loss

    @torch.no_grad()
    def generate_samples(self, num_samples, num_steps=None, eps=1e-5):
        """Generate samples using flow map."""
        if num_steps is None:
            num_steps = self.config.sampling.steps
        B = num_samples
        V = self.vocab_size
        L = self.num_tokens
        device = self.device

        z = torch.randn((num_samples, L, V), device=device, dtype=self.dtype)
        tau_vals = torch.linspace(0.0, 1.0, num_steps + 1, device=device)

        for i in range(num_steps):
            tau_curr = tau_vals[i]
            tau_next = tau_vals[i + 1]
            tau_in = tau_curr.expand(B)
            t_in = self._tau_to_t(tau_in)
            dt_in = self._tau_to_t(tau_next.expand(B)) - t_in

            x_1_pred_fm = self.teacher_forward(z, tau_in).exp()
            x_1_pred_sc = self.forward_no_softmax(z, tau_in, tau_next.expand(B))
            v_pred = (x_1_pred_fm - z) / (1.0 - t_in.view(-1, 1, 1) + eps)
            z = (z + v_pred * dt_in.view(-1, 1, 1)
                 + 0.5 * (dt_in.view(-1, 1, 1) ** 2) * x_1_pred_sc)

        return z.argmax(dim=-1)


class FMLM_TwoStage(FLMBase):
    """FMLM two-stage distillation (appendix: second stage compresses two-model teacher into single model)."""

    def __init__(self, config, tokenizer):
        super().__init__(config, tokenizer)
        self.teacher_model_f = None
        self.teacher_model_g = None
        self._is_resuming = (
            config.checkpointing.resume_from_ckpt
            and config.checkpointing.resume_ckpt_path is not None
            and utils.fsspec_exists(config.checkpointing.resume_ckpt_path)
        )

    def setup(self, stage: str):
        if self.teacher_model_f is None or self.teacher_model_g is None:
            self._initialize_teacher_f()
            self._initialize_teacher_g()
        if stage == 'fit' and not self._is_resuming:
            print(">>> Initializing student from teacher...")
            self._initialize_student_from_teacher()
        elif self._is_resuming:
            print(">>> Skipping student initialization (resuming from checkpoint).")

    def _initialize_teacher_f(self):
        self.teacher_model_f = self._load_teacher_model(
            self.config.algo.teacher_f_path, use_plain_config=True)

    def _initialize_teacher_g(self):
        """Load the residual teacher model (uses current config, no plain override)."""
        self.teacher_model_g = self._load_teacher_model(
            self.config.algo.teacher_g_path, use_plain_config=False)

    def _initialize_student_from_teacher(self):
        self._copy_teacher_weights_to_student(self.teacher_model_f.state_dict())


    def on_train_start(self):
        super().on_train_start()
        if self.teacher_model_f is None or self.teacher_model_g is None:
            print("Initializing teacher models...")
            self._initialize_teacher_f()
            self._initialize_teacher_g()
            print("Initializing student from teacher")
            self._initialize_student_from_teacher()


    def teacher_f_forward(self, xt, tau=None, d=None, use_jvp_attn=False):
        del d, use_jvp_attn
        sigma = tau.unsqueeze(-1) if tau.ndim == 1 else tau
        sigma = self._process_sigma(sigma)
        with torch.no_grad():
            with torch.amp.autocast(device_type=self.device.type, dtype=torch.float32):
                model_output = self.teacher_model_f(xt, sigma)
        return self._process_model_output(model_output, xt, sigma)

    def teacher_g_forward(self, xt, tau, tau_prime=None, **kwargs):
        sigma = self._process_sigma(tau)
        if tau_prime is not None:
            sigma_prime = self._process_sigma(tau_prime)
        else:
            sigma_prime = None
        with torch.amp.autocast(device_type=self.device.type, dtype=torch.float32):
            model_output = self.teacher_model_g(
                xt, sigma, sigma_prime, **kwargs)
        return model_output

    def loss(self, x1, output_tokens,
             current_accumulation_step=None, train_mode=False,
             xT=None, given_t=None, not_sampling_t=False):
        del given_t, not_sampling_t, output_tokens
        B = x1.shape[0]

        d_tau = self._sample_t_interval(
            B, current_accumulation_step,
            t_min=0.0, t_max=1.0).clamp(min=1e-5, max=1.0)
        tau_s = torch.rand(B, device=self.device) * (1.0 - d_tau)

        if self.config.algo.add_boundary:
            p_boundary = 1.0 / self.config.algo.boundary_prob
            is_boundary = torch.rand(B, device=self.device) < p_boundary
            tau_s = torch.where(is_boundary, torch.tensor(0.0, device=self.device), tau_s)
            d_tau = torch.where(is_boundary, torch.tensor(1.0, device=self.device), d_tau)

        t_s = self._tau_to_t(tau_s)
        x_t_s, _ = self.corrupt_continuous(x1, t_s)
        dt = self._tau_to_t(tau_s + d_tau) - t_s

        f_final_f = self.teacher_f_forward(x_t_s, tau_s).exp()
        v_f = (f_final_f - x_t_s) / (1.0 - t_s.view(-1, 1, 1) + 1e-5)
        f_final_g = self.teacher_g_forward(x_t_s, tau_s, tau_s + d_tau)
        F_s_t = (x_t_s + v_f * dt.view(-1, 1, 1)
                 + 0.5 * (dt.view(-1, 1, 1) ** 2) * f_final_g)

        student_pred = self.forward(x_t_s, tau_s, tau_s + d_tau).exp()
        student_v = (student_pred - x_t_s) / (1.0 - t_s.view(-1, 1, 1) + 1e-5)
        F_s_t_distilled = x_t_s + dt.view(-1, 1, 1) * student_v

        error = F_s_t - F_s_t_distilled
        loss = (error ** 2).mean(dim=-1) * self.vocab_size

        return loss

    @torch.no_grad()
    def generate_samples(self, num_samples, num_steps=None,
                        eps=1e-5):
        if num_steps is None:
            num_steps = self.config.sampling.steps
        gamma = getattr(self.config.sampling, 'gamma', 0.0)
        print(f"Sampling with {num_steps} steps")
        B = num_samples
        V = self.vocab_size
        L = self.num_tokens
        device = self.device

        tau_vals = torch.linspace(0.0, 1.0, num_steps + 1, device=device)  
        z = torch.randn((num_samples, L, V), device=device, dtype=self.dtype)
        
        for i in range(num_steps):
            tau_curr = tau_vals[i]
            tau_next = tau_vals[i + 1]

            t_curr = self._tau_to_t(tau_curr.expand(B))

            t_next = self._tau_to_t(tau_next.expand(B))
            sigma_target = 1.0 - t_next

            sigma_tilde = sigma_target * torch.sqrt(torch.tensor(1.0 - gamma**2))
            t_tilde = 1.0 - sigma_tilde
            tau_tilde = self._t_to_tau(t_tilde)

            log_D_st_pred = self.forward(z, tau_curr.expand(B), tau_tilde)
            D_st_pred = log_D_st_pred.exp()

            if i == num_steps - 1:
                z = D_st_pred
                break

            weight_z = (1.0 - t_tilde.view(-1, 1, 1)) / (1.0 - t_curr.view(-1, 1, 1))
            weight_D = ((t_tilde.view(-1, 1, 1) - t_curr.view(-1, 1, 1))
                        / (1.0 - t_curr.view(-1, 1, 1)))
            z_tilde = weight_z * z + weight_D * D_st_pred

            if gamma > 0:
                noise_std = gamma * sigma_target.view(-1, 1, 1)
                mean_adjustment = sigma_tilde.view(-1, 1, 1) - sigma_target.view(-1, 1, 1)
                z = z_tilde + mean_adjustment * D_st_pred + noise_std * torch.randn_like(z)
            else:
                z = z_tilde
                
        return z.argmax(dim=-1)