Dekodowanie z ograniczeniami: gramatyki i schemat JSON, żeby wywołania narzędzi zawsze się parsowały

Najczęstszy błąd agentów, jaki debuguję, to nie zła odpowiedź. To dobra odpowiedź, której harness nie umie sparsować: brakujący nawias klamrowy, doklejony na końcu komentarz, nazwa pola, której schemat nie zna. Grzeczna prośba do modelu sprawia, że zdarza się to rzadziej. Dekodowanie z ograniczeniami sprawia, że nie może się zdarzyć.

Mechanizm: maskowanie logitów

Na każdym kroku silnik wie, które tokeny zachowają poprawność wyjścia według gramatyki. Logity wszystkich pozostałych tokenów ustawia przed próbkowaniem na −∞. Model nadal wybiera, ale tylko spośród tokenów, które mogą prowadzić do poprawnego wyjścia. Próbkowanie, temperatura i wszystko inne z notatki o próbkowaniu działają na zamaskowanym rozkładzie.

Poprawne wyjście jest wtedy gwarantowane przez konstrukcję, a nie przez nadzieję.

Formaty regularne: automat skończony

Jeśli format jest regularny, jak data, enum albo rekord o stałej strukturze, kompiluje się do automatu skończonego. Silnik śledzi bieżący stan i dopuszcza tylko tokeny, których znaki prowadzą do poprawnego następnego stanu.

Outlines (Willard i Louf, 2023) przyspieszyło to, wyliczając z góry dla każdego stanu automatu, które tokeny ze słownika są dozwolone. Każdy krok staje się odczytem z tablicy zamiast przeglądania całego słownika.

Formaty zagnieżdżone: potrzebny jest stos

JSON nie jest regularny. Obiekty zagnieżdżają się w tablicach wewnątrz obiektów na dowolną głębokość, a dopasowanie nawiasów wymaga pamięci, której automat skończony nie ma. Prawdziwa obsługa JSON-a wymaga gramatyki bezkontekstowej, wykonywanej jako automat ze stosem: maszyna stanów plus stos.

llama.cpp używa GBNF, własnego formatu gramatyk (--grammar, --grammar-file). Potrafi też za Ciebie przekonwertować schemat JSON na gramatykę (--json-schema albo response_format na serwerze). Gramatyka jest egzekwowana token po tokenie w trakcie generowania.

Problem granic tokenów

Gramatyki są zapisane w znakach. Modele emitują tokeny, a jeden token może obejmować kilka symboli gramatyki: "},{" może być jednym tokenem, który zamyka łańcuch znaków oraz obiekt i otwiera kolejny. Silnik musi sprawdzać każdy kandydujący token znak po znaku względem automatu, od bieżącego stanu, co przy słowniku 150K tokenów jest kosztowne.

Współczesne implementacje dzielą tę pracę. Poprawność większości tokenów zależy tylko od bieżącego stanu automatu i można ją policzyć z góry. W trakcie działania sprawdzane są tylko tokeny, których poprawność zależy od zawartości stosu. XGrammar (Dong i in., 2024) jest zbudowany wokół tego podziału i to dzięki niemu dekodowanie ograniczone gramatyką może dziś działać niemal z prędkością dekodowania bez ograniczeń.

Ograniczenia mogą pogorszyć rozumowanie

Wymuszanie formatu od pierwszego tokenu ma swoją cenę. Tam i in. (2024) stwierdzili, że ścisłe ograniczenia formatu mierzalnie pogarszały rozumowanie w kilku zadaniach w porównaniu ze swobodną odpowiedzią modelu i późniejszą konwersją. Model zmuszony do wyemitowania {"answer":, zanim do czegokolwiek doszedł, nie ma gdzie myśleć.

Wzorzec, który sprawdza się u agentów:

1. Pozwól modelowi rozumować w wolnym tekście albo w osobnym bloku rozumowania.
2. Ograniczaj tylko wywołanie narzędzia: gramatyka albo schemat, które zaczynają obowiązywać, gdy model rozpoczyna wywołanie, już po zakończeniu rozumowania.
3. I tak waliduj względem prawdziwego schematu. Gramatyka gwarantuje składnię, a nie to, że "path" wskazuje istniejący plik.

Co to naprawia w praktyce

W moim własnym benchmarku agentów jednym z trzech błędów harnessu, które kosztowały 52,5 punktu, było wyciąganie gołego JSON-a: harness nie potrafił wyciągnąć odpowiedzi, gdy model odpowiadał gołym JSON-em (historia jest w notatce o amber-bench). Wywołania narzędzi generowane z ograniczeniami usuwają całą tę klasę błędów. Parser i generator zgadzają się co do formatu, bo format egzekwuje ta sama gramatyka, której używają obie strony.

Zasady, których się trzymam

1. Ograniczaj wywołania narzędzi, nie myślenie. Najpierw swobodne rozumowanie, a dla wywołania ścisła gramatyka.
2. Generuj gramatyki z tego samego schematu, względem którego walidują Twoje narzędzia, żeby istniało jedno źródło prawdy.
3. Trzymaj schematy płytkie i zamknięte. Bez additionalProperties, enumy wszędzie tam, gdzie zbiór jest znany. Każde ograniczenie, które niesie gramatyka, to błąd, którego model nie może popełnić.
4. Nadal obsługuj odmowę i przekroczenia czasu. Gramatyka nie zmusi modelu do skończenia; ustaw limit tokenów i rozsądną ścieżkę błędu.

Źródła: Willard i Louf, Efficient Guided Generation for Large Language Models (2023). Dong i in., XGrammar (2024). Tam i in., Let Me Speak Freely? A Study on the Impact of Format Restrictions on Performance of Large Language Models (2024). Dokumentacja GBNF i schematów JSON w llama.cpp.