One-step optimization method collapses when applied to math tasks (GSM8K and MATH)

Hi! Great work！
I'm having issues applying it to generative tasks like GSM8K and MATH - the model outputs completely collapse.
Could you help with:

1.  What modifications are needed for hard generative tasks? (Code examples would be helpful)
2. Do these tasks require the NGD method? If so, could you share the implementation?

Thanks for any guidance!