Incorrect active parameter count for MoE models

[TimFelixBeyer]

Hi,

first of all thank you for the very interesting paper!

I believe there is a mistake in the code to calculate the FLOPs for Mixtral. Mixtral is a 8x7b MoE model with 2 active experts.
Therefore, in the forward pass, only 1/4 of the non-embedding parameters are actually used for computation, yet the model_size dict (which is used in all flop calculations) contains the full number of ~47B tokens:

llm-threat-model/aggregate_csv.py

Line 12 in 3de6c6a

"mixtral_8x7b": 46702792704,

This likely leads to inflated FLOP counts for results using the MoE model (e.g., PAIR, AutoDAN).

Best
Tim

[kotekjedi]

Dear Tim,

Sorry for the late reply, and thank you for pointing this out — you are absolutely right. The FLOP count for Mixtral should indeed be divided by 4; the value in our model_size dictionary was mistakenly set to the full parameter count.

That said, we found that this only affected the PAIR results — for AutoDAN, we had used the correct FLOP count in our plots. We also want to confirm that this correction does not affect any of the conclusions we draw in the paper.

We have fixed the issue in the camera-ready version, and all affected plots have been updated accordingly. You can find the latest version here: https://arxiv.org/abs/2410.16222. We will also update the codebase in the next (hopefully soon!) release.

Thanks again for catching this!

Best regards,
Alexander Panfilov