Flow of Spans: Generalizing Language Models to Dynamic Span-Vocabulary via GFlowNets

AI summary

Authors: Bo Xue, Yunchong Song, Fanghao Shao, Xuekai Zhu, Lin Chen, Luoyi Fu, Xinbing Wang, Zhouhan Lin (LUMIA Lab, Shanghai Jiao Tong University; Shanghai Artificial Intelligence Laboratory; Shanghai Innovation Institute) arXiv: 2602.10583v1 [cs.AI], 11 Feb 2026 · Code: https://github.com/sappho-x/Flow-of-Spans

Overview

Research area: Language modeling and text generation, sitting at the intersection of dynamic-vocabulary (span/retrieval-based) language models and Generative Flow Networks (GFlowNets).

Technical level: Intermediate. The text-generation framing and empirical results are accessible; the GFlowNets training objective (subtrajectory balance) and the flow-matching formalism are advanced.

Scope: The paper proposes FoSS (Flow of SpanS), a span-generation language model that constructs an explicitly directed-acyclic-graph (DAG) state space and trains it with GFlowNets, then evaluates it on open-ended generation, domain adaptation, scaling, and knowledge-intensive question answering.

What This Paper Is About

Standard autoregressive language models generate one token at a time from a fixed vocabulary, which forms a tree-shaped state space where every state has exactly one predecessor. Recent dynamic-vocabulary models sample retrieved multi-word spans instead, but they ignore that the same sentence can be assembled from spans of different lengths, meaning the real state space is a DAG with many paths to the same text — a structure those methods never model explicitly, so they are biased toward whichever segmentation path appeared in training.

FoSS addresses this by building a dynamic span vocabulary from flexibly segmented retrieved text, guaranteeing a DAG-structured state space, and training the model with GFlowNets so it explores many compositional paths at once and generalizes better.

Key Contributions

1. The first GFlowNets framework for span-level language modeling. Prior GFlowNets-based language models operated at the token level and remained confined to tree-structured state spaces; FoSS explicitly constructs a DAG-structured state space so GFlowNets can explore diverse compositional paths.

2. A DAG-Inducing span segmentation algorithm. Built on standard forward maximum matching with a controlled stochastic early-stopping mechanism, it produces multiple distinct segmentation trajectories per training sentence whose paths share common sub-sequences, explicitly modeling the DAG.

3. Specialized reward models for general text generation. Unlike prior GFlowNets work on tasks with clear answers (e.g., reasoning), FoSS designs a reward combining a language-model likelihood term and a preference-model term, trained with a Bradley-Terry log-likelihood loss and a score-centering regularizer, to balance fluency and human-text alignment.

4. A hybrid online–offline training procedure. Mini-batches mix trajectories from the forward policy, a reward-prioritized replay buffer of past high-reward trajectories, and the training set itself, with an initial fine-tuning phase on training data to reduce variance in sparse-reward regimes.

Main Findings

• In-domain text generation. On WikiText-103 (1,801,350 training, 3,760 development, 4,358 test samples), FoSS reaches MAUVE 30.78 under greedy decoding and 31.65 under nucleus sampling. This is a 5.51 absolute MAUVE improvement over CoG and 8.22 over the fine-tuned Transformer with nucleus sampling, and 2.04 over GFlowNets-FT.
• Greedy search stays competitive. Despite greedy search's known degeneration tendency, FoSS with greedy search beats the fine-tuned Transformer and CoG with nucleus sampling by 7.35 and 4.64 MAUVE, respectively.
• Out-of-domain adaptation. Trained on WikiText-103 and tested on the English portion of Law-MT (389,292 training samples; 2,000 development and 2,000 test), FoSS reaches MAUVE 27.84 (greedy) and 32.17 (nucleus), improving by 6.96 and 4.03 over the fine-tuned Transformer and CoG, without domain-specific training.
• GPT-4 preference evaluation. FoSS is preferred over all baselines across fluency, coherence, informativeness, and grammatical correctness. In domain, it is rated better than Transformer in 53% of cases and kNN-LM in 67%; out of domain, better than Transformer in 56% and kNN-LM in 75%. Against CoG, FoSS wins 42% (in domain, with 31% neutral) and 74% (out of domain); against GFlowNets-FT, 55% and 49%.
• Knowledge-intensive tasks. Across TruthfulQA, OpenBookQA, ARC-Challenge, MedMCQA, and Med-USMLE, FoSS leads most benchmarks, hitting 30.45, 26.20, 24.63, 27.44, and 25.27 accuracy. Versus the fine-tuned Transformer it improves TruthfulQA from 28.76% to 30.45% and OpenBookQA from 22.71% to 26.20%. CoG is stronger on MedMCQA (27.55 vs. 27.44).
• Ablation: the DAG structure matters most. A variant with all phrases removed from the vocabulary (forcing a tree) scores MAUVE 29.61/28.62 versus FoSS's 31.65/32.17, a 9.65% relative MAUVE drop and a 27.46% Diversity drop. Removing the PM or LM reward components also lowers MAUVE in every setting; the PM-only variant yields the highest Diversity scores (94.72 out of domain), while full FoSS gives the highest MAUVE.
• Scaling across three axes. With an En-Wiki memory, FoSS reaches MAUVE 27.88 (greedy) and 33.79 (nucleus). Quality improves consistently with larger retrieval datastores, more offline training data (FoSS beats fully-trained CoG and Transformer even trained on a 0.47% data fraction), and larger prefix encoders (GPT-2 base, Medium, Large, XL).
• Aggregate claim from the abstract. FoSS improves MAUVE by up to 12.5% over Transformer and gains 3.5% on knowledge-intensive tasks.

Methodology in Plain English

The researchers reframe text generation as a sequence of span choices rather than token choices. A state is the prefix plus everything generated so far; an action is appending a span — either a retrieved multi-word phrase or a single word from a fixed vocabulary or the termination symbol. Because a sentence like "ABCDEFGH" can be built as "AB"→"CD"→"EFGH" or "AB"→"C"→"DE"→"FG"→"H", several trajectories reach the same state, which makes the graph a DAG rather than a tree.

Training uses the subtrajectory balance objective, which matches forward and backward path probabilities against a state-flow function, and which the authors chose over flow matching, detailed balance, and trajectory balance for stability on longer sequences. The backward policy is a uniform distribution over all possible suffixes of the next state in the dynamic vocabulary. The loss sums only over pairs of states that are both complete sentences, focusing the learning signal on meaningful transitions.

The forward policy is a span language model inspired by CoG: a transformer prefix encoder maps the current state to a contextual vector, a span encoder produces embeddings for candidate phrases (fine-tuned BERT-base-cased, using start/end position embeddings combined through MLPs), and the policy distribution is a softmax over the dot product of the prefix vector and span embeddings. Prefix encoders for FoSS and all baselines are initialized from pre-trained GPT-2.

The reward is R(s_n) = exp(α·log p_LM(s_n|c) + (1−α)·log p_PM(s_n)) with α in (0,1). The LM term favors fluent, contextually relevant text; the PM term, a discriminator trained to score human-written references above policy-generated text, pushes toward human-like diversity and counteracts the repetitive patterns that high-likelihood LM regions tend to produce. Evaluation uses MAUVE and Diversity, with greedy decoding and nucleus sampling at top_p 0.95, plus GPT-4 pairwise judgments and classified-with-options accuracy on QA benchmarks (likelihood approximated by dynamic programming over generation paths).

Why This Matters

Impact on research. The paper makes the case that the state space of a span-based language model is genuinely a DAG, not a tree, and that this structure has to be modeled deliberately rather than left implicit in training data. It also extends GFlowNets beyond token-level language modeling and beyond tasks with verifiable answers, showing that a learned preference model can serve as the reward signal for open-ended generation. The ablations — especially removing the DAG structure or either reward term — give concrete evidence that the DAG state space is the load-bearing component, not an incidental detail.

Real-world applications:

• Domain adaptation without retraining: swapping the source text collection moves the model to a new domain (legal text here), which suits settings where labeled or fine-tuned data is scarce.
• Retrieval-augmented assistants and knowledge-grounded generation, where a richer datastore directly improves output quality.
• Knowledge-intensive question answering in specialized fields such as medicine and science, where the paper reports accuracy gains.
• Memory-scaling deployments where compute budget is better spent on retrieval corpora than on additional training runs.

Industry relevance. The plug-and-play memory scaling result matters for systems that already maintain large retrieval indexes: quality improves with a larger datastore without additional training. The 0.47% data-fraction result likewise speaks to sample efficiency, which is relevant when training data or annotation is the bottleneck.

Future Directions

• The appendix content is truncated mid-sentence in the supplied text, so several implementation details for the span encoder, dynamic vocabulary construction, and reward initialization are not fully reported here.
• The paper reports no setting where FoSS with an En-Wiki memory was matched against GFlowNets-FT or GDV — those two baselines are marked "-" in the scaling columns of Table 1, leaving the scaling comparison incomplete.
• MedMCQA is the one downstream benchmark where a baseline (CoG, 27.55) exceeds FoSS (27.44), an exception the paper does not analyze.
• The work opens the question of how far the DAG construction could extend to other generative domains, and how the reward balance α trades off between fluency and human-likeness in settings beyond those tested.

Target Audience

Researchers and practitioners in language modeling, retrieval-augmented generation, and generative modeling who are interested in moving beyond fixed token vocabularies or in applying GFlowNets to sequence generation. It is also relevant to engineers building retrieval-backed generation systems who want to understand how datastore size, training data volume, and model scale interact with output quality. Readers without prior exposure to GFlowNets will need background in Markov decision processes, flow matching, and the basics of autoregressive language models.

Authors’ abstract

Standard autoregressive language models generate text token-by-token from a fixed vocabulary, inducing a tree-structured state space when viewing token sampling as an action, which limits flexibility and expressiveness. Recent work introduces dynamic vocabulary by sampling retrieved text spans but overlooks that the same sentence can be composed of spans of varying lengths, lacking explicit modeling of the directed acyclic graph (DAG) state space. This leads to restricted exploration of compositional paths and is biased toward the chosen path. Generative Flow Networks (GFlowNets) are powerful for efficient exploring and generalizing over state spaces, particularly those with a DAG structure. However, prior GFlowNets-based language models operate at the token level and remain confined to tree-structured spaces, limiting their potential. In this work, we propose Flow of SpanS (FOSS), a principled GFlowNets framework for span generation. FoSS constructs a dynamic span vocabulary by segmenting the retrieved text flexibly, ensuring a DAG-structured state space, which allows GFlowNets to explore diverse compositional paths and improve generalization. With specialized reward models, FoSS generates diverse, high-quality text. Empirically, FoSS improves MAUVE scores by up to 12.5% over Transformer on text generation and achieves 3.5% gains on knowledge-intensive tasks, consistently outperforming state-of-the-art methods. Scaling experiments further demonstrate FoSS benefits from larger models, more data, and richer retrieval corpora, retaining its advantage over strong baselines.