Question about the method

Awesome work !!! I wonder have you tried using simpler token selection module such as MLP followed by Gumbel-Softmax (like DynamicViT did), and design the FLOPs loss term ? I think the FLOPs loss will also propagate to the token selection module since Gumbel-Softmax is differentiable. Hopes to get your reply, thanks !!!