langset 0.16.0

A short path to a world model in your LLM: few-shot fine-tune it to predict in latent space (a JEPA world model) — or emit a single latent as a bespoke embedding model, SetFit-ready.

langset — a short path to a world model in your LLM

langset few-shot fine-tunes a pretrained LLM to predict in latent space. Bolt a tiny head onto the backbone and it emits a sequence of latents in its own token stream — one per step, with a learned STOP — where each latent holds a calibrated superposition of possible next states. Language models and world models are usually posed as a choice between them; this is the other answer — the same pretrained LLM predicts its own latent futures, with no new primitive required. You get there few-shot, by describing the states in words.

A trained langset world model rolling out a maze — each tick is one emitted latent holding a whole set of next states (the lime frontier), the caption counting how many. It predicts the distribution of where the search could be, not a single guess. See the superposition example.

The same machinery with a single latent per input is a bespoke embedding model, SetFit-ready — that tier still works and is documented below. But the reason to reach for langset is the world model.

The idea

Language models and world models are usually posed as a choice. It's a false binary — and underneath it is an architecture question: does moving from next-token prediction to latent prediction require a new primitive? It doesn't.

LeCun names what a language model lacks — an abstract latent variable, discrete, low-capacity, collapse-resistant. langset gives a pretrained model exactly that, inside its own vocabulary: a self-supervised JEPA latent emitted as FSQ digit tokens through the ordinary softmax, trained toward the model's own latent (kept from collapsing by a standard self-supervised mechanism) — no decoder, no target outside the model. Latents and text co-train under one next-token cross-entropy. Quantize the latent, keep the token.

In practice you write the states of a process as target_texts, and the model learns to emit the set of admissible next states — no labels, no hand-built simulator. What that buys:

• 🔎 Readout, not acquisition. The emitted latents surface world knowledge the pretrained model already held. Scramble the vocabulary so the task keeps its dynamics but loses its meaning, and much of that knowledge falls away with it — evidence the latents read it out of pretraining, not out of the task.
• 🎲 One latent, several futures. A single emission can hold several live continuations at once — its uncertainty rising with how many remain, dropping none of them. A calibrated superposition, read from the digit softmax (see the maze-superposition example).
• 💬 Reason in words, then emit. When holding those futures takes reasoning rather than perception, one emission isn't enough on its own — but because the latent shares the stream, the model can reason in tokens first, and that reasoning is what calibrates its uncertainty.
• 🧭 Steer with text. The target_text defines what the latent tracks; rewrite it to re-steer — no architecture changes, no relabeling.

Install

pip install langset

Quickstart — a world model

Rows are input_text (a state) → target_texts (the set of possible next states). The model learns to emit that set, deciding how many via a learned STOP; at inference rollout(..., return_soft=True) reads the emitted set back plus its per-step entropy — the model's own calibrated uncertainty.

from langset import LangSetModel, Trainer, TrainingArguments

rows = [{"input_text": "<a state>", "target_texts": ["<next state A>", "<next state B>", "<next state C>"]},
        # ...one row per state; target_texts is the SET of admissible next states
       ]
model = LangSetModel.from_pretrained("HuggingFaceTB/SmolLM2-135M", multi_latent=True)   # FSQ set-emission head
Trainer(model, TrainingArguments(epochs=15), rows).train()

lat, lengths, soft, ent = model.rollout("<a state>", return_soft=True)
# soft = the expected latent SET (the superposition); ent = per-step entropy (higher = a more open future)

examples/maze-superposition is the end-to-end reference: it trains a world model on a maze search-frontier and measures the headline property with langset.probes — the emitted latent's entropy tracks the true number of possible next states (a calibrated superposition, not one guess).

How it works — JEPA in the token stream

1. Predict in latent space (JEPA). Each emitted latent is trained to match the stop-grad target latents of its next states (in-batch negatives keep distinct states apart). Predicting a target-encoder's latents — not pixels, not tokens — is exactly the JEPA objective, here run inside a pretrained LLM.
2. Token-native emission. Each latent is finite-scalar-quantized (FSQ) into per-dimension digits the model predicts, with STOP folded into dim-0's softmax; every emitted latent is fed back into the stream so the next one is conditioned on those already emitted. The latent is literally a token — text and latents share one softmax/CE interface.
3. Anti-collapse is the JEPA apparatus. A stop-grad EMA target twin (BYOL/JEPA) supplies the targets by default; inject SIGRegTarget for the EMA-free LeJEPA alternative (details below).
4. Collapse-aware selection. langset selects on held-out retrieval/reconstruction with a hard penalty on any collapse of the geometry — never on the training loss (which collapse can game).

World-model knobs

Every knob below is a strategy injected into TrainingArguments, not a boolean on a monolith — the defaults give you the FSQ + EMA-twin world model, and each injection swaps one piece.

Named state emission — ConceptObjective

FSQ learns an implicit discrete code. When the state vocabulary is known, ConceptObjective instead constructs each emitted latent from a soft mixture over named members. Lists split probability mass equally; dictionaries provide explicit weights. Each facet gets its own softmax, so values in stage compete with other stages without competing with values in result.

from langset import LangSetModel, Trainer, TrainingArguments
from langset.strategies import ConceptObjective

rows = [{
    "input_text": "release 2.4 has entered canary deployment",
    "target_texts": [
        "the canary receives a small traffic share",
        "health checks remain stable",
        "the release expands to all regions",
    ],
    "concepts": [
        {
            "phase": ["canary"],
            "traffic": {"low": 0.9, "medium": 0.1},
            "health": ["observing"],
        },
        {
            "phase": ["canary"],
            "traffic": ["medium"],
            "health": ["healthy"],
        },
        {
            "phase": ["complete"],
            "traffic": ["full"],
            "health": ["healthy"],
        },
    ],
}]

model = LangSetModel.from_pretrained(
    "HuggingFaceTB/SmolLM2-135M",
    multi_latent=True,
    code_emit=True,
    n_codes=4,
    res_dim=32,
)
Trainer(model, TrainingArguments(
    emission=ConceptObjective,
    concept_field="concepts",
    code_source="random",
), rows).train()

The emitted vector is [named state | residual]. The state portion is a mixture over a fixed codebook and can be read back by projection; the residual carries context that the named alphabet cannot express. Set res_dim=0 for pure named state. StateResidualObjective exposes the same mechanism through indexed per-tick member lists when a dataset already has a stable integer vocabulary. Codebooks can be random orthonormal, model-embedded, orthogonalized model embeddings, or vectors encoded through the target path.

Anti-collapse: EMA twin (default) vs SIGReg (LeJEPA)

By default the multi-latent path prevents representation collapse with an EMA target twin — a stop-grad copy of the model whose slowly-moving latents are the targets, so the online model can't trivially match a target that moves with it. Injecting target_source=SIGRegTarget swaps this for SIGReg (Sketched Isotropic Gaussian Regularization, from LeJEPA, Balestriero & LeCun 2025, arXiv:2511.08544): no twin, no stop-grad — targets come from the live encoder, and collapse is prevented instead by regularizing the pre-quantization latent z = down_proj(·) toward an isotropic Gaussian (an Epps–Pulley goodness-of-fit test over random 1-D projections). The isotropic Gaussian is the distribution LeJEPA proves is uniquely privileged for downstream identifiability (arXiv:2605.26379).

from langset.strategies import SIGRegTarget
Trainer(model, TrainingArguments(target_source=SIGRegTarget, sigreg_lambda=0.3), rows).train()

Anti-collapse is chosen by injecting a different target-source strategy, not a boolean flag — the default target_source=EMATwinTarget and SIGRegTarget are interchangeable implementations (see strategies.py).

Trade-offs:

	EMA twin (default)	SIGReg (target_source=SIGRegTarget)
memory	a full frozen copy of the backbone in VRAM	none — no twin
per-step cost	one extra target forward	one Gaussian-regularizer pass (cheap)
anti-collapse	stop-grad target	isotropic-Gaussian penalty on pre-quant z
separation term	pairs with in-batch InfoNCE (lam_multi_nce)	replaces it (InfoNCE auto-gated off)

Empirically SIGReg matches or beats the twin on local calibration but retains a small gap on global structure at high sigreg_lambda; tune sigreg_lambda ≈ 0.3 (over-diversification washes out global structure; too small under-constrains). Implementation note (see sigreg.py): the test is center-only, not standardized — standardizing per-dim before the Gaussianity test is scale-invariant and silently defeats anti-collapse (a collapsed batch passes with zero gradient). SIGReg is research-grade; the EMA twin remains the validated default.

CoT-conditioned emission — build_cot_loss_terms + cot_seed_texts

By default the model emits latents straight from the input seed. Injecting the CoT strategy pair inserts a chain-of-thought step: the model is co-trained to generate a per-row reasoning string (a cot_text column) before it emits the latents, and the latent forward is conditioned on seed + CoT. Two objectives share one optimizer step — the FSQ latent loss, and a next-token cross-entropy on the CoT string (weight lam_cot) through the tied embedding, i.e. the same CE machinery the latents already use.

It's selected by injecting two strategies (not a flag): loss_terms=build_cot_loss_terms adds the isolated CoTGenTerm, and seed_builder=cot_seed_texts conditions the emission on the reasoning:

from langset.strategies import build_cot_loss_terms, cot_seed_texts
rows = [{"input_text": "...", "target_texts": ["...", "..."], "cot_text": "step-by-step reasoning ..."}, ...]
model = LangSetModel.from_pretrained("Qwen/Qwen3-1.7B-Base", multi_latent=True)
Trainer(model, TrainingArguments(loss_terms=build_cot_loss_terms, seed_builder=cot_seed_texts, lam_cot=1.0),
        rows).train()

Why: some target latents aren't a direct function of the surface input — they need an intermediate inference the model can do but doesn't surface in one hop. Letting the model think in tokens first, then emit, lets that reasoning inform the latent, in the spirit of chain-of-thought reasoning (Wei et al., arXiv:2201.11903) — but here the reasoning and the latent are co-trained in one token stream rather than the reasoning being an external prompt, and unlike COCONUT's continuous latent thoughts (Hao et al., arXiv:2412.06769) the CoT stays in readable token space sharing the softmax/CE interface.

Trade-offs: the two forward+backward passes run separately so their autograd graphs never coexist (peak activation is max(latent, cot), not the sum) — but you still pay a second forward per step, and CoT blocks are long (train-time cost scales with CoT length, not the short latents). You need a cot_text column: at train time it's a teacher-forcing target; the measured result is the lift from the model learning to produce its own CoT (self-generated reasoning helps even when the CoT text itself came from a stronger teacher, which is not a fair ceiling to compare against). Without the injected strategies (or with an absent cot_text column) the path is byte-identical to the plain FSQ emission — CoTGenTerm self-skips on empty reasoning.

Superposition — one latent, several futures

To make one emitted latent hold a calibrated superposition — several possible next states at once, its uncertainty rising with how many are live — supervise it directly: write a target text that describes the whole set of futures. There's no special loss or batching trick — the set is in the target, so the single latent that matches it has to hold the set, and its per-step entropy (read via rollout(..., return_soft=True)) tracks how open the future is. The only knob is selector=last_epoch_selector — retrieval MRR is meant to fall as the latent spreads over the set, so it must not gate checkpoint selection.

from langset import LangSetModel, Trainer, TrainingArguments
from langset.strategies import last_epoch_selector

# one row per state; each target text names the SET of admissible next states — a wider set should read back hotter
rows = [
    {"input_text": "a fair coin is flipped",         "target_texts": ["it lands heads or tails"]},                    # 2 live
    {"input_text": "a six-sided die is rolled",       "target_texts": ["it shows one, two, three, four, five, or six"]},# 6 live
    {"input_text": "water at sea level reaches 100°C", "target_texts": ["it boils to steam"]},                         # 1 live
    # ...many states, each target enumerating its own live next-states
]
model = LangSetModel.from_pretrained("HuggingFaceTB/SmolLM2-135M", multi_latent=True)
Trainer(model, TrainingArguments(selector=last_epoch_selector), rows).train()

# read the superposition back: entropy should rise with the number of live outcomes (die > coin > boil)
_, _, soft, ent = model.rollout("a six-sided die is rolled", return_soft=True)   # soft = expected latent, ent = per-step entropy

The same shape scales to a real world model: in examples/maze-superposition, each tick's target is the set of active search-frontier cells, and the emitted latent's entropy provably tracks the branch count (corr(entropy, k) = +0.34–0.39) — the example measures that calibration with langset.probes.

Why last_epoch_selector: the default checkpoint selection (retr_mrr) rewards a collapsed one-future-per-state geometry — exactly the wrong signal here, since you want retrieval MRR to fall as the latent spreads over the set. (There is also a plain snapshot_every=N knob that saves the online weights to {output_dir}_ep{N}, {output_dir}_ep{2N}, … after every N epochs — independent of the eval cadence, separate from the best-so-far restore — to keep a checkpoint trajectory for offline evaluation.)

Text replay — learn_field / learn_ratio

Fine-tuning a backbone on the emit objective can slowly erode its plain next-token ability. Text replay interleaves ordinary language-model steps into training to rehearse it: tag some rows learn (via a learn_field column) and, with probability learn_ratio before each normal batch, the trainer runs a next-token cross-entropy on those rows (input_text → target, through the tied input embedding — no separate LM head) as its own optimizer step. It's the standard rehearsal remedy for catastrophic forgetting (Robins 1995) applied to the backbone.

rows = [
    {"input_text": "...", "target_texts": ["...", "..."]},          # normal emit rows
    {"input_text": "domain fact to keep fluent", "target_texts": ["..."] , "tag": "learn"},   # rehearsed as text
]
Trainer(model, TrainingArguments(learn_field="tag", learn_ratio=0.2), rows).train()

Works on both the single-latent and multi-latent paths (multi-latent uses target_texts[0] as the replay target). learn_ratio=0 or no learn rows = off (byte-identical). Use it when the backbone must stay fluent on a domain while you retrain its emission geometry; leave it off for pure embedding tasks.

Also: a bespoke embedding model

The single-latent path emits one vector per input instead of a set — a bespoke embedding model in a geometry you define, and a drop-in Sentence-Transformer body for SetFit. Same one idea: the target_text is the geometry (describe instrumentation and the space clusters by instrumentation; describe vocals and emotion and it clusters by those). Rows are input_text → a single target_text:

from langset import LangSetModel, Trainer, TrainingArguments

rows = [{"input_text": "an hour-long track of detuned riffs that never break stride, at the pace of continental drift",
         "target_text": "glacial detuned doom-metal, sludgy and hypnotic, buried roared vocals"},
        # ...
       ]
model = LangSetModel.from_pretrained("HuggingFaceTB/SmolLM2-135M")   # any HF causal LM, single-latent
Trainer(model, TrainingArguments(), train_dataset=rows).train()

z = model.encode(["a wall of downtuned fuzz that buries the vocals under sheer volume"])
print(z.shape)   # (1, 576) — a latent in the backbone's own space

See examples/sounds_like/ for the full reference task (album review → "how it sounds" latent).

With SetFit

The name is the chain: lang·set·fit — a language model emits into the set geometry (langset, usable on its own), which then fits a classifier. model.as_sentence_transformer() is a drop-in SetFit model_body. The clean distinction: SetFit answers with a label; langset answers with a latent.

	reach for SetFit	reach for langset
your answer is	a label (fixed classes)	a point in a space — retrieval, "find similar", ranking, clustering
you define the target by	enumerating classes	a description of the geometry ("how it sounds")
your input	text to classify	text or an identifier — leans on the LLM's world knowledge

Use SetFit alone for plain few-shot classification; use langset when the answer is a geometry you'll retrieve / rank / cluster in; chain langset → SetFit when a task-shaped body helps the classifier.

pip install "langset[setfit]"      # pins the verified composition window

from sklearn.linear_model import LogisticRegression
from setfit import SetFitModel

clf = SetFitModel(model_body=langset_model.as_sentence_transformer(),
                  model_head=LogisticRegression(max_iter=2000),    # direct construction needs an explicit head
                  labels=[...])
clf.fit(x_train, y_train, num_epochs=1)        # frozen body + head — the robust path
clf.predict(["..."])

Dependency alignment. langset runs modern backbones (transformers ≥4.41, incl. Qwen3). langset[setfit] pins transformers 4.46.x / torch <2.5 / Python 3.10–3.12: SetFit imports training_args.default_logdir (removed after 4.46), and 4.46 + torch ≥2.5 trips a torch.distributed.tensor bug. Use the frozen-body SetFitModel.fit/predict path above; the full setfit.Trainer (fine-tunes the body) is fragile in this window.

A set of latents, outside world models

The variable-length set emission is useful even when the set isn't a set of futures — anywhere one input carries an unknown number of things, and a single averaged vector would blur them together:

• Multi-item extraction — a latent per entity / keyphrase / skill / ingredient, retrieved against a reference bank (examples/ner-multi-latent/ does this for named entities).
• Multi-vector retrieval — represent a query or document as a set of latents (ColBERT-style late interaction) instead of one averaged vector.
• Multi-aspect / multi-label — one latent per facet ({brand, category, material}) or applicable label.

from langset import LangSetModel
import torch.nn.functional as F

m = LangSetModel.load("path/to/checkpoint", device="cpu")
bank = ["PER: Barack Obama", "LOC: Berlin", "PER: Angela Merkel", "ORG: Apple", "LOC: California"]
zb = F.normalize(m.emit(bank).float(), dim=-1)
lat = F.normalize(m.rollout("Barack Obama visited Berlin to meet Angela Merkel.").float(), dim=-1)
for v in lat:                                                        # one latent per entity, count set by a learned STOP
    print(bank[int((v @ zb.T).argmax())])                           # PER: Barack Obama / LOC: Berlin / PER: Angela Merkel

Retrieval-preserving multi-vector — emission=QueryBridgeEmission

The default emission is FSQ: discrete digits, emitted autoregressively (each latent conditioned on the last), an ordered sequence. When the output is instead an unordered set of items in a real vector space — a document's facts, a passage's claims, the entities in a sentence — inject emission=QueryBridgeEmission for the parallel-query family: N learned query tokens cross-attend the backbone's per-token states in one forward pass and emit an unordered set of L2-normalized vectors, each with a validity gate that decides how many are real (≤ N). No autoregression, no digit softmax — continuous vectors, one shot.

Pair it with a frozen backbone and the frozen FrozenEncoderTarget, and it becomes a pure add-on to an existing embedding model — the emitted vectors are trained to match the encoder's own embedding of each target text (DETR-style set-matching), so they live in the same space the base model already retrieves in:

from langset import LangSetModel, Trainer, TrainingArguments
from langset.bridge_emission import QueryBridgeEmission, FrozenEncoderTarget

# each row: a passage -> the SET of facts it states (as text). target_texts defines the set; the bridge learns
# to emit one vector per fact, each matched to the frozen encoder's embedding of that fact.
rows = [{"input_text": "The X200 ships in slate and sand, weighs 1.2 kg, and runs 14 hours on a charge.",
         "target_texts": ["available in slate and sand", "weighs 1.2 kg", "14-hour battery life"]},
        # ...one row per passage; target_texts = the facts to surface
       ]
model = LangSetModel.from_pretrained("Qwen/Qwen3-Embedding-0.6B", multi_latent=True, freeze_backbone=True)
Trainer(model, TrainingArguments(
    emission=QueryBridgeEmission,        # one-pass parallel-query set-emitter (vs the default FSQ rollout)
    target_source=FrozenEncoderTarget,   # targets = the frozen encoder's own embedding of each fact
    lam_multi_nce=0.0,                   # the bridge's base loss already runs the per-fact contrastive term
    n_queries=16,                        # max facts per input; the validity gate picks the actual count
), rows).train()
# at inference the model emits a validity-gated set in one pass — one unit vector per fact, count decided by the gate.

Why reach for it:

• 🔒 Retrieval preserved by construction. With freeze_backbone=True the base embedder's single-vector geometry is untouched — the bridge is a pure add-on, so model.encode() and its cosine behavior stay exactly the base model's. You bolt a multi-vector head onto an existing retriever without disturbing it.
• 🧷 A set, in the same space, cosine-ready. Every emitted vector is a normal unit vector in the encoder's space, so a query can retrieve the specific fact instead of a blurred whole-document average — and the set drops straight into any cosine ANN / multi-vector index (ColBERT-style late interaction), no new primitive.
• ⚡ One pass, unordered. Unlike the autoregressive FSQ rollout, all items emit together — cheaper at inference and the natural shape when order is meaningless.
• 🎯 Hard negatives fold in. Set hard_neg_field="..." (a per-row list of confusable near-misses — e.g. a fact that differs from the real one in a single attribute) and they join the emitter's own contrastive denominator, so each vector learns to separate the look-alikes a single averaged vector blurs together.

Reach for FSQ (the default) when the emission is an ordered rollout whose steps condition each other — a world model. Reach for QueryBridgeEmission when it's an unordered set of vectors in an existing embedding space — multi-vector retrieval or multi-item extraction where the base retriever's geometry must stay intact. Injected via emission= / target_source=; with no emission set the default FSQ path is byte-identical.

Big batches without the memory — GradCache

The contrastive loss wants a large batch — every other row is an in-batch negative, and more negatives means a sharper geometry. But holding the whole batch's activation graph at once OOMs. GradCache (Gao et al. 2021) gives you the exact large-batch gradient while peak activation stays at one small chunk: forward the full batch in no_grad sub-chunks (only the [B, d] embeddings survive), run the loss on the full set and cache each embedding's gradient, then re-forward each chunk with grad and inject the cached gradients. batch_size becomes the big effective batch; gc_chunk is the memory unit you tune to the card.

# single-latent (embedding model): the whole contrastive loss is cached -> BIT-EXACT vs the direct big-batch step
Trainer(model, TrainingArguments(
    grad_cache=True, batch_size=1024, gc_chunk=32,   # 1024 in-batch negatives, activation = 32 rows
    lam_recon=0,                                     # required: recon needs the per-token backbone graph, not the pooled embedding
)).train()

# multi-latent (world model): the cross-batch InfoNCE is cached exact; scheduled sampling is preserved
Trainer(model, TrainingArguments(
    grad_cache=True, batch_size=512, gc_chunk=16,
    ss_prob=0.25,                                    # kept — a SHARED self-feed mask makes the rollout replay identically across phases
)).train()

Requirements and what stays exact:

	single-latent	multi-latent
requires	dropout=0, lam_recon=0	dropout=0, EMA twin (not SIGRegTarget)
scheduled sampling (ss_prob>0)	n/a	supported — the coin flips are sampled once and shared across both phases
exactness	bit-exact (the whole loss is cached)	cross-batch InfoNCE exact; the per-row base loss (stop/dims/recon) is accumulated per chunk (a small gc_chunk-dependent reweighting, negligible in practice)

gc_chunk=0 (default) or grad_cache=False leaves training byte-identical. GradCache is only defined for the contrastive path, so SIGRegTarget (a cross-batch regularizer) and the FSQ label subspace (which reads the rollout logits, not the cached latent) are rejected with a clear error when grad_cache=True.

Long rollouts without the memory — KV cache

The multi-latent world model rolls out autoregressively at train time: it emits one latent per tick and feeds each back to condition the next, so a rollout of T ticks runs the backbone T times over the growing sequence. With no cache that re-encodes the shared prompt every tick, and because every tick carries its own emission loss, autograd keeps all T full-sequence forwards alive at once — activation memory scales with ticks, which forces gradient checkpointing (and its recompute tax) on any long-rollout corpus.

kv_cache=True forwards the prompt once, then feeds each latent token alone against the cached prefix K/V (one new position per tick). Peak activation becomes ~one prompt forward plus T single tokens instead of T full-prefix forwards, so long rollouts train without checkpointing.

# world model on a long-rollout corpus: no grad_ckpt needed
Trainer(model, TrainingArguments(
    kv_cache=True,        # prompt forwarded once, single-token hops vs full-prefix recompute
    ss_prob=0.25,         # scheduled sampling preserved (same self-feed decisions as the recompute path)
)).train()

The cache is not detached, so gradients flow through it to the backbone exactly as in a full-sequence forward — the emitted logits are numerically identical to the recompute rollout (verified to ~1e-7, tests/test_kv_cache.py), under both teacher forcing and self-feed. kv_cache=False (default) leaves training byte-identical. Two constraints: it is mutually exclusive with gradient checkpointing (HF disables the cache under checkpointing — use one or the other), and it targets standard (non-PLE) backbones. Both are checked with a clear error.