pCoMole/flow_matching/loss/editflows_loss.py

# Copyright (c) Meta Platforms, Inc.
# All rights reserved.

from typing import List, Optional, Sequence
import torch
import torch.nn.functional as F
from torch import Tensor
from torch.nn.modules.loss import _Loss

# If you kept the adapter around only for scheduler access you can drop it entirely,
# because we pass precomputed_weight from training.step. Keeping it is harmless though.
class EditFlowsLoss(_Loss):
    """
    Edit Flows loss (Eq. 23), ragged version:
      L_i =  (sum_j λ_ins[i][j] + sum_j λ_del[i][j] + sum_j λ_sub[i][j])
            - w_i * ( sum_{ins events e} [log λ_ins[i][slot_e] + log Q_ins[i](y_e)]
                    + sum_{del j}       [log λ_del[i][j]]
                    + sum_{sub j}       [log λ_sub[i][j] + log Q_sub[i](y_j)] )
    """

    def __init__(self, reduction: str = "mean") -> None:
        super().__init__(None, None, reduction)

    def forward(
        self,
        lam_ins: torch.Tensor,     # (B, L)
        logits_ins: torch.Tensor,  # (B, L, V)
        lam_del: torch.Tensor,     # (B, L)
        lam_sub: torch.Tensor,     # (B, L)
        logits_sub: torch.Tensor,  # (B, L, V)
        z_t: torch.Tensor,         # (B, N) aligned, with eps_id
        z_1: torch.Tensor,         # (B, N) aligned target, with eps_id
        x_t: torch.Tensor,         # (B, L)
        valid_mask: torch.Tensor,  # (B, L) bool, False==padding, True==valid
        precomputed_weight: torch.Tensor,  # (B,) or ()  = kappa_dot/(1-kappa)
        eps_id: int,
        bos_id: int,
        eos_id: int,
    ) -> torch.Tensor:
        """
        Implements Eq. 23 style loss for Edit Flows.

        We:
        1) penalize total outgoing rate on x_t
        2) for every column in (z_t, z_1) that still differs, we map it to an x_t edit
            and add  - w * log u_required

        BOS/EOS handling:
        - z_t and z_1 are aligned and already contain BOS/EOS
        - we do NOT allow edits that delete/replace BOS/EOS
        - we also skip columns whose target token is BOS/EOS (they should already match)
        """
        device = lam_ins.device
        B, L = x_t.shape
        _, N = z_t.shape

        # 1. ----- RATE TERM -----
        # valid_mask: True = real token, False = pad
        valid_f = valid_mask.to(lam_ins.dtype)  # (B, L) 1.0 for valid, 0.0 for pad

        # total outgoing rate at each position
        total_rate_pos = lam_ins + lam_del + lam_sub     # (B, L)
        total_rate_pos = total_rate_pos * valid_f        # zero out pads
        loss_rate = total_rate_pos.sum(dim=1)            # (B,)

        # 2. ----- EDIT TERM -----
        # precompute log-softmax for tokens (better numerics)
        logp_ins = F.log_softmax(logits_ins, dim=-1)   # (B, L, V)
        logp_sub = F.log_softmax(logits_sub, dim=-1)   # (B, L, V)

        # Make weight shape nice
        if precomputed_weight.dim() == 0:
            precomputed_weight = precomputed_weight.view(1).expand(B).to(device)
        else:
            precomputed_weight = precomputed_weight.to(device)

        # use float for accumulation
        loss_edit = torch.zeros(B, dtype=torch.float32, device=device)

        for b in range(B):
            # how many *valid* tokens in x_t[b]
            valid_len = int(valid_mask[b].sum().item())

            prefix_non_eps = 0  # = number of non-ε seen in z_t[b, :i]

            for i in range(N):
                zt = int(z_t[b, i].item())
                z1 = int(z_1[b, i].item())

                # map this column to an x_t position
                if zt != eps_id:
                    # this column corresponds to x_t[b, prefix_non_eps]
                    x_pos = prefix_non_eps
                    is_token = True
                    prefix_non_eps += 1
                else:
                    # this is a gap column, sits BETWEEN tokens
                    x_pos = prefix_non_eps
                    is_token = False

                # if already matched -> no term
                if zt == z1:
                    continue

                # --- BOS/EOS guards on the target side ---
                # if the target token is BOS or EOS, we shouldn't try to force an edit here
                if z1 == bos_id or z1 == eos_id:
                    # skip this column; aligned BOS/EOS should already match
                    continue

                # figure out which edit we need
                # CASE 1: deletion: z_t has token, z_1 has ε
                if is_token and (z1 == eps_id):
                    # delete token at x_pos
                    if x_pos >= valid_len:
                        # out of range (shouldn't happen if alignment & mask match)
                        raise NotImplementedError

                    x_token = int(x_t[b, x_pos].item())
                    # do NOT delete BOS/EOS
                    if x_token == bos_id or x_token == eos_id:
                        continue

                    lam = lam_del[b, x_pos].clamp_min(1e-12)
                    log_u_req = torch.log(lam)

                # CASE 2: substitution: token -> different token
                elif is_token and (z1 != eps_id) and (zt != z1):
                    if x_pos >= valid_len:
                        raise NotImplementedError

                    x_token = int(x_t[b, x_pos].item())
                    # do NOT substitute BOS/EOS
                    if x_token == bos_id or x_token == eos_id:
                        continue

                    lam = lam_sub[b, x_pos].clamp_min(1e-12)
                    logp_tok = logp_sub[b, x_pos, z1]
                    log_u_req = torch.log(lam) + logp_tok

                # CASE 3: insertion: ε -> token
                elif (not is_token) and (z1 != eps_id):
                    # insertion in the gap after token (x_pos - 1)
                    # if x_pos == 0, we insert after "BOS"/at start -> map to position 0
                    ins_pos = x_pos - 1
                    if ins_pos < 0:
                        ins_pos = 0

                    # clamp to last valid position if needed
                    if valid_len == 0:
                        # degenerate, but avoid -1
                        ins_pos = 0
                    elif ins_pos >= valid_len:
                        ins_pos = valid_len - 1

                    # also don't insert "after" EOS if ins_pos currently points to EOS
                    x_token = int(x_t[b, ins_pos].item())
                    if x_token == eos_id:
                        # simplest policy: skip this insertion supervision
                        continue

                    lam = lam_ins[b, ins_pos].clamp_min(1e-12)
                    logp_tok = logp_ins[b, ins_pos, z1]
                    log_u_req = torch.log(lam) + logp_tok

                else:
                    # unknown pattern (shouldn't happen with proper alignment)
                    raise NotImplementedError

                w = precomputed_weight[b]
                loss_edit[b] += - w * log_u_req

        # 3. ----- COMBINE -----
        loss = loss_rate + loss_edit  # (B,)
        loss = loss.mean()
        return loss


class EditFlowsLossReParam(_Loss):
    """
    Edit Flows loss (Eq. 23) but parameterized by:
      - lam_total: (B, L) >= 0
      - pi_type:   (B, L, 3) prob over {ins, del, sub}, sums to 1

    Token distributions are still given by logits_ins/logits_sub.

    Required-edit log rate:
      del: log u = log lam_total[pos] + log pi_del[pos]
      sub: log u = log lam_total[pos] + log pi_sub[pos] + log Q_sub(token)
      ins: log u = log lam_total[pos] + log pi_ins[pos] + log Q_ins(token)
    """

    def __init__(self, reduction: str = "mean", gamma_rate: float = 1.0, eps: float = 1e-12) -> None:
        super().__init__(None, None, reduction)
        self.gamma_rate = float(gamma_rate)
        self.eps = float(eps)

    def forward(
        self,
        lam_total: torch.Tensor,     # (B, L)
        pi_type: torch.Tensor,       # (B, L, 3) probabilities over {ins, del, sub}
        logits_ins: torch.Tensor,    # (B, L, V)
        logits_sub: torch.Tensor,    # (B, L, V)
        z_t: torch.Tensor,           # (B, N) aligned, with eps_id
        z_1: torch.Tensor,           # (B, N) aligned target, with eps_id
        x_t: torch.Tensor,           # (B, L)
        valid_mask: torch.Tensor,    # (B, L) bool, False==padding, True==valid   (same as your old loss)
        precomputed_weight: torch.Tensor,  # (B,) or ()  = kappa_dot/(1-kappa)
        eps_id: int,
        bos_id: int,
        eos_id: int,
    ) -> torch.Tensor:
        device = lam_total.device
        B, L = x_t.shape
        _, N = z_t.shape

        # --- RATE TERM ---
        valid_f = valid_mask.to(lam_total.dtype)                 # True=valid -> 1.0
        loss_rate = (lam_total * valid_f).sum(dim=1)             # (B,)
        loss_rate = self.gamma_rate * loss_rate

        # --- EDIT TERM ---
        logp_ins = F.log_softmax(logits_ins, dim=-1)             # (B, L, V)
        logp_sub = F.log_softmax(logits_sub, dim=-1)             # (B, L, V)

        # normalize pi_type defensively (in case caller passes slightly off-simplex)
        pi_sum = pi_type.sum(dim=-1, keepdim=True).clamp_min(self.eps)
        pi_type = pi_type / pi_sum

        log_lam = torch.log(lam_total.clamp_min(self.eps))       # (B, L)
        log_pi = torch.log(pi_type.clamp_min(self.eps))          # (B, L, 3)

        # weight shape
        if precomputed_weight.dim() == 0:
            precomputed_weight = precomputed_weight.view(1).expand(B).to(device)
        else:
            precomputed_weight = precomputed_weight.to(device)

        loss_edit = torch.zeros(B, dtype=torch.float32, device=device)

        for b in range(B):
            valid_len = int(valid_mask[b].sum().item())
            prefix_non_eps = 0

            for i in range(N):
                zt = int(z_t[b, i].item())
                z1 = int(z_1[b, i].item())

                # map this alignment column to an x_t position
                if zt != eps_id:
                    x_pos = prefix_non_eps
                    is_token = True
                    prefix_non_eps += 1
                else:
                    x_pos = prefix_non_eps
                    is_token = False

                # already matched -> no supervision
                if zt == z1:
                    continue

                # don't force edits where target is BOS/EOS
                if z1 == bos_id or z1 == eos_id:
                    continue

                # CASE 1: deletion (token -> eps)
                if is_token and (z1 == eps_id):
                    if x_pos >= valid_len:
                        raise NotImplementedError

                    x_token = int(x_t[b, x_pos].item())
                    if x_token == bos_id or x_token == eos_id:
                        continue

                    log_u_req = log_lam[b, x_pos] + log_pi[b, x_pos, 1]  # del index = 1

                # CASE 2: substitution (token -> different token)
                elif is_token and (z1 != eps_id) and (zt != z1):
                    if x_pos >= valid_len:
                        raise NotImplementedError

                    x_token = int(x_t[b, x_pos].item())
                    if x_token == bos_id or x_token == eos_id:
                        continue

                    log_u_req = (
                        log_lam[b, x_pos]
                        + log_pi[b, x_pos, 2]           # sub index = 2
                        + logp_sub[b, x_pos, z1]
                    )

                # CASE 3: insertion (eps -> token), map gap to "after token (x_pos-1)"
                elif (not is_token) and (z1 != eps_id):
                    ins_pos = x_pos - 1
                    if ins_pos < 0:
                        ins_pos = 0

                    if valid_len == 0:
                        ins_pos = 0
                    elif ins_pos >= valid_len:
                        ins_pos = valid_len - 1

                    x_token = int(x_t[b, ins_pos].item())
                    if x_token == eos_id:
                        continue

                    log_u_req = (
                        log_lam[b, ins_pos]
                        + log_pi[b, ins_pos, 0]         # ins index = 0
                        + logp_ins[b, ins_pos, z1]
                    )

                else:
                    raise NotImplementedError

                w = precomputed_weight[b]
                loss_edit[b] += -w * log_u_req

        loss = loss_rate + loss_edit
        return loss.mean()