[ENH] Add optional margin for CDW-CE

[conradbkay]

Describe the feature or idea you want to propose

Paper by some of the original CDW-CE authors:
https://arxiv.org/abs/2412.01246

The formula doesn't make sense to me assuming log(1 − max(1, yˆi + m)) is always undefined

Describe your proposed solution

I'm guessing they mean something like

if self.margin is not None:
    probs = F.softmax(y_pred, dim=1)
    probs = torch.clamp(probs + self.margin, max=1 - 1e-8)
    log_term = torch.log(1 - probs)
    
    loss = log_term * weights

Describe alternatives you have considered, if relevant

I couldn't find any code implementations for the newer paper

Additional context

No response

[victormvy]

Hi Conrad, thanks for pointing this out, it makes sense.
The goal here is to implement methods as they are described in the papers, so it would be interesting if the original authors could clarify this point or release source code to confirm the intended formulation.