Using CCE for image classification as loss

[MichaelMMeskhi]

Out of curiosity, would CCE be easily interchanged with CE for instance in image classification problems?

[erikwijmans]

Should be. We developed CCE in the context of LLMs, but the only requirement of CCE is that there is a final affine transformation that goes from an embedding space to logits over the label (vocab) space. So as long as your loss computation is something like

logits = self.classifier(embeddings)
loss = torch.nn.functional.cross_entropy(logits, labels)

where self.classifier = nn.Linear(...)

you can instead do

from cut_cross_entropy import linear_cross_entropy
loss = linear_cross_entropy(embeddings, self.classifier.weight, labels, bias=self.classifier.bias)

[nleroy917]

Hello @erikwijmans! I might hijack this.

Thank you for this code, I'm very excited about the prospect of CCE; I have a data modality with extraordinarily large vocab sizes.

I tried implementing what you suggested above, and while it works (it runs), the loss appears to stay constant (right at 2). Im curious if you have any off-rip thoughts about this?

Here is exactly what I am running:

def compute_cce_loss(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    """
    Compute the loss for the batch using the LigerFusedCrossEntropyLoss.

    Inspired by: https://github.com/apple/ml-cross-entropy/issues/22

    :param embeddings: The embeddings from the model of shape (batch_size, seq_length, d_model).
    :param labels: The labels tensor of shape (batch_size, seq_length).

    :return: The computed loss.
    """
    labels = labels.view(-1)
    embeddings = embeddings.view(-1, self.r2v_model.d_model)

    loss = linear_cross_entropy(
        embeddings, self.linear.weight, labels, bias=self.linear.bias, ignore_index=-100
    )

    return loss

FWIW, I am running with lightning and using wandb to track experiments. This is a single training step:

def training_step(
    self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int
) -> torch.Tensor:
    """
    Perform a training step.

    :param batch: The batch containing (tokens, masked_tokens, labels, attention_mask).
    :param batch_idx: The batch index.
    :return: The computed loss.
    """
    _, masked_tokens, labels, attention_mask = batch
    masked_token_positions = labels != -100
    # forward pass
    embeddings = self.forward(
        masked_tokens, mask=attention_mask, masked_token_positions=masked_token_positions
    )
    # compute loss
    # loss = self.compute_loss(embeddings, labels)
    # loss = self.compute_fused_loss(embeddings, labels)
    loss = self.compute_cce_loss(embeddings, labels)

    # log loss
    self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

    return loss

[nleroy917]

Coming in to update with my configuration, and mention that the loss seems to be working now! I switched to LinearCrossEntropy and/or module form, instead of the function form. Not sure what I was doing wrong, but I am getting more conventional loss curves.

class MLMAdapter(L.LightningModule):

    def __init__(
        self, model: AtacformerExModel, loss_type: MLM_CE_LOSS_TYPES = "vanilla", **kwargs
    ):
        """
        Instantiate the masked-language-modeling adapter.

        :param Atacformer model: The model to train.
        :param MLM_CE_LOSS_TYPES loss_type: The type of loss to use for training. (default: "vanilla")
        :param kwargs: Additional arguments to pass to the LightningModule constructor.
        """
        super().__init__(**kwargs)
        self.r2v_model = model._model
        self.loss_type = loss_type

        # linear layer acts as a classification layer for training
        # the model on the masked language modeling task
        self.linear = nn.Linear(model._model.d_model, model._model.vocab_size)

        self.tokenizer = model.tokenizer
        self.init_lr = kwargs.get("init_lr", 1e-5)

        # tie the weights of the linear layer with the token embeddings
        # this is done primarily to reduce the number of parameters
        # during training significantly
        self.linear.weight = self.r2v_model.embedding.weight

        # verify vocab_size consistency
        assert self.r2v_model.vocab_size == len(
            self.tokenizer
        ), f"Model vocab_size {self.r2v_model.vocab_size} does not match tokenizer vocab_size {len(self.tokenizer)}"

        self.masked_token_id = self.tokenizer.mask_token_id()
        self.passing_token_id = self.tokenizer.padding_token_id()

        self.ce_fn = nn.CrossEntropyLoss()
        self.cce_fn = self._init_cce_loss()

        # The dictionary that picks the correct “worker”
        self.loss_fns = {
            "vanilla": self.compute_vanilla_loss,
            "cce": self.compute_cce_loss,
        }

    def _init_cce_loss(self):
        try:
            from cut_cross_entropy import LinearCrossEntropy

            return LinearCrossEntropy(ignore_index=-100)
        except ImportError:
            return None

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor = None,
        masked_token_positions: torch.Tensor = None,
    ) -> torch.Tensor:
        token_embeddings = self.r2v_model(x, mask=mask)
        return token_embeddings

    def on_fit_start(self):
        # move positional embeddings to the device
        self.r2v_model.positional_encoding = self.r2v_model.positional_encoding.to(self.device)

    def compute_cce_loss(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """
        Compute the loss for the batch using the LigerFusedCrossEntropyLoss.

        Inspired by: https://github.com/apple/ml-cross-entropy/issues/22

        :param embeddings: The embeddings from the model of shape (batch_size, seq_length, d_model).
        :param labels: The labels tensor of shape (batch_size, seq_length).

        :return: The computed loss.
        """
        if self.cce_fn is None:
            raise RuntimeError("LinearCrossEntropy is not available. Something went wrong.")

        labels = labels.view(-1)
        embeddings = embeddings.view(-1, self.r2v_model.d_model)

        # compute loss using the linear cross entropy implementation
        # this is done to drastically reduce the memory footprint
        # inspiration: https://github.com/apple/ml-cross-entropy/issues/22
        loss = self.cce_fn(embeddings, self.linear.weight, labels, bias=self.linear.bias)

        return loss

    def compute_loss(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """
        Compute the loss for the batch.

        :param embeddings: The embeddings from the model of shape (batch_size, seq_length, d_model).
        :param labels: The labels tensor of shape (batch_size, seq_length).

        :return: The computed loss.
        """
        # actual single "public" interface
        # that dispatches (no if/else needed).
        return self.loss_fns[self.loss_type](embeddings, labels)

    def training_step(
        self, batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int
    ) -> torch.Tensor:
        """
        Perform a training step.

        :param batch: The batch containing (tokens, masked_tokens, labels, attention_mask).
        :param batch_idx: The batch index.
        :return: The computed loss.
        """
        _, masked_tokens, labels, attention_mask = batch
        masked_token_positions = labels != -100
        # forward pass
        embeddings = self.forward(
            masked_tokens, mask=attention_mask, masked_token_positions=masked_token_positions
        )
        # compute loss
        loss = self.compute_loss(embeddings, labels)

        # log loss
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

        return loss