SLIM-RL: Efektywny training dyfuzyjnych LLM-ów bez rozbijania trajektorii

Naukowcy zaprezentowali SLIM-RL, metodę znacznie bardziej efektywną niż dotychczasowy standard TraceRL dla reinforcement learningu w dyfuzyjnych LLM-ach. Podczas gdy TraceRL rekonstruuje pełną trajektorię modelu podczas treningu poprzez rozbijanie każdego rollout-u na aż K/s próbek treningowych (gdzie K to rozmiar bloku), SLIM-RL osiąga podobne wyniki bez tego kosztownego procesu. Kluczową ideą jest kontrolowanie ryzyka za pomocą dekodera z tau-budżetem, który ogranicza commit risk każdego kroku i zmniejsza łączne ryzyko w danych treningowych.

Metoda SLIM-RL łączy kilka zaawansowanych technik: sequence-level importance sampling, deterministic quadrature w kontekście różnych poziomów maskingu oraz nowy schemat maski monotonnicznie malejący na blok. Podczas optymalizacji model trenuje się na rolloutach kontrolowanych pod względem ryzyka, wykorzystując random-masking objective bez potrzeby rekonstrukcji trajektorii. To podejście znacząco zmniejsza wymagania obliczeniowe.

Wyniki na modelu SDAR-4B pokazują, że SLIM-RL dorównuje najlepszym wynikom TraceRL-a dla MATH500 używając zaledwie 46 procent próbek treningowych przy rozmiarze bloku 16. Przy rozmiarze bloku 4, czteromiardowy SLIM-RL przewyższa większe modele dyfuzyjne LLaDA-8B i Dream-7B, osiągając o 10,76 procent wyższy wynik na MATH500 od LLaDA-8B, pozostając wciąż poniżej autoregresyjnego Qwen2.5-7B. Na zadaniach kodowania metoda poprawia wyniki o 4,20 procent na MBPP i 3,65 procent na HumanEval w porównaniu z TraceRL.