gelatwo-xlam-function-monarch-16384

A hidden Markov model distilled from meta-llama/Llama-3.1-8B-Instruct, used as the tractable proposal for constrained tool calling on xLAM (function-call template) in Mitigating Bias in Locally Constrained Decoding via Tractable Proposals (arXiv:2606.01926).

Locally constrained decoding (LCD) enforces a constraint by masking next tokens one step at a time, which biases the samples it draws. P-GCD instead combines this HMM with a tensorized finite automaton to form a proposal that carries both the logical and the probabilistic information about the target distribution, and samples from it with sequential Monte Carlo. The HMM supplies the lookahead the automaton alone cannot: how likely the constraint is to be satisfiable from here.

Architecture

Transition	Monarch-factorized, 128x128 blocks
Hidden states	16384
Vocabulary	128256 (matches meta-llama/Llama-3.1-8B-Instruct)
Class	gelatwo.hmm.MonarchHMM

This is not a fine-tune of the base LM and shares none of its weights: it is a separate model, trained by EM on sequences sampled from that LM, and it reuses the LM's tokenizer so their vocabularies line up at decoding time.

Usage

Install the reference implementation, github.com/MhDang/gelatwo, then:

from gelatwo.hmm import MonarchHMM

hmm = MonarchHMM.from_pretrained("mhdang/gelatwo-xlam-function-monarch-16384")

configs/xlam.yaml in the reference implementation already points at this checkpoint, so the xLAM (function template) runs need no override.

Citation

@inproceedings{dang2026mitigating,
  title = {Mitigating Bias in Locally Constrained Decoding via Tractable Proposals},
  author = {Dang, Meihua and Song, Linxin and Zhang, Honghua and Zhao, Jieyu and Van den Broeck, Guy and Ermon, Stefano},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year = {2026},
}