langset 0.17.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.

LangSet adds a learned vector-emission interface to a pretrained language model. A model may emit one ordinary continuous embedding, an unordered continuous vector set through QueryBridgeEmission, or an autoregressive sequence of named-state superpositions through a fixed concept codebook. The output geometry is explicit rather than hidden in quantizer digits.

In practice, target_texts describes the ordered ticks of a process. For named-state training, a parallel concepts entry describes the distribution of admissible states within each tick. The model learns both the trajectory and the uncertainty inside each step — 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, with uncertainty rising as probability mass spreads across named concepts. 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 (the current state) → target_texts (the ordered ticks to emit). A matching concepts entry describes the named-state distribution inside each tick. The learned STOP decides how many ticks to emit; at inference, rollout(..., return_soft=True) returns the trajectory, its soft named-state mixtures, and the entropy of each mixture.

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

rows = [{
    "input_text": "<a state>",
    "target_texts": ["<the next tick>"],
    "concepts": [{"future": ["next state A", "next state B", "next state C"]}],
}]
model = LangSetModel.from_pretrained(
    "HuggingFaceTB/SmolLM2-135M",
    multi_latent=True,
    code_emit=True,
    n_codes=1,  # placeholder; ConceptObjective discovers the alphabet from `concepts`
)
Trainer(model, TrainingArguments(
    epochs=15,
    emission=ConceptObjective,
    concept_field="concepts",
), rows).train()

lat, lengths, soft, ent = model.rollout("<a state>", return_soft=True)
# soft = named-state superposition; ent = uncertainty over the concepts at each tick

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. Named-state emission. The autoregressive path predicts a distribution over a fixed concept codebook plus an independent STOP decision. The committed superposition is fed back into the stream, so each later state is conditioned on those already emitted.
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 multi-vector run chooses an explicit emission strategy in TrainingArguments: a named-state objective for ordered autoregressive state, or QueryBridgeEmission for an unordered continuous vector set. The target source and auxiliary terms remain independently injectable.

desired output	model shape	emission strategy
one continuous embedding	default single-latent model	default
ordered ticks, each a named-state superposition	multi_latent=True, code_emit=True	ConceptObjective
unordered set of continuous vectors	multi_latent=True	QueryBridgeEmission

Named state emission — ConceptObjective

ConceptObjective 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. target_texts[i] and concepts[i] describe the same tick: the text supervises its continuous context, while the concepts supervise its readable named state.

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=1,  # placeholder; replaced with the discovered concept alphabet
    res_dim=32,
)
Trainer(model, TrainingArguments(
    emission=ConceptObjective,
    concept_field="concepts",
    code_source="random",
), rows).train()

Each tick emits one concatenated vector: [named state | residual]. ConceptObjective discovers the concept alphabet from the training rows and installs a fixed codebook. The named-state portion is the differentiable mixture of that tick's concept distributions and can be read back by projection. When res_dim>0, a learned residual projection fills the trailing dimensions and is trained against the corresponding portion of the target-text embedding, preserving context outside the named alphabet. The two portions are normalized together before the vector is fed back for the next autoregressive tick. Set res_dim=0 for pure named state.

StateResidualObjective exposes the same representation through indexed per-tick member lists when a dataset already has a stable integer vocabulary. CodeSoftmaxObjective is the pure-state variant and therefore requires res_dim=0. 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 autoregressive named-state 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 concept 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 ConceptObjective, build_cot_loss_terms, cot_seed_texts

rows = [{
    "input_text": "...",
    "target_texts": ["...", "..."],
    "concepts": [{"phase": ["starting"]}, {"phase": ["complete"]}],
    "cot_text": "step-by-step reasoning ...",
}, ...]
model = LangSetModel.from_pretrained(
    "Qwen/Qwen3-1.7B-Base", multi_latent=True, code_emit=True, n_codes=1
)
Trainer(model, TrainingArguments(
    emission=ConceptObjective,
    concept_field="concepts",
    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 concept emission — CoTGenTerm self-skips on empty reasoning.

Superposition — one latent, several futures

To make one tick hold several possible states at once, put those alternatives in that tick's concepts entry. Lists assign equal mass; dictionaries provide explicit probabilities. This is distinct from target_texts: multiple target_texts entries mean multiple ordered ticks, while multiple concept members mean uncertainty inside one tick. rollout(..., return_soft=True) exposes the resulting mixture and its entropy.

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 is observed",
        "the release proceeds based on health",
    ],
    "concepts": [
        {"health": {"healthy": 0.7, "degraded": 0.2, "failing": 0.1}},
        {"outcome": {"expand": 0.7, "hold": 0.2, "rollback": 0.1}},
    ],
}]
model = LangSetModel.from_pretrained(
    "HuggingFaceTB/SmolLM2-135M",
    multi_latent=True,
    code_emit=True,
    n_codes=1,
)
Trainer(model, TrainingArguments(
    emission=ConceptObjective,
    concept_field="concepts",
), rows).train()

lat, lengths, soft, ent = model.rollout(
    "release 2.4 has entered canary deployment",
    return_soft=True,
)
# Two ordered ticks; each `soft` entry is a distribution over that tick's named states.

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.

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.
• 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

Named-state emission is an ordered autoregressive sequence: each state conditions the next. 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 continuous set emitter
    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 concept 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 ConceptObjective when the emission is an ordered rollout whose steps condition each other. Reach for QueryBridgeEmission when it is 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=; multi-vector training requires an explicit emission= choice.

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(
    emission=ConceptObjective,
    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) is 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.