Extending BCEWithLogitsLoss to non-scalar outputs (multi label)

[Leon0402]

Hi,

currently BCEWithLogitsLoss supports only binary classification:

backpack/backpack/core/derivatives/bcewithlogitsloss.py

Line 78 in 8aad99d

if module.input0.shape[1] != 1:

I have the VQAv2 taks, which is multi label, but also has soft labels i.e. multiple correct labels. Thus, I think currently it does not work with this library, because of:

• Extending BCEWithLogitsLoss to non-binary labels #281
• And because only scalar outputs are supported

What would be necessary to support my setup here?

[f-dangel]

Hi, thanks for reaching out.

It will depend what extension you would like to add support to.
Happy to give more details if you let me know what quantity you would like to compute.

Felix

[Leon0402]

Hi @f-dangel,

I am using backpack as part of laplace, which you can see here: https://github.com/aleximmer/Laplace/blob/main/laplace/curvature/backpack.py. The extensions used are:

KFAC,
KFLR,
BatchGrad,
DiagGGNExact,
DiagGGNMC,
SumGradSquared,

with BatchGrad and DiagGGNExact the most important for now. But if we can support all, of course that would be great!

One unrelated question to this issue, but maybe you can quickly help here too. Laplace has the following line:
https://github.com/aleximmer/Laplace/blob/main/laplace/curvature/backpack.py#L65
E.g. dict like input we often see for Hugging Face models seems to not be supported. I just removed that line and used my Hugging Face Model with the existing BCEWithLogitsLoss extension, and it "worked". Worked in the sense that something would train, and the results look, at first glance, not completely horrible.
Could you comment if Backpack in deed does not support dict-like input? This is definitely necessary for me, and my assumption was that this hint might be incorrect, and the only remaining issue is that I am using the wrong distribution in BCEWithLogitsLoss.

Thanks!
Leon

[f-dangel]

Hi, thanks for your patience.

Let me try to break down your request and let's focus on the two most important use cases for you.

1. Supporting multi-outputs for BCE with BatchGrad
2. Supporting multi-outputs for BCE with DiagGGNExact

Regarding 1: BackPACK should already be able to compute per-datum gradients if you are using BCE loss. You can follow the example here in the docs.

Regarding 2: To get this to work, you have to modify the function that computes the symmetric factorization of the BCE loss Hessian (here). The general documentation for what this function does is here.