Reinforcement Learning optimiert Token‑Entmaskierung in Diffusions‑Sprachmodellen

Diffusions‑Sprachmodelle (dLLMs) haben in den letzten Monaten die Leistung ihrer autoregressiven Gegenstücke auf vielen Aufgaben erreicht und gleichzeitig das Potenzial, bei der Inferenz effizienter zu sein. Besonders vielversprechend ist die Variante „masked discrete diffusion“, bei der ein Puffer aus Masken‑Tokens schrittweise durch Tokens aus dem Vokabular ersetzt wird.

Die Effizienz kann gesteigert werden, wenn mehrere Tokens gleichzeitig entmaskiert werden. Zu viele Entmaskierungen auf einmal führen jedoch zu Qualitätsverlusten. Deshalb ist die Auswahl der Tokens, die in jedem Diffusionsschritt ersetzt werden, ein entscheidender Designfaktor.

Frühere Ansätze nutzten heuristische Strategien wie Confidence‑Thresholding, die sowohl die Qualität als auch die Token‑Durchsatzrate verbesserten. Diese Heuristiken erfordern jedoch manuelle Feinabstimmung und zeigen bei größeren Puffern eine Leistungsabnahme.

In der vorliegenden Arbeit wird stattdessen ein Reinforcement‑Learning‑Ansatz vorgeschlagen. Die Entmaskierung wird als Markov‑Entscheidungsprozess formuliert, wobei das dLLM die Umgebung darstellt. Eine kompakte Policy‑Architektur – ein einzelner Transformer‑Layer – wandelt die Token‑Konfidenzen des Modells in Entmaskierungsentscheidungen um.

Experimentelle Ergebnisse zeigen, dass die trainierten Policies die Leistung der besten heuristischen Verfahren erreichen, wenn sie mit semi‑autoregressiver Generierung kombiniert werden, und in der vollständigen Diffusions‑Umgebung sogar überlegen sind.