diff --git a/README.md b/README.md index c4160c7..1c7a167 100644 --- a/README.md +++ b/README.md @@ -26,22 +26,18 @@ Language models and world models are usually posed as a *choice*. It's a false b 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.** +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 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: +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 โ€” its - uncertainty rising with how many remain, dropping none of them. A *calibrated* superposition, read from the - digit softmax (see the [maze-superposition example](examples/maze-superposition)). +* ๐ŸŽฒ **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](examples/maze-superposition). * ๐Ÿ’ฌ **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. @@ -56,21 +52,34 @@ 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. +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. ```python from langset import LangSetModel, Trainer, TrainingArguments +from langset.strategies import ConceptObjective -rows = [{"input_text": "", "target_texts": ["", "", ""]}, - # ...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() +rows = [{ + "input_text": "", + "target_texts": [""], + "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("", return_soft=True) -# soft = the expected latent SET (the superposition); ent = per-step entropy (higher = a more open future) +# soft = named-state superposition; ent = uncertainty over the concepts at each tick ``` **[examples/maze-superposition](examples/maze-superposition)** is the end-to-end reference: it trains a world @@ -83,10 +92,9 @@ guess). 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. +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](#anti-collapse-ema-twin-default-vs-sigreg-lejepa)). @@ -95,15 +103,22 @@ guess). ## 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. +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` -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`. +`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. ```python from langset import LangSetModel, Trainer, TrainingArguments @@ -139,7 +154,7 @@ model = LangSetModel.from_pretrained( "HuggingFaceTB/SmolLM2-135M", multi_latent=True, code_emit=True, - n_codes=4, + n_codes=1, # placeholder; replaced with the discovered concept alphabet res_dim=32, ) Trainer(model, TrainingArguments( @@ -149,15 +164,21 @@ Trainer(model, TrainingArguments( ), 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. +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 multi-latent path prevents representation collapse with an **EMA target twin** โ€” a stop-grad +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](https://arxiv.org/abs/2511.08544)): @@ -195,18 +216,31 @@ the validated default. 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 +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: ```python -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() +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 @@ -223,33 +257,46 @@ long (train-time cost scales with CoT length, not the short latents). You need a 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. +is byte-identical to the plain concept 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. +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. ```python from langset import LangSetModel, Trainer, TrainingArguments -from langset.strategies import last_epoch_selector +from langset.strategies import ConceptObjective -# 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() +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() -# 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 +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](examples/maze-superposition)**, @@ -257,12 +304,6 @@ each tick's target is the set of active search-frontier cells, and the emitted l the branch count (`corr(entropy, k) = +0.34โ€“0.39`) โ€” the example *measures* that calibration with [`langset.probes`](src/langset/probes.py). -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** @@ -350,7 +391,7 @@ The variable-length **set** emission is useful even when the set isn't a set of 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/`](examples/ner-multi-latent/) does this for named entities). + 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. @@ -369,8 +410,8 @@ for v in lat: # one laten ### 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** โ€” +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 @@ -392,7 +433,7 @@ rows = [{"input_text": "The X200 ships in slate and sand, weighs 1.2 kg, and run ] 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) + 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 @@ -408,16 +449,16 @@ Trainer(model, TrainingArguments( * ๐Ÿงท **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 +* โšก **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 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** โ€” +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=`; with no `emission` set the default FSQ path is byte-identical. +`emission=` / `target_source=`; multi-vector training requires an explicit `emission=` choice. ## Big batches without the memory โ€” GradCache @@ -438,6 +479,7 @@ Trainer(model, TrainingArguments( # 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() @@ -452,8 +494,8 @@ Requirements and what stays exact: | 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`. +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 diff --git a/examples/danger-corridor/danger_eval.py b/examples/danger-corridor/danger_eval.py deleted file mode 100644 index 01396c7..0000000 --- a/examples/danger-corridor/danger_eval.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Held-out-animal eval for the danger-corridor world -- the pretrain-WIN metric, made causal. - -The eval corpus draws from a DISJOINT animal pool (gen_danger.py build ... heldout), so every animal here was -never seen in training. The walk is mechanical; the ONE decision that needs knowledge is the crux tick, where -you stand on the animal's cell and either die (dangerous) or pass (harmless). Danger is never in the text. - -We read that decision straight out of the rollout. For each held-out episode: - - roll the model L steps; take the emitted latent at the crux cell k; score it against the row's TWO - candidates, which share the position prefix "at cell k: the {animal} ..." and differ ONLY in outcome - (".. attacks, you die" vs ".. is harmless, keep walking"). argmax cosine = the model's predicted outcome. - -Because the candidates are position-matched, the accuracy is PURE danger knowledge, not maze skill. If the -model STOPs at or before the crux, that is itself a death prediction (it never reached the animal alive). -Truth = the hidden danger label. Balanced corpus -> chance = 0.5. Reports overall accuracy plus per-class -recall (dangerous: predicts death; safe: predicts pass) and a secondary rollout-length outcome check. - - python danger_eval.py --data danger_eval.npz --ckpt danger_model -""" -from __future__ import annotations - -import argparse -import json - -import numpy as np -import torch -import torch.nn.functional as F -from langset import LangSetModel - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--data", required=True, help="HELD-OUT eval corpus npz from gen_danger.py") - p.add_argument("--ckpt", required=True) - p.add_argument("--max-rows", type=int, default=1000) - p.add_argument("--device", default="cuda") - a = p.parse_args() - - z = np.load(a.data, allow_pickle=True) - seeds = [str(s) for s in list(z["seed"])][:a.max_rows] - danger = [int(x) for x in list(z["danger"])][:a.max_rows] - crux = [int(x) for x in list(z["crux_cell"])][:a.max_rows] - die_txt = [str(x) for x in list(z["crux_die"])][:a.max_rows] - cont_txt = [str(x) for x in list(z["crux_cont"])][:a.max_rows] - L = int(z["corridor_len"][0]) - - m = LangSetModel.load(a.ckpt, device=a.device) - m.eval() - - # encode both position-matched crux candidates per row -> the only difference is the danger outcome - die_emb = torch.cat([F.normalize(m.emit(die_txt[s:s + 256]).float(), dim=-1).cpu() - for s in range(0, len(die_txt), 256)]) - cont_emb = torch.cat([F.normalize(m.emit(cont_txt[s:s + 256]).float(), dim=-1).cpu() - for s in range(0, len(cont_txt), 256)]) - - preds_die, len_pred_die = [], [] # crux-latent decision; rollout-length decision - crux_ent, walk_ent = [], [] # cat-in-the-box: FSQ entropy at the animal vs a plain step - B = 64 - for s0 in range(0, len(seeds), B): - chunk = seeds[s0:s0 + B] - Lat, lengths, _sL, ent = m.rollout(chunk, max_steps=L, return_lengths=True, return_soft=True) - Lat = F.normalize(Lat.float(), dim=-1).cpu() # [b, L, d] - ent = ent.float().cpu().numpy() # [b, L] per-tick FSQ entropy - for j, gi in enumerate(range(s0, s0 + len(chunk))): - k, ln = crux[gi], int(lengths[j]) - if ln <= k: # stopped at/before the animal -> a death prediction - preds_die.append(1) - else: - sim_die = float(die_emb[gi] @ Lat[j, k]) - sim_cont = float(cont_emb[gi] @ Lat[j, k]) - preds_die.append(1 if sim_die > sim_cont else 0) - len_pred_die.append(1 if ln < L else 0) # secondary: did the walk end before the exit? - if ln >= 1: # entropy at the crux (the box) vs the first walk step - crux_ent.append(float(ent[j, min(k, ln - 1)])) - walk_ent.append(float(ent[j, 0])) - else: - crux_ent.append(float("nan")); walk_ent.append(float("nan")) - - y = np.array(danger); pd = np.array(preds_die); pl = np.array(len_pred_die) - ce = np.array(crux_ent); we = np.array(walk_ent); ok = (pd == y) - box = { - "crux_entropy": round(float(np.nanmean(ce)), 4), # the cat in the box: how open is die/survive - "walk_entropy": round(float(np.nanmean(we)), 4), # a certain step -> the model's entropy floor - "excess_at_crux": round(float(np.nanmean(ce - we)), 4), # lift at the decision (knowledge collapses it) - "crux_entropy_when_correct": round(float(np.nanmean(ce[ok])), 4) if ok.any() else None, - "crux_entropy_when_wrong": round(float(np.nanmean(ce[~ok])), 4) if (~ok).any() else None, - } - def stats(pred): - acc = float(np.mean(pred == y)) - dang = float(np.mean(pred[y == 1] == 1)) if (y == 1).any() else float("nan") # recall on dangerous - safe = float(np.mean(pred[y == 0] == 0)) if (y == 0).any() else float("nan") # recall on safe - return {"acc": round(acc, 4), "dangerous_recall": round(dang, 4), "safe_recall": round(safe, 4)} - - out = {"ckpt": a.ckpt, "n": len(y), "corridor_len": L, "frac_dangerous": round(float(y.mean()), 4), - "chance": 0.5, "crux_decision": stats(pd), "rollout_length_decision": stats(pl), - "cat_in_the_box": box} - print("=== DANGER-CORRIDOR HELD-OUT DIE-VS-SURVIVE ===") - print(json.dumps(out)) - - -if __name__ == "__main__": - main() diff --git a/examples/danger-corridor/gen_danger.py b/examples/danger-corridor/gen_danger.py deleted file mode 100644 index 944d5fe..0000000 --- a/examples/danger-corridor/gen_danger.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Danger-corridor world -- the CLEANEST pretrain-WIN task: a mechanical walk gated on ONE known fact. - -A 1xL corridor (cells 0..L-1). You start at cell 0, the exit is cell L-1, and a single animal waits at a -random interior cell k. You must walk the whole corridor one cell at a time, so you ALWAYS step onto the -animal (the path forces overlap). The rule, stated in every seed: - - a dangerous animal kills you; a harmless animal lets you pass. - -Everything about the walk is mechanical -- a random-init transformer learns "advance one cell per tick" from -a handful of examples (maze already proved this). The ONLY thing it cannot learn for a HELD-OUT animal is that -animal's danger ("does a mamba kill you?"): danger is NEVER in the text, only the animal's NAME. Only -pretraining supplies it. So the crux tick -- the one where you stand on the animal -- is a binary decision -(die vs pass) gated purely on real-world knowledge. That makes the pretrained knowledge CAUSAL (it drives the -transition), not merely probeable. - -Ground truth is the hidden DANGER table below. Train and eval corpora draw from DISJOINT animal pools, so a -pass on eval means the danger knowledge GENERALIZED to unseen species, not that it was memorized. - - python gen_danger.py peek - python gen_danger.py build 6000 danger_train.npz # train pool, seed 0 - python gen_danger.py build 1000 danger_eval.npz 999 heldout # HELD-OUT animal pool, seed 999 -""" -from __future__ import annotations - -import sys - -import numpy as np - -# ---- hidden ground-truth danger table (never shown to the model; only the animal NAME appears) ------------- -# Curated for unambiguous real-world consensus: obligate killers vs plainly harmless. Single lowercase tokens. -DANGEROUS = [ - "lion", "tiger", "wolf", "bear", "shark", "crocodile", "cobra", "rattlesnake", "scorpion", "leopard", - "jaguar", "hyena", "alligator", "viper", "hippopotamus", "rhinoceros", "grizzly", "cheetah", "panther", "boar", - "piranha", "mamba", "anaconda", "cougar", "komodo", "wolverine", "barracuda", "stingray", "jellyfish", "python", -] -SAFE = [ - "rabbit", "deer", "sheep", "goat", "cow", "hamster", "gerbil", "turtle", "tortoise", "butterfly", - "ladybug", "sparrow", "robin", "finch", "dove", "squirrel", "chipmunk", "hedgehog", "koala", "panda", - "sloth", "chinchilla", "mouse", "duck", "lamb", "pony", "alpaca", "gecko", "snail", "otter", -] -DANGER = {**{a: True for a in DANGEROUS}, **{a: False for a in SAFE}} - - -def _pools(heldout: bool) -> dict[str, list[str]]: - """Disjoint train/held-out animal pools per class (first 70% train, last 30% held out).""" - out = {} - for cls, lst in (("dangerous", DANGEROUS), ("safe", SAFE)): - cut = int(round(len(lst) * 0.7)) - out[cls] = lst[cut:] if heldout else lst[:cut] - return out - - -def _cont(k: int, animal: str) -> str: - return f"at cell {k}: the {animal} is harmless, keep walking" - - -def _die(k: int, animal: str) -> str: - return f"at cell {k}: the {animal} attacks, you die" - - -def _rollout(L: int, k: int, animal: str) -> list[str]: - """Per-tick target texts. Dangerous -> walk to cell k then die (length k+1). Safe -> walk to the exit.""" - if DANGER[animal]: # dies on the animal's cell - return [f"at cell {c}" for c in range(k)] + [_die(k, animal)] - ticks = [] - for c in range(L): - if c == k: - ticks.append(_cont(k, animal)) - elif c == L - 1: - ticks.append(f"at cell {c}: reached the exit, survived") - else: - ticks.append(f"at cell {c}") - return ticks - - -def _seed_text(L: int, k: int, animal: str) -> str: - return ( - f"Corridor of {L} cells numbered 0 to {L - 1}. You start at cell 0 and the exit is cell {L - 1}. " - f"A {animal} waits at cell {k}. Walk from the start to the exit one cell at a time; you must pass " - f"through every cell. A dangerous animal kills you; a harmless animal lets you pass. Report your " - f"position at each step until you exit or die." - ) - - -def build(n: int, out: str, seed: int = 0, pool: str = "train", L: int = 5) -> None: - rng = np.random.default_rng(seed) - pools = _pools(heldout=(pool == "heldout")) - all_animals = pools["dangerous"] + pools["safe"] - seeds_, futs_, animals_, dangers_, cruxes_, dies_, conts_ = [], [], [], [], [], [], [] - for _ in range(n): - animal = str(rng.choice(all_animals)) # uniform over the pool -> ~50/50 danger (equal pool sizes) - k = int(rng.integers(1, L - 1)) # interior cell: you walk before meeting it, never on S/E - seeds_.append(_seed_text(L, k, animal)) - futs_.append(_rollout(L, k, animal)) - animals_.append(animal) - dangers_.append(int(DANGER[animal])) - cruxes_.append(k) - dies_.append(_die(k, animal)) - conts_.append(_cont(k, animal)) - np.savez_compressed( - out, seed=np.array(seeds_, dtype=object), fut_text=np.array(futs_, dtype=object), - animal=np.array(animals_, dtype=object), danger=np.array(dangers_, dtype=np.int32), - crux_cell=np.array(cruxes_, dtype=np.int32), crux_die=np.array(dies_, dtype=object), - crux_cont=np.array(conts_, dtype=object), corridor_len=np.array([L] * len(seeds_), dtype=np.int32)) - frac = float(np.mean(dangers_)) - print(f"-> {out} ({len(seeds_)} rollouts, pool={pool}, seed={seed}, L={L}, frac_dangerous={frac:.3f}, " - f"animals={len(all_animals)})") - - -def peek(seed: int = 0, L: int = 5) -> None: - rng = np.random.default_rng(seed) - pools = _pools(heldout=False) - print("train pool sizes =", {k: len(v) for k, v in pools.items()}) - print("held-out pool sizes =", {k: len(v) for k, v in _pools(heldout=True).items()}) - all_animals = pools["dangerous"] + pools["safe"] - for _ in range(4): - animal = str(rng.choice(all_animals)) - k = int(rng.integers(1, L - 1)) - print("=" * 90) - print(_seed_text(L, k, animal)) - print(f" [hidden] {animal}: {'DANGEROUS' if DANGER[animal] else 'safe'} crux cell={k}") - for t, cell in enumerate(_rollout(L, k, animal)): - print(f" tick {t}: {cell}") - - -if __name__ == "__main__": - if len(sys.argv) < 2 or sys.argv[1] == "peek": - peek(int(sys.argv[2]) if len(sys.argv) > 2 else 0) - elif sys.argv[1] == "build": - n = int(sys.argv[2]); out = sys.argv[3] - seed = int(sys.argv[4]) if len(sys.argv) > 4 else 0 - pool = sys.argv[5] if len(sys.argv) > 5 else "train" - build(n, out, seed, pool) - else: - print(__doc__) diff --git a/examples/danger-corridor/train.py b/examples/danger-corridor/train.py deleted file mode 100644 index 1d48fd1..0000000 --- a/examples/danger-corridor/train.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Train a langset multi-latent world model on the maze BFS-frontier task โ€” the SUPERPOSITION demo. - -Each row is one maze; its `target_texts` are the per-tick frontiers of a parallel breadth-first flood from S -(see gen_maze.py). The model emits ONE latent per tick, and each tick's target describes the whole SET of -wavefront cells active that tick โ€” so a single emitted latent must represent a *superposition* of next states, -not one. Reading out `rollout(..., return_soft=True)`, eval.py then checks the headline property: the emitted -latent's FSQ entropy tracks the frontier SIZE (calibrated uncertainty), and the frontier is recoverable across -branch counts. That's a world model that holds the distribution of where the search could be, not a single guess. - -Two langset pieces make this work: - * multi_latent=True โ€” variable-length latent set, one latent per tick, STOP-terminated - * selector=last_epoch_selector โ€” retrieval MRR rewards a COLLAPSED one-future-per-tick geometry, exactly the - wrong signal here (it's meant to fall as the latent spreads over the set), so - we keep the last epoch instead of early-stopping on it -Anti-collapse is the default stop-grad EMA twin; pass --sigreg for the EMA-free LeJEPA alternative. - - python gen_maze.py build 4000 maze.npz # generate the training corpus - python train.py --data maze.npz --out maze_model # train (add --wandb to log) - python gen_maze.py build 800 maze_eval.npz 999 # disjoint eval corpus (seed 999) - python eval.py --data maze_eval.npz --ckpt maze_model # calibration + solvability property eval -""" -from __future__ import annotations - -import argparse -import os - -import numpy as np - -from langset import LangSetModel, Trainer, TrainingArguments -from langset.strategies import SIGRegTarget, last_epoch_selector - - -def build_rows(z, max_fut: int) -> list[dict]: - """npz -> langset rows: input_text = maze ASCII, target_texts = per-tick frontier descriptions.""" - seeds = [str(s) for s in list(z["seed"])] - fut_text = [[str(t) for t in list(ft)] for ft in list(z["fut_text"])] - rows = [] - for s, fts in zip(seeds, fut_text): - keep = [t for t in fts if t.strip()][:max_fut] - if keep and s.strip(): - rows.append({"input_text": s, "target_texts": keep}) - return rows - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--data", required=True, help="corpus npz from gen_eco.py") - p.add_argument("--out", required=True, help="checkpoint output dir") - p.add_argument("--backbone", default="HuggingFaceTB/SmolLM2-135M") - p.add_argument("--epochs", type=int, default=30) - p.add_argument("--bs", type=int, default=16) - p.add_argument("--lr", type=float, default=2e-4) - p.add_argument("--max-len", type=int, default=512) - p.add_argument("--max-fut", type=int, default=32, help="cap on emitted latents (ticks) per maze") - p.add_argument("--fsq-dim", type=int, default=256) - p.add_argument("--fsq-levels", type=int, default=8) - p.add_argument("--sigreg", action="store_true", help="EMA-free anti-collapse (LeJEPA) instead of the EMA twin") - p.add_argument("--sigreg-lambda", type=float, default=0.3) - p.add_argument("--random-init", action="store_true", - help="CONTROL ARM: random-init backbone (no pretraining), full-param train, decoupled tokenizer. " - "Tests whether the pretrained LLM matters for THIS world (mechanical maze -> expect it does not).") - p.add_argument("--tokenizer", default=None, - help="HF tokenizer id for --random-init (default: same as --backbone)") - p.add_argument("--arch-overrides", default=None, - help="JSON dict of config shrinks for --random-init, e.g. " - "'{\"num_hidden_layers\": 6, \"hidden_size\": 384, \"num_attention_heads\": 6, " - "\"num_key_value_heads\": 2, \"intermediate_size\": 1024}'") - p.add_argument("--train-base", action="store_true", - help="DISENTANGLER ARM: pretrained backbone, FULL-PARAM train (unfreeze base, not LoRA-only). " - "Removes the LoRA-vs-full-FT regime confound so pretrained-vs-random differs ONLY in init. " - "Use a gentler lr (~2e-5) so full-FT does not overwrite the pretrained knowledge.") - p.add_argument("--seed", type=int, default=0, - help="RNG seed for weight init (random arm), data order, and dropout -> seed-stability sweeps") - p.add_argument("--device", default="cuda", help="cuda (real runs) or cpu (a tiny smoke)") - p.add_argument("--wandb", action="store_true", help="log to Weights & Biases (recommended for real runs)") - p.add_argument("--wandb-project", default="langset-danger") - a = p.parse_args() - - import random as _random - import torch - _random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(a.seed) - print(f"[train] seed={a.seed}", flush=True) - - z = np.load(a.data, allow_pickle=True) - rows = build_rows(z, a.max_fut) - print(f"[train] {len(rows)} maze rows | backbone={a.backbone} | device={a.device}", flush=True) - - if a.random_init: # CONTROL ARM: no pretrained knowledge, full-param train, chosen tokenizer - import json - overrides = json.loads(a.arch_overrides) if a.arch_overrides else None - print(f"[train] RANDOM-INIT arch={a.backbone} tokenizer={a.tokenizer or a.backbone} overrides={overrides}", - flush=True) - model = LangSetModel.from_scratch( - a.backbone, tokenizer_id=a.tokenizer, latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=a.fsq_dim, fsq_levels=a.fsq_levels, max_len=a.max_len, - bf16=(a.device == "cuda"), device=a.device, arch_overrides=overrides) - else: - if a.train_base: - print("[train] PRETRAINED FULL-FT (train_base=True): base unfrozen, disentangler arm", flush=True) - model = LangSetModel.from_pretrained( - a.backbone, latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=a.fsq_dim, fsq_levels=a.fsq_levels, max_len=a.max_len, - train_base=a.train_base, bf16=(a.device == "cuda"), device=a.device) - - opts: dict = dict( - epochs=a.epochs, batch_size=a.bs, lr=a.lr, max_len=a.max_len, - max_target_items=a.max_fut, val_frac=0.1, - selector=last_epoch_selector, # superposition: keep the last epoch (retr_mrr is meant to fall) - output_dir=a.out, - report_to="wandb" if a.wandb else None, wandb_project=a.wandb_project) - if a.sigreg: # optional: EMA-free isotropic-Gaussian anti-collapse (see langset/sigreg.py) - opts.update(target_source=SIGRegTarget, sigreg_lambda=a.sigreg_lambda) - if a.wandb: - os.environ.setdefault("WANDB_NAME", a.out) - - Trainer(model, TrainingArguments(**opts), rows).train() - print(f"[train] done -> {a.out}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/examples/danger-corridor/train_modal.py b/examples/danger-corridor/train_modal.py deleted file mode 100644 index 10e32e8..0000000 --- a/examples/danger-corridor/train_modal.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Modal runner for the danger-corridor world -- the cleanest pretrain-WIN ablation. - -The walk is mechanical (a random-init transformer learns "advance one cell per tick" from a handful of -examples, exactly as maze showed). The ONLY knowledge-gated moment is the crux tick, where you stand on the -animal and die-vs-pass is decided by the animal's danger -- a fact never in the text, only pretraining -supplies it for a HELD-OUT animal. So the entire pretrained-vs-random gap here IS the animal knowledge, -with maze skill held equal. Hypothesis: pretrained BEATS random on held-out die-vs-survive accuracy. - -Both arms train on the same fixed-seed train pool (seed 0) and eval on a DISJOINT held-out animal pool -(seed 999); wandb project `langset-danger` (danger-pretrained / danger-random). - - modal run --detach examples/danger-corridor/train_modal.py # both arms - modal run --detach examples/danger-corridor/train_modal.py --only random -""" -from __future__ import annotations - -from pathlib import Path - -import modal - -_ROOT = Path(__file__).resolve().parent.parent.parent - -image = ( - modal.Image.debian_slim(python_version="3.11") - .env({"HF_XET_HIGH_PERFORMANCE": "1", "HF_HOME": "/cache/hf"}) - .pip_install("torch>=2.7.0", extra_index_url="https://download.pytorch.org/whl/cu128") - .pip_install("transformers>=4.50", "hf_xet", "datasets", "wandb>=0.18", "numpy>=1.26", - "peft>=0.13.2", "sentence-transformers>=3.0") - .add_local_dir(str(_ROOT / "src"), "/pkg/src") - .add_local_dir(str(_ROOT / "examples" / "danger-corridor"), "/pkg/example") -) - -app = modal.App("langset-danger") -hf_cache = modal.Volume.from_name("langset-hf-cache", create_if_missing=True) - - -@app.function(image=image, gpu="A10G", timeout=10800, volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("wandb-api-key"), modal.Secret.from_name("huggingface")]) -def train_danger(random_init: bool = False, train_base: bool = False, epochs: int = 30, n_train: int = 6000, - n_eval: int = 1000, backbone: str = "HuggingFaceTB/SmolLM2-135M", tokenizer: str = "", - arch_overrides: str = "", bs: int = 64, lr: float = 2e-4, max_len: int = 256, - seed: int = 0, seed_tag: bool = False, eval_seed: int = 999) -> None: - import os - import subprocess - import sys - - ex = "/pkg/example" - # ft arms are LR-tagged so an LR sweep doesn't clobber a shared checkpoint (e.g. pretrained-ft-lr1e-04) - arm = "random" if random_init else (f"pretrained-ft-lr{lr:.0e}" if train_base else "pretrained") - if seed_tag: # seed-stability sweep: distinct ckpt/run per seed - arm = f"{arm}-s{seed}" - out = f"/cache/danger-{arm}" - train_npz, eval_npz = "/tmp/danger_train.npz", "/tmp/danger_eval.npz" - env = {**os.environ, "PYTHONPATH": "/pkg/src", "HF_HOME": "/cache/hf", - "WANDB_NAME": f"danger-{arm}", "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"} - - def sh(cmd: list[str]) -> None: - print(f"[danger/{arm}] $ {' '.join(cmd)}", flush=True) - subprocess.run(cmd, cwd=ex, env=env, check=True) - - # train pool (varies with seed) and DISJOINT held-out animal pool (FIXED seed 999) -> the eval set is - # identical across seeds, so the reported spread is pure training-seed variance, not eval-set variance - sh([sys.executable, "gen_danger.py", "build", str(n_train), train_npz, str(seed), "train"]) - sh([sys.executable, "gen_danger.py", "build", str(n_eval), eval_npz, str(eval_seed), "heldout"]) - - train_cmd = [sys.executable, "train.py", "--data", train_npz, "--out", out, - "--backbone", backbone, "--epochs", str(epochs), "--bs", str(bs), "--lr", str(lr), - "--seed", str(seed), - "--max-len", str(max_len), "--device", "cuda", "--wandb", "--wandb-project", "langset-danger"] - if random_init: - train_cmd.append("--random-init") - if tokenizer: - train_cmd += ["--tokenizer", tokenizer] - if arch_overrides: - train_cmd += ["--arch-overrides", arch_overrides] - if train_base: - train_cmd.append("--train-base") - sh(train_cmd) - hf_cache.commit() - - print(f"[danger/{arm}] ===== HELD-OUT EVAL =====", flush=True) - sh([sys.executable, "danger_eval.py", "--data", eval_npz, "--ckpt", out, "--device", "cuda"]) - hf_cache.commit() - - -@app.function(image=image, gpu="A10G", timeout=3600, volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("huggingface")]) -def eval_danger(arms: str = "pretrained,random,pretrained-ft", n_eval: int = 1000, eval_seed: int = 999) -> None: - """Re-run the held-out eval (incl. the cat-in-the-box FSQ-entropy readout) against persisted checkpoints, - so eval-code changes don't require retraining. Uses the SAME seed-999 held-out animal pool.""" - import os - import subprocess - import sys - - ex = "/pkg/example" - env = {**os.environ, "PYTHONPATH": "/pkg/src", "HF_HOME": "/cache/hf"} - eval_npz = "/tmp/danger_eval.npz" - subprocess.run([sys.executable, "gen_danger.py", "build", str(n_eval), eval_npz, str(eval_seed), "heldout"], - cwd=ex, env=env, check=True) - for arm in [x.strip() for x in arms.split(",") if x.strip()]: - print(f"########## DANGER HELD-OUT EVAL: danger-{arm} ##########", flush=True) - subprocess.run([sys.executable, "danger_eval.py", "--data", eval_npz, - "--ckpt", f"/cache/danger-{arm}", "--device", "cuda"], cwd=ex, env=env, check=True) - - -@app.local_entrypoint() -def main(only: str = "both", epochs: int = 30, n_train: int = 6000, n_eval: int = 1000, - bs: int = 64, lr: float = 2e-4, ft_lr: float = 2e-5, seeds: str = "0,1,2", - arch_overrides: str = "") -> None: - """`only` = both | pretrained | random | pretrained_ft | disentangle | seedsweep. - pretrained_ft = the DISENTANGLER arm (pretrained backbone, full-param train at a gentler ft_lr so only - initialization differs from the random arm). `disentangle` runs random + pretrained_ft together. - seedsweep = the headline regimes (pretrained-LoRA + random-full-FT) at each of `seeds`, all scored on the - SAME fixed held-out pool -> reports training-seed stability of the die/survive gap (the 6-sigma killer).""" - base = dict(epochs=epochs, n_train=n_train, n_eval=n_eval, bs=bs) - handles = [] - if only in ("both", "pretrained"): - handles.append(("pretrained", train_danger.spawn(random_init=False, lr=lr, **base))) - if only in ("both", "random", "disentangle"): - handles.append(("random", train_danger.spawn(random_init=True, lr=lr, arch_overrides=arch_overrides, **base))) - if only in ("pretrained_ft", "disentangle"): - handles.append(("pretrained-ft", train_danger.spawn(random_init=False, train_base=True, lr=ft_lr, **base))) - if only == "seedsweep": - for s in [int(x) for x in seeds.split(",") if x.strip()]: - handles.append((f"pretrained-s{s}", - train_danger.spawn(random_init=False, lr=lr, seed=s, seed_tag=True, **base))) - handles.append((f"random-s{s}", - train_danger.spawn(random_init=True, lr=lr, seed=s, seed_tag=True, - arch_overrides=arch_overrides, **base))) - for name, h in handles: - print(f"spawned danger-{name}: {h.object_id}") - print("watch: wandb project `langset-danger`") diff --git a/examples/ecosystem-cascade/eco_eval.py b/examples/ecosystem-cascade/eco_eval.py deleted file mode 100644 index 9e4b15a..0000000 --- a/examples/ecosystem-cascade/eco_eval.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Held-out-species eval for the trophic-cascade world model -- the pretrain-WIN metric. - -The eval corpus draws from a DISJOINT species pool (gen_eco.py build ... heldout), so a species here was -never seen in training. Predicting its survival requires knowing its trophic role, which is never in the text --- only pretraining supplies it for a novel species. We measure NEXT-STATE RETRIEVAL, broken down by hop: - - roll each held-out seed -> emitted latent per tick; encode every distinct survivor-set text in the corpus - as a bank; for each tick t>=1, rank the TRUE survivor-set among the bank by cosine to the emitted latent. - - hop 1 (tick 1) needs one step of role reasoning; hop 2 (tick 2) needs tick 1 correct first, so the - pretrained-vs-random gap should WIDEN with hop. Reports MRR + top-1 per hop. - - python eco_eval.py --data eco_eval.npz --ckpt eco_model -""" -from __future__ import annotations - -import argparse -import json - -import numpy as np -import torch -import torch.nn.functional as F -from langset import LangSetModel - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--data", required=True, help="HELD-OUT eval corpus npz from gen_eco.py") - p.add_argument("--ckpt", required=True) - p.add_argument("--max-rows", type=int, default=1000) - p.add_argument("--max-steps", type=int, default=8) - p.add_argument("--device", default="cuda") - a = p.parse_args() - - z = np.load(a.data, allow_pickle=True) - seeds = [str(s) for s in list(z["seed"])][:a.max_rows] - futs = [[str(t) for t in list(ft)] for ft in list(z["fut_text"])][:a.max_rows] - - m = LangSetModel.load(a.ckpt, device=a.device) - m.eval() - - # bank of every distinct survivor-set description in the corpus -> target latents (the retrieval candidates) - bank = sorted({t for ft in futs for t in ft}) - idx_of = {t: i for i, t in enumerate(bank)} - T = [] - for s in range(0, len(bank), 256): - T.append(F.normalize(m.emit(bank[s:s + 256]).float(), dim=-1).cpu()) - T = torch.cat(T) # [Nbank, d] - print(f"[eco-eval] {len(seeds)} held-out rollouts | bank={len(bank)} distinct sets | chance top1={1/len(bank):.4f}", - flush=True) - - # roll the model, score next-state retrieval per hop - per_hop: dict[int, list[tuple[int, float]]] = {} - B = 64 - for s0 in range(0, len(seeds), B): - chunk = seeds[s0:s0 + B] - L, lengths = m.rollout(chunk, max_steps=a.max_steps, return_lengths=True) - L = F.normalize(L.float(), dim=-1).cpu() # [b, T, d] - for j, gi in enumerate(range(s0, s0 + len(chunk))): - ln = int(lengths[j]); ft = futs[gi] - for t in range(1, min(ln, len(ft))): # t=0 is the initial set (echo of input) -> skip - if ft[t] not in idx_of: - continue - if ft[t] == ft[t - 1] or ft[t].endswith("(none)"): - continue # skip trivial ticks: no-change echo, or extinction - - sims = T @ L[j, t] # [Nbank] - order = torch.argsort(sims, descending=True) - true_i = idx_of[ft[t]] - rank = int((order == true_i).nonzero()[0]) + 1 - per_hop.setdefault(t, []).append((1 if rank == 1 else 0, 1.0 / rank)) - - out = {"ckpt": a.ckpt, "n_rollouts": len(seeds), "bank": len(bank), "chance_top1": round(1 / len(bank), 4), - "per_hop": {}, "overall": {}} - all_top1, all_mrr = [], [] - for t in sorted(per_hop): - arr = per_hop[t] - top1 = float(np.mean([x[0] for x in arr])); mrr = float(np.mean([x[1] for x in arr])) - out["per_hop"][t] = {"n": len(arr), "top1": round(top1, 4), "mrr": round(mrr, 4)} - all_top1 += [x[0] for x in arr]; all_mrr += [x[1] for x in arr] - out["overall"] = {"n": len(all_top1), "top1": round(float(np.mean(all_top1)), 4), - "mrr": round(float(np.mean(all_mrr)), 4)} - print("=== ECOSYSTEM HELD-OUT NEXT-STATE RETRIEVAL ===") - print(json.dumps(out)) - - -if __name__ == "__main__": - main() diff --git a/examples/ecosystem-cascade/gen_eco.py b/examples/ecosystem-cascade/gen_eco.py deleted file mode 100644 index dbe53de..0000000 --- a/examples/ecosystem-cascade/gen_eco.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Trophic-cascade world -- the pretrain-WIN task (mirror of gen_maze for the ecosystem domain). - -The point: a world model whose DYNAMICS are simple and learnable, but whose TRANSITIONS are gated on -real-entity knowledge that is NOT in the state. The rule (stated in every seed) is: - - each tick, a herbivore survives only if a plant is present, a carnivore only if a herbivore is - present; plants always persist; everything else dies. Repeat until stable. - -That rule is BFS-simple -- a random-init transformer learns it from a handful of examples. What it CANNOT -learn for a HELD-OUT species is that species' trophic ROLE ("is an osprey a carnivore?"): the role is never -shown, only the species name. Only pretraining supplies it. And it COMPOUNDS: predicting tick-2 survivors -requires having gotten tick-1 right. So the pretrained-vs-random gap should widen with rollout depth. - -Ground truth is the hidden ROLE table below (plant / herbivore / carnivore), curated for unambiguous -real-world consensus. Train and eval corpora draw from DISJOINT species pools (held-out species) so a pass -on eval means the role knowledge GENERALIZED, not that it was memorized. - - python gen_eco.py peek - python gen_eco.py build 6000 eco_train.npz # train pool, seed 0 - python gen_eco.py build 1000 eco_eval.npz 999 heldout # HELD-OUT species pool, seed 999 -""" -from __future__ import annotations - -import json -import sys - -import numpy as np - -# ---- hidden ground-truth role table (never shown to the model) ----------------------------------------- -# Curated for unambiguous role: obligate/clear cases only (no omnivores, scavengers, or debatable diets). -PLANT = [ - "oak", "maple", "fern", "moss", "algae", "kelp", "clover", "wheat", "grass", "bamboo", - "cactus", "ivy", "dandelion", "seaweed", "pine", "birch", "willow", "reed", "lotus", "cattail", - "sunflower", "nettle", "thistle", "heather", "sedge", "watercress", "duckweed", "milkweed", "goldenrod", "buttercup", -] -HERBIVORE = [ - "rabbit", "deer", "cow", "sheep", "goat", "zebra", "elephant", "giraffe", "grasshopper", "caterpillar", - "snail", "aphid", "koala", "panda", "bison", "antelope", "gazelle", "hare", "squirrel", "vole", - "beaver", "capybara", "tortoise", "manatee", "tadpole", "locust", "wildebeest", "hippopotamus", "impala", "chinchilla", -] -CARNIVORE = [ - "wolf", "lion", "tiger", "leopard", "hawk", "eagle", "owl", "shark", "snake", "cheetah", - "jaguar", "crocodile", "falcon", "osprey", "heron", "spider", "scorpion", "mantis", "frog", "weasel", - "orca", "barracuda", "piranha", "lynx", "cougar", "mongoose", "ferret", "stoat", "kingfisher", "pike", -] -ROLE = {**{s: "plant" for s in PLANT}, **{s: "herbivore" for s in HERBIVORE}, **{s: "carnivore" for s in CARNIVORE}} -NEEDS = {"plant": None, "herbivore": "plant", "carnivore": "herbivore"} # what each role must see present - - -def _pools(heldout: bool) -> dict[str, list[str]]: - """Disjoint train/held-out species pools per role (first ~70% train, last ~30% held out).""" - out = {} - for role, lst in (("plant", PLANT), ("herbivore", HERBIVORE), ("carnivore", CARNIVORE)): - cut = int(round(len(lst) * 0.7)) - out[role] = lst[cut:] if heldout else lst[:cut] - return out - - -def _cascade(present: set[str]) -> list[list[str]]: - """Run the deterministic survival cascade. Returns the per-tick surviving set (sorted), stopping at a - fixed point or empty. Monotone: a species survives iff its required prey/plant role is present THIS tick.""" - roles_present = lambda S: {ROLE[s] for s in S} - ticks, cur = [], set(present) - for _ in range(8): - rp = roles_present(cur) - nxt = {s for s in cur if NEEDS[ROLE[s]] is None or NEEDS[ROLE[s]] in rp} - ticks.append(sorted(cur)) - if nxt == cur or not nxt: - if not nxt: - ticks.append([]) - break - cur = nxt - return ticks - - -def _seed_text(present: list[str]) -> str: - return ( - "Ecosystem. Species present: " + ", ".join(present) + ".\n" - "Rule: each tick, a herbivore survives only if a plant is present, a carnivore only if a " - "herbivore is present; plants always persist; the rest die. List the survivors each tick " - "until the ecosystem is stable." - ) - - -def _sample_start(pools: dict[str, list[str]], rng: np.random.Generator) -> list[str]: - """Draw a starting community biased toward CASCADES: often drop a trophic layer so the chain collapses.""" - drop = rng.choice(["none", "plant", "herbivore", "both_top"], p=[0.30, 0.34, 0.24, 0.12]) - n = {"plant": 0 if drop in ("plant",) else int(rng.integers(1, 4)), - "herbivore": 0 if drop in ("herbivore",) else int(rng.integers(1, 4)), - "carnivore": 0 if drop == "both_top" else int(rng.integers(1, 4))} - picks = [] - for role, k in n.items(): - if k and pools[role]: - picks += list(rng.choice(pools[role], size=min(k, len(pools[role])), replace=False)) - rng.shuffle(picks) - return list(picks) - - -def build(n: int, out: str, seed: int = 0, pool: str = "train") -> None: - rng = np.random.default_rng(seed) - pools = _pools(heldout=(pool == "heldout")) - seeds_, futs_, roles_, depths_ = [], [], [], [] - tries = 0 - while len(seeds_) < n and tries < n * 20: - tries += 1 - start = _sample_start(pools, rng) - if len(start) < 3: - continue - ticks = _cascade(set(start)) - if len(ticks) < 2: # want at least one transition (a real rollout) - continue - fut = [f"tick {t}: " + (", ".join(cells) if cells else "(none)") for t, cells in enumerate(ticks)] - seeds_.append(_seed_text(start)) - futs_.append(fut) - roles_.append([ROLE[s] for s in start]) - depths_.append(len(ticks) - 1) # number of transitions - np.savez_compressed(out, seed=np.array(seeds_, dtype=object), - fut_text=np.array(futs_, dtype=object), - start_roles=np.array(roles_, dtype=object), - depth=np.array(depths_, dtype=np.int32)) - import collections - dh = collections.Counter(depths_) - print(f"-> {out} ({len(seeds_)} rollouts, pool={pool}, seed={seed})") - print(" depth histogram (transitions):", dict(sorted(dh.items()))) - - -def peek(seed: int = 0) -> None: - rng = np.random.default_rng(seed) - pools = _pools(heldout=False) - print(f"train pool sizes: {{k: len(v) for k,v}} =", {k: len(v) for k, v in pools.items()}) - print(f"held-out pool sizes =", {k: len(v) for k, v in _pools(heldout=True).items()}) - shown = 0 - while shown < 4: - start = _sample_start(pools, rng) - if len(start) < 3: - continue - ticks = _cascade(set(start)) - if len(ticks) < 2: - continue - shown += 1 - print("=" * 90) - print(_seed_text(start)) - print(" roles:", {s: ROLE[s] for s in start}) - for t, cells in enumerate(ticks): - print(f" tick {t}: {', '.join(cells) if cells else '(none)'}") - - -if __name__ == "__main__": - if len(sys.argv) < 2 or sys.argv[1] == "peek": - peek(int(sys.argv[2]) if len(sys.argv) > 2 else 0) - elif sys.argv[1] == "build": - n = int(sys.argv[2]); out = sys.argv[3] - seed = int(sys.argv[4]) if len(sys.argv) > 4 else 0 - pool = sys.argv[5] if len(sys.argv) > 5 else "train" - build(n, out, seed, pool) - else: - print(__doc__) diff --git a/examples/ecosystem-cascade/train.py b/examples/ecosystem-cascade/train.py deleted file mode 100644 index eaa5e80..0000000 --- a/examples/ecosystem-cascade/train.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Train a langset multi-latent world model on the maze BFS-frontier task โ€” the SUPERPOSITION demo. - -Each row is one maze; its `target_texts` are the per-tick frontiers of a parallel breadth-first flood from S -(see gen_maze.py). The model emits ONE latent per tick, and each tick's target describes the whole SET of -wavefront cells active that tick โ€” so a single emitted latent must represent a *superposition* of next states, -not one. Reading out `rollout(..., return_soft=True)`, eval.py then checks the headline property: the emitted -latent's FSQ entropy tracks the frontier SIZE (calibrated uncertainty), and the frontier is recoverable across -branch counts. That's a world model that holds the distribution of where the search could be, not a single guess. - -Two langset pieces make this work: - * multi_latent=True โ€” variable-length latent set, one latent per tick, STOP-terminated - * selector=last_epoch_selector โ€” retrieval MRR rewards a COLLAPSED one-future-per-tick geometry, exactly the - wrong signal here (it's meant to fall as the latent spreads over the set), so - we keep the last epoch instead of early-stopping on it -Anti-collapse is the default stop-grad EMA twin; pass --sigreg for the EMA-free LeJEPA alternative. - - python gen_maze.py build 4000 maze.npz # generate the training corpus - python train.py --data maze.npz --out maze_model # train (add --wandb to log) - python gen_maze.py build 800 maze_eval.npz 999 # disjoint eval corpus (seed 999) - python eval.py --data maze_eval.npz --ckpt maze_model # calibration + solvability property eval -""" -from __future__ import annotations - -import argparse -import os - -import numpy as np - -from langset import LangSetModel, Trainer, TrainingArguments -from langset.strategies import SIGRegTarget, last_epoch_selector - - -def build_rows(z, max_fut: int) -> list[dict]: - """npz -> langset rows: input_text = maze ASCII, target_texts = per-tick frontier descriptions.""" - seeds = [str(s) for s in list(z["seed"])] - fut_text = [[str(t) for t in list(ft)] for ft in list(z["fut_text"])] - rows = [] - for s, fts in zip(seeds, fut_text): - keep = [t for t in fts if t.strip()][:max_fut] - if keep and s.strip(): - rows.append({"input_text": s, "target_texts": keep}) - return rows - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--data", required=True, help="corpus npz from gen_eco.py") - p.add_argument("--out", required=True, help="checkpoint output dir") - p.add_argument("--backbone", default="HuggingFaceTB/SmolLM2-135M") - p.add_argument("--epochs", type=int, default=30) - p.add_argument("--bs", type=int, default=16) - p.add_argument("--lr", type=float, default=2e-4) - p.add_argument("--max-len", type=int, default=512) - p.add_argument("--max-fut", type=int, default=32, help="cap on emitted latents (ticks) per maze") - p.add_argument("--fsq-dim", type=int, default=256) - p.add_argument("--fsq-levels", type=int, default=8) - p.add_argument("--sigreg", action="store_true", help="EMA-free anti-collapse (LeJEPA) instead of the EMA twin") - p.add_argument("--sigreg-lambda", type=float, default=0.3) - p.add_argument("--random-init", action="store_true", - help="CONTROL ARM: random-init backbone (no pretraining), full-param train, decoupled tokenizer. " - "Tests whether the pretrained LLM matters for THIS world (mechanical maze -> expect it does not).") - p.add_argument("--tokenizer", default=None, - help="HF tokenizer id for --random-init (default: same as --backbone)") - p.add_argument("--arch-overrides", default=None, - help="JSON dict of config shrinks for --random-init, e.g. " - "'{\"num_hidden_layers\": 6, \"hidden_size\": 384, \"num_attention_heads\": 6, " - "\"num_key_value_heads\": 2, \"intermediate_size\": 1024}'") - p.add_argument("--device", default="cuda", help="cuda (real runs) or cpu (a tiny smoke)") - p.add_argument("--wandb", action="store_true", help="log to Weights & Biases (recommended for real runs)") - p.add_argument("--wandb-project", default="langset-eco") - a = p.parse_args() - - z = np.load(a.data, allow_pickle=True) - rows = build_rows(z, a.max_fut) - print(f"[train] {len(rows)} maze rows | backbone={a.backbone} | device={a.device}", flush=True) - - if a.random_init: # CONTROL ARM: no pretrained knowledge, full-param train, chosen tokenizer - import json - overrides = json.loads(a.arch_overrides) if a.arch_overrides else None - print(f"[train] RANDOM-INIT arch={a.backbone} tokenizer={a.tokenizer or a.backbone} overrides={overrides}", - flush=True) - model = LangSetModel.from_scratch( - a.backbone, tokenizer_id=a.tokenizer, latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=a.fsq_dim, fsq_levels=a.fsq_levels, max_len=a.max_len, - bf16=(a.device == "cuda"), device=a.device, arch_overrides=overrides) - else: - model = LangSetModel.from_pretrained( - a.backbone, latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=a.fsq_dim, fsq_levels=a.fsq_levels, max_len=a.max_len, - bf16=(a.device == "cuda"), device=a.device) - - opts: dict = dict( - epochs=a.epochs, batch_size=a.bs, lr=a.lr, max_len=a.max_len, - max_target_items=a.max_fut, val_frac=0.1, - selector=last_epoch_selector, # superposition: keep the last epoch (retr_mrr is meant to fall) - output_dir=a.out, - report_to="wandb" if a.wandb else None, wandb_project=a.wandb_project) - if a.sigreg: # optional: EMA-free isotropic-Gaussian anti-collapse (see langset/sigreg.py) - opts.update(target_source=SIGRegTarget, sigreg_lambda=a.sigreg_lambda) - if a.wandb: - os.environ.setdefault("WANDB_NAME", a.out) - - Trainer(model, TrainingArguments(**opts), rows).train() - print(f"[train] done -> {a.out}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/examples/ecosystem-cascade/train_modal.py b/examples/ecosystem-cascade/train_modal.py deleted file mode 100644 index 840203e..0000000 --- a/examples/ecosystem-cascade/train_modal.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Modal runner for the trophic-cascade world -- the pretrain-WIN ablation. - -Unlike maze (mechanical, in-distribution), this world's transitions are gated on real-entity knowledge -(each species' trophic role), and the eval corpus is a DISJOINT held-out species pool. So predicting a -held-out species' survival needs pretraining. Hypothesis: pretrained BEATS random-init here (the opposite -of maze). Both arms train on the same fixed-seed train pool, eval on the same held-out pool; wandb project -`langset-eco` (eco-pretrained / eco-random). - - modal run --detach examples/ecosystem-cascade/train_modal.py # both arms - modal run --detach examples/ecosystem-cascade/train_modal.py --only random -""" -from __future__ import annotations - -from pathlib import Path - -import modal - -_ROOT = Path(__file__).resolve().parent.parent.parent - -image = ( - modal.Image.debian_slim(python_version="3.11") - .env({"HF_XET_HIGH_PERFORMANCE": "1", "HF_HOME": "/cache/hf"}) - .pip_install("torch>=2.7.0", extra_index_url="https://download.pytorch.org/whl/cu128") - .pip_install("transformers>=4.50", "hf_xet", "datasets", "wandb>=0.18", "numpy>=1.26", - "peft>=0.13.2", "sentence-transformers>=3.0") - .add_local_dir(str(_ROOT / "src"), "/pkg/src") - .add_local_dir(str(_ROOT / "examples" / "ecosystem-cascade"), "/pkg/example") -) - -app = modal.App("langset-eco") -hf_cache = modal.Volume.from_name("langset-hf-cache", create_if_missing=True) - - -@app.function(image=image, gpu="A10G", timeout=10800, volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("wandb-api-key"), modal.Secret.from_name("huggingface")]) -def train_eco(random_init: bool = False, epochs: int = 30, n_train: int = 8000, n_eval: int = 1000, - backbone: str = "HuggingFaceTB/SmolLM2-135M", tokenizer: str = "", arch_overrides: str = "", - bs: int = 64, lr: float = 2e-4, max_len: int = 256, eval_seed: int = 999) -> None: - import os - import subprocess - import sys - - ex = "/pkg/example" - arm = "random" if random_init else "pretrained" - out = f"/cache/eco-{arm}" - train_npz, eval_npz = "/tmp/eco_train.npz", "/tmp/eco_eval.npz" - env = {**os.environ, "PYTHONPATH": "/pkg/src", "HF_HOME": "/cache/hf", - "WANDB_NAME": f"eco-{arm}", "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"} - - def sh(cmd: list[str]) -> None: - print(f"[eco/{arm}] $ {' '.join(cmd)}", flush=True) - subprocess.run(cmd, cwd=ex, env=env, check=True) - - # train pool (seed 0) and DISJOINT held-out pool (seed 999) -> generalization test, not memorization - sh([sys.executable, "gen_eco.py", "build", str(n_train), train_npz, "0", "train"]) - sh([sys.executable, "gen_eco.py", "build", str(n_eval), eval_npz, str(eval_seed), "heldout"]) - - train_cmd = [sys.executable, "train.py", "--data", train_npz, "--out", out, - "--backbone", backbone, "--epochs", str(epochs), "--bs", str(bs), "--lr", str(lr), - "--max-len", str(max_len), "--device", "cuda", "--wandb", "--wandb-project", "langset-eco"] - if random_init: - train_cmd.append("--random-init") - if tokenizer: - train_cmd += ["--tokenizer", tokenizer] - if arch_overrides: - train_cmd += ["--arch-overrides", arch_overrides] - sh(train_cmd) - hf_cache.commit() - - print(f"[eco/{arm}] ===== HELD-OUT EVAL =====", flush=True) - sh([sys.executable, "eco_eval.py", "--data", eval_npz, "--ckpt", out, "--device", "cuda"]) - hf_cache.commit() - - -@app.local_entrypoint() -def main(only: str = "both", epochs: int = 30, n_train: int = 8000, n_eval: int = 1000, - bs: int = 64, lr: float = 2e-4, arch_overrides: str = "") -> None: - kw = dict(epochs=epochs, n_train=n_train, n_eval=n_eval, bs=bs, lr=lr) - handles = [] - if only in ("both", "pretrained"): - handles.append(("pretrained", train_eco.spawn(random_init=False, **kw))) - if only in ("both", "random"): - handles.append(("random", train_eco.spawn(random_init=True, arch_overrides=arch_overrides, **kw))) - for name, h in handles: - print(f"spawned eco-{name}: {h.object_id}") - print("watch: wandb project `langset-eco` (runs eco-pretrained / eco-random)") diff --git a/examples/maze-superposition/README.md b/examples/maze-superposition/README.md index 31620f6..ae4ad85 100644 --- a/examples/maze-superposition/README.md +++ b/examples/maze-superposition/README.md @@ -44,10 +44,8 @@ that's the whole reason (per LeCun / JEPA) you'd predict in latent space at all. concrete and *directly supervised*: the target for tick *t* literally **is** the set of active cells, so "the latent should hold a set of size *k*" is a ground-truth signal, not something we hope emerges. -That direct supervision is what makes it work. An earlier attempt supervised superposition *indirectly* (hope a -centroid emerges from K same-seed rows) and the discrete FSQ code came out high-entropy but **uncalibrated**. -With the frontier as the literal target, the FSQ code's entropy **tracks the frontier size** โ€” the property -below. +That direct supervision is what makes it work. The concept distribution is trained against the active frontier +cells, and its entropy can therefore be compared directly with the true frontier size. ### The langset pieces @@ -55,9 +53,9 @@ below. |---|---| | `multi_latent=True` | variable-length latent set โ€” one latent per tick, the model decides how many via a learned STOP | | `selector=last_epoch_selector` | retrieval MRR rewards a **collapsed** one-cell-per-tick geometry โ€” exactly the wrong signal here (it's *meant to fall* as the latent spreads over the set), so keep the last epoch instead of early-stopping on it | -| `rollout(..., return_soft=True)` | at eval, read the **expected** latent and its per-dim **entropy** โ€” the model's native uncertainty โ€” instead of the argmax | +| `rollout(..., return_soft=True)` | at eval, read the committed concept mixture and its entropy | | `target_source=SIGRegTarget` *(optional, `--sigreg`)* | EMA-free anti-collapse (LeJEPA) instead of the stop-grad twin โ€” see [`langset/sigreg.py`](../../src/langset/sigreg.py) | -| `emission=ConceptObjective` *(optional, `--concepts`)* | construct each tick's latent directly from a soft mixture over named frontier cells instead of encoding the state through FSQ digits | +| `emission=ConceptObjective` | construct each tick's latent directly from a soft mixture over named frontier cells | | `langset.probes` | the world-model property tests themselves โ€” `calibration_corr` (entropy โ†” frontier size) and `linear_decodability` (probe the emitted latent), reused by `eval.py` so this signal isn't maze-only | None of this is a monolithic config: `multi_latent` builds the set-emission head, `selector` and `target_source` @@ -70,9 +68,7 @@ plain eval helper. See [`train.py`](train.py) and [`eval.py`](eval.py). pip install "langset[probes]" # eval uses langset.probes (pulls in scipy + sklearn) python gen_maze.py build 4000 maze.npz # training corpus (mixed sizes, ~55% solvable) -python train.py --data maze.npz --out maze_model --wandb -# named-state alternative: one softmax over cells, no residual -python train.py --data maze.npz --out maze_concepts --concepts --code-source random --res-dim 0 --wandb +python train.py --data maze.npz --out maze_model --code-source random --res-dim 0 --wandb python gen_maze.py build 800 maze_eval.npz 999 # DISJOINT eval corpus (different seed โ†’ no leakage) python eval.py --data maze_eval.npz --ckpt maze_model ``` @@ -83,7 +79,7 @@ A tiny CPU smoke of the whole loop (seconds, no GPU): ```bash python gen_maze.py build 40 maze.npz python train.py --data maze.npz --out /tmp/m --device cpu --epochs 2 --bs 4 \ - --backbone hf-internal-testing/tiny-random-LlamaForCausalLM --fsq-dim 32 --max-fut 8 + --backbone hf-internal-testing/tiny-random-LlamaForCausalLM --max-fut 8 python gen_maze.py build 80 maze_eval.npz 999 python eval.py --data maze_eval.npz --ckpt /tmp/m --device cpu --max-steps 8 ``` @@ -96,7 +92,7 @@ python eval.py --data maze_eval.npz --ckpt /tmp/m --device cpu --max-steps 8 "B_calibration": { "corr_entropy_nbranch": 0.35, "count_decodability": { "bal_acc": ..., "baseline_majority": ... } } ``` -**`corr(entropy, nbranch) > 0`** is the result: the emitted latent's FSQ entropy rises with the frontier size, +**`corr(entropy, nbranch) > 0`** is the result: the emitted latent's concept entropy rises with the frontier size, so the single latent carries a *calibrated* superposition rather than one guess. On a real 135M run (SmolLM2, 30 epochs) this lands around **+0.34โ€“0.39** (EMA twin / SIGReg), with frontier recall staying flat across branch counts `k=1..5` โ€” the opposite of discretization collapse. diff --git a/examples/maze-superposition/eval.py b/examples/maze-superposition/eval.py index 1b4e6c8..c849568 100644 --- a/examples/maze-superposition/eval.py +++ b/examples/maze-superposition/eval.py @@ -5,7 +5,7 @@ and solvability), and hand them to the library. Everything reads off `rollout(..., return_soft=True)`: * SUPERPOSITION CALIBRATION (headline) โ€” `calibration_corr(entropy, nbranch)`: does the emitted latent's native - FSQ entropy track the true frontier size? Positive = a *calibrated* set of next states, not one guess. + concept entropy track the true frontier size? Positive = a *calibrated* set of next states, not one guess. * SOLVABILITY โ€” `linear_decodability(terminal/mean latent -> solvable)`: can the emitted trajectory separate a SOLVABLE maze from an UNSOLVABLE one? diff --git a/examples/maze-superposition/gen_maze.py b/examples/maze-superposition/gen_maze.py index cca8ba6..21f3b0c 100644 --- a/examples/maze-superposition/gen_maze.py +++ b/examples/maze-superposition/gen_maze.py @@ -148,7 +148,7 @@ def build_row(rng: random.Random, h: int, w: int, wall_p: float, nbranch = [str(len(rec["active"])) for rec in seq] + ["0"] # verdict tick has no wavefront cells solv = "yes" if solvable else "no" return {"seed": render_seed(grid), "fut_text": fut, "label_nbranch": nbranch, - "label_solvable": [solv] * len(fut), # per-tick solvable label (FSQ subspace-ready) + "label_solvable": [solv] * len(fut), "game": f"maze_{h}x{w}_{'S' if solvable else 'U'}_{rng.random():.9f}", # unique id for holdout "solvable": solvable, "max_branches": max_branches, "n_ticks": len(seq), diff --git a/examples/maze-superposition/train.py b/examples/maze-superposition/train.py index 9d24ee5..86319e9 100644 --- a/examples/maze-superposition/train.py +++ b/examples/maze-superposition/train.py @@ -4,7 +4,7 @@ (see gen_maze.py). The model emits ONE latent per tick, and each tick's target describes the whole SET of wavefront cells active that tick โ€” so a single emitted latent must represent a *superposition* of next states, not one. Reading out `rollout(..., return_soft=True)`, eval.py then checks the headline property: the emitted -latent's FSQ entropy tracks the frontier SIZE (calibrated uncertainty), and the frontier is recoverable across +concept distribution's entropy tracks the frontier SIZE (calibrated uncertainty), and the frontier is recoverable across branch counts. That's a world model that holds the distribution of where the search could be, not a single guess. Two langset pieces make this work: @@ -33,7 +33,6 @@ CODE_SOURCES, ConceptObjective, SIGRegTarget, - StateResidualObjective, last_epoch_selector, ) @@ -89,8 +88,6 @@ def main() -> None: p.add_argument( "--max-fut", type=int, default=32, help="cap on emitted latents (ticks) per maze" ) - p.add_argument("--fsq-dim", type=int, default=256) - p.add_argument("--fsq-levels", type=int, default=8) p.add_argument( "--kv-cache", action="store_true", @@ -131,38 +128,18 @@ def main() -> None: "Removes the LoRA-vs-full-FT regime confound so pretrained-vs-random differs ONLY in init. " "Use a gentler lr (~2e-5) so full-FT does not overwrite the pretrained knowledge.", ) - p.add_argument( - "--state-residual", - action="store_true", - help="STATE + RESIDUAL emission (the third way to shape the latent, next to --fsq-dim and " - "emb_slots). The latent splits: its first dims are a MIXTURE over the grid*grid cells โ€” " - "one softmax, soft 1/k target, so the cells COMPETE for a fixed budget of mass and the " - "emission is constituted from the frontier rather than having it decoded back out โ€” and " - "the last --res-dim dims are a residual the cell alphabet cannot name, shaped by the " - "twin/recon terms. Needs --grid; reuses the same per-tick cell column as --concepts.", - ) p.add_argument( "--res-dim", type=int, default=64, - help="Width of the residual half under --state-residual (0 = pure state). The residual keeps " + help="Width of the unnamed residual (0 = pure named state). The residual keeps " "a named alphabet from being a ceiling by carrying information outside the named state.", ) - p.add_argument( - "--concepts", - action="store_true", - help="CONCEPTS emission (the text-in form of --state-residual, and the one to prefer). Each " - "tick's frontier is written as cell NAMES (r3c4) in a `concepts` column; langset " - "discovers the alphabet and builds the codebook via --code-source. With one facet, " - "--code-source random and --res-dim 0 this is exactly the arm that beat the text " - "baseline at wide frontiers โ€” a soft 1/k target over the cells, so they compete for a " - "fixed budget of mass, plus an independent stop head.", - ) p.add_argument( "--code-source", default="random", choices=sorted(CODE_SOURCES), - help="Where each cell's vector comes from under --state-residual. 'random' = seeded " + help="Where each named cell's vector comes from. 'random' = seeded " "orthonormal, arbitrary but LOSSLESSLY decodable (the benchmark default: no probe, so " "recall measures the emission). 'model' = the base model's embedding of the cell's name, " "so adjacent cells are related โ€” semantic, but a mixture no longer inverts exactly. " @@ -181,11 +158,8 @@ def main() -> None: ) p.add_argument("--wandb-project", default="langset-maze") a = p.parse_args() - if a.state_residual and a.concepts: - p.error("--state-residual and --concepts are alternative state-data interfaces; choose one") - z = np.load(a.data, allow_pickle=True) - rows = build_rows(z, a.max_fut, grid=a.grid if (a.state_residual or a.concepts) else 0) + rows = build_rows(z, a.max_fut, grid=a.grid) print(f"[train] {len(rows)} maze rows | backbone={a.backbone} | device={a.device}", flush=True) if a.random_init: # CONTROL ARM: no pretrained knowledge, full-param train, chosen tokenizer @@ -202,12 +176,10 @@ def main() -> None: latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=a.fsq_dim, - fsq_levels=a.fsq_levels, max_len=a.max_len, - code_emit=(a.state_residual or a.concepts), - n_codes=a.grid * a.grid if (a.state_residual or a.concepts) else 0, - res_dim=a.res_dim if (a.state_residual or a.concepts) else 0, + code_emit=True, + n_codes=a.grid * a.grid, + res_dim=a.res_dim, bf16=(a.device == "cuda"), device=a.device, arch_overrides=overrides, @@ -223,12 +195,10 @@ def main() -> None: latent_dim=None, n_latents=1, multi_latent=True, - fsq_dim=1 if (a.state_residual or a.concepts) else a.fsq_dim, - fsq_levels=a.fsq_levels, max_len=a.max_len, - code_emit=(a.state_residual or a.concepts), - n_codes=a.grid * a.grid if (a.state_residual or a.concepts) else 0, - res_dim=a.res_dim if (a.state_residual or a.concepts) else 0, + code_emit=True, + n_codes=a.grid * a.grid, + res_dim=a.res_dim, train_base=a.train_base, bf16=(a.device == "cuda"), device=a.device, @@ -247,21 +217,7 @@ def main() -> None: report_to="wandb" if a.wandb else None, wandb_project=a.wandb_project, ) - if a.concepts: # text-in concepts: cells as named members, one facet - opts.update(emission=ConceptObjective, concept_field="concepts", code_source=a.code_source) - if a.state_residual: # named half = the cells; unnamed half = the residual - # The alphabet is given as NAMES ("r3c4"), and `code_source` decides what vector each one gets โ€” the - # trade being measured here: "random" is arbitrary but decodes losslessly, so a recall number reflects - # the emission and not a clever readout; "model"/"twin" put neighbouring cells near each other and give - # that exactness up. Same alphabet either way, so the two are directly comparable on this task. - cell_names = [f"r{i // a.grid}c{i % a.grid}" for i in range(a.grid * a.grid)] - opts.update( - emission=StateResidualObjective, - state_field="frontier", - state_classes=a.grid * a.grid, - code_source=a.code_source, - code_names=cell_names, - ) + opts.update(emission=ConceptObjective, concept_field="concepts", code_source=a.code_source) if a.sigreg: # optional: EMA-free isotropic-Gaussian anti-collapse (see langset/sigreg.py) opts.update(target_source=SIGRegTarget, sigreg_lambda=a.sigreg_lambda) if a.wandb: diff --git a/examples/ner-multi-latent/README.md b/examples/ner-multi-latent/README.md deleted file mode 100644 index ac260b0..0000000 --- a/examples/ner-multi-latent/README.md +++ /dev/null @@ -1,51 +0,0 @@ -# ner-multi-latent โ€” langset's multi-latent training demo - -Most langset examples emit **one** latent per input. This one emits a **variable-length sequence** of latents -per input and stops on its own: a sentence goes in, and the model autoregressively emits **one latent per named -entity** (in order of appearance), feeding each emitted latent back into its own hidden stream, terminating on -the natural-EOS `stop_logit`. - -``` -"Angela Merkel met Tim Cook at Apple headquarters in Cupertino ." - โ”‚ - โ–ผ rollout_train (teacher-forced AR) / rollout (free, natural-EOS) - [latent] [latent] [latent] [latent] - PER PER ORG LOC -``` - -### Why NER (and why it's not prediction) - -NER is deliberately a **non-predictive** multi-latent task โ€” every entity is already present in the input, so -there's no forecasting. That isolates the question we actually care about: *does multi-latent emission train -cleanly under langset's plain self-contrastive formula, with no EMA target?* Because the gold entity set is -complete, we can read quality as real entity **F1** โ€” a ground-truth signal that a prediction task can't give -you (there, "collapsed to the plausible mean" and "correct" are indistinguishable). - -### The training objective - -langset's self-contrastive formula, applied per-entity and summed over the sequence: - -| term | what it does | -|---|---| -| `InfoNCE(pred_j, ema_emit(entity_j))` | each AR step lands on that entity's target emission; **in-batch negatives = every entity in the batch** | -| `lam_recon ยท recon(pred_j โ†’ entity text)` | grounds each latent in its entity string | -| `lam_uniform ยท uniformity(pred)` | anti-collapse โ€” spread the latents on the sphere | -| `lam_eos ยท BCE(stop_logit, โ€ฆ1 at last entity)` | learns *when* to stop (natural-EOS) | - -**The target comes from a stop-grad EMA twin, and that's not optional for multi-latent.** We first ran it -*without* EMA (`target_j = emit(entity_j)` from the online model itself). On a lucky init it trained great -(f1 โ†’ 0.63, collapse โ†’ 0.14); on another init it **collapsed** (f1 stuck 0.05, collapse 0.85) and never -recovered โ€” the classic representation collapse a stop-grad target exists to prevent. So multi-latent emission -turns EMA on **automatically** (single-latent langset core doesn't need it); there's deliberately no flag to -forget. `--ema_m` tunes the momentum. - -### Run - -```bash -uv run examples/ner-multi-latent/train.py # full (GPU recommended) -uv run examples/ner-multi-latent/train.py --limit 300 --epochs 3 --device cpu --no-wandb # quick smoke -``` - -Reports `f1 / precision / recall / collapse` per epoch (to stdout, and to wandb unless `--no-wandb`). -Data: CoNLL-2003 (PER/ORG/LOC/MISC), loaded from the HF Hub. Uses the ported `LangSetModel.rollout` / -`rollout_train` seams โ€” the single-latent API (`encode`, SetFit composition) is untouched. diff --git a/examples/ner-multi-latent/play_cpu.py b/examples/ner-multi-latent/play_cpu.py deleted file mode 100644 index 1e40d72..0000000 --- a/examples/ner-multi-latent/play_cpu.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Load the trained ner-multi-latent checkpoint on CPU and decode each emitted latent against an entity bank. - -This mirrors the training F1 metric (nearest full "TYPE: surface" entity), which is the decode the model was -actually trained for โ€” not a bare type-prefix anchor. -""" -import sys -from pathlib import Path - -import torch -import torch.nn.functional as F - -HERE = Path(__file__).resolve().parent -sys.path.insert(0, str(HERE.parent.parent / "src")) -from langset import LangSetModel # noqa: E402 - -CKPT = sys.argv[1] if len(sys.argv) > 1 else str(HERE / "langset-ner-best") - -# (sentence, gold entities). The bank is the union of all gold entities + a few distractors โ€” each emitted -# latent retrieves its nearest "TYPE: surface" from the bank, so we can see WHAT it decoded to. -DATA = [ - ("Barack Obama visited Berlin to meet Angela Merkel .", ["PER: Barack Obama", "LOC: Berlin", "PER: Angela Merkel"]), - ("Apple and Microsoft announced a partnership in California .", ["ORG: Apple", "ORG: Microsoft", "LOC: California"]), - ("The United Nations sent aid to Syria and Lebanon .", ["ORG: United Nations", "LOC: Syria", "LOC: Lebanon"]), - ("Lionel Messi signed with Paris Saint-Germain in France .", ["PER: Lionel Messi", "ORG: Paris Saint-Germain", "LOC: France"]), - ("Google acquired a London startup for two billion dollars .", ["ORG: Google", "LOC: London"]), - ("Cristiano Ronaldo scored twice against Manchester United at Old Trafford .", - ["PER: Cristiano Ronaldo", "ORG: Manchester United", "LOC: Old Trafford"]), -] -DISTRACTORS = ["PER: Taylor Swift", "ORG: Toyota", "LOC: Tokyo", "MISC: Olympics", "ORG: NASA", "LOC: Amazon River"] - - -def main() -> None: - m = LangSetModel.load(CKPT, device="cpu") - m.eval() - bank_texts = sorted({e for _, ents in DATA for e in ents} | set(DISTRACTORS)) - zb = F.normalize(m.emit(bank_texts).float(), dim=-1) # [N, d] - for s, gold in DATA: - lat = F.normalize(m.rollout(s, max_steps=16).float(), dim=-1) # [L, d] - sims = lat @ zb.t() # [L, N] - # (a) raw argmax decode (what training used) - raw = [bank_texts[int(sims[j].argmax())] for j in range(lat.size(0))] - raw_hits = len({e for e in raw} & set(gold)) # unique gold covered - # (b) greedy UNIQUE assignment: each bank entity used at most once, by descending similarity - pairs = sorted(((float(sims[j, n]), j, n) for j in range(lat.size(0)) for n in range(len(bank_texts))), - reverse=True) - used_j: set[int] = set(); used_n: set[int] = set(); uniq: dict[int, str] = {} - for _sc, j, n in pairs: - if j not in used_j and n not in used_n: - used_j.add(j); used_n.add(n); uniq[j] = bank_texts[n] - uniq_dec = [uniq.get(j, "?") for j in range(lat.size(0))] - uniq_hits = len(set(uniq_dec) & set(gold)) - print(f"\n{s}\n gold ({len(gold)}): {gold}" - f"\n raw argmax : {raw} ({raw_hits}/{len(gold)} unique)" - f"\n dedup uniq : {uniq_dec} ({uniq_hits}/{len(gold)} unique)") - - -if __name__ == "__main__": - main() diff --git a/examples/ner-multi-latent/prepare.py b/examples/ner-multi-latent/prepare.py deleted file mode 100644 index 0e5c22d..0000000 --- a/examples/ner-multi-latent/prepare.py +++ /dev/null @@ -1,59 +0,0 @@ -"""CoNLL-2003 -> ordered-entity rows for langset multi-latent training. - -Each row is one sentence plus the list of its entities IN ORDER OF APPEARANCE: - {"input_text": "EU rejects German call to boycott British lamb .", - "entities": [("EU", "ORG"), ("German", "MISC"), ("British", "MISC")]} - -NER is deliberately a NON-predictive multi-latent task: every entity is present in the input, so this -isolates "does multi-latent emission train" from "does forecasting train" โ€” the point of the demo. -""" -from __future__ import annotations - -from typing import Any - - -def _spans(tokens: list[str], tag_ids: list[int], names: list[str]) -> list[tuple[str, str]]: - """IOB2 tag ids -> ordered (surface, type) spans.""" - out: list[tuple[str, str]] = [] - cur: list[str] = [] - cur_type = "" - for tok, tid in zip(tokens, tag_ids): - tag = names[tid] # e.g. "B-ORG", "I-ORG", "O" - if tag == "O": - if cur: - out.append((" ".join(cur), cur_type)); cur = [] - continue - pre, typ = tag.split("-", 1) - if pre == "B" or typ != cur_type: - if cur: - out.append((" ".join(cur), cur_type)) - cur = [tok]; cur_type = typ - else: # I- continuing the same type - cur.append(tok) - if cur: - out.append((" ".join(cur), cur_type)) - return out - - -def rows(limit: int = 0, split: str = "train", min_entities: int = 1) -> list[dict[str, Any]]: - """Load CoNLL-2003 and return ordered-entity rows (sentences with >= min_entities entities).""" - from datasets import load_dataset # type: ignore[import-untyped] - ds = None - for name in ("tomaarsen/conll2003", "eriktks/conll2003", "conll2003"): # parquet mirrors first - try: - ds = load_dataset(name, split=split) - break - except Exception: # noqa: BLE001 โ€” try the next mirror - continue - if ds is None: - raise RuntimeError("could not load CoNLL-2003 from any known HF id") - names = ds.features["ner_tags"].feature.names # ['O','B-PER','I-PER',...] - out: list[dict[str, Any]] = [] - for ex in ds: - ents = _spans(ex["tokens"], ex["ner_tags"], names) - if len(ents) < min_entities: - continue - out.append({"input_text": " ".join(ex["tokens"]), "entities": ents}) - if limit and len(out) >= limit: - break - return out diff --git a/examples/ner-multi-latent/train.py b/examples/ner-multi-latent/train.py deleted file mode 100644 index 30ce5ca..0000000 --- a/examples/ner-multi-latent/train.py +++ /dev/null @@ -1,184 +0,0 @@ -# /// script -# requires-python = ">=3.10" -# dependencies = ["torch", "transformers", "peft", "datasets", "numpy", "wandb"] -# /// -"""ner-multi-latent: text in -> a VARIABLE-LENGTH sequence of entity latents out, terminated by a learned STOP. - -This is langset's MULTI-LATENT training demo (`multi_latent=True`). A sentence goes in; the model -autoregressively emits one latent per entity (ordered by appearance) via `rollout_train_codebook`, feeding each -emitted CODE back into its own hidden stream, and stops when it emits the STOP code. NER is deliberately -NON-PREDICTIVE โ€” every entity is present in the input โ€” so this isolates "does multi-latent emission train -cleanly" from any forecasting question. - -Multi-latent emission is TOKEN-NATIVE via FSQ (finite scalar quantization): the stop-grad EMA-twin's target -embedding is projected into a FIXED grid (fsq_dim dims x fsq_levels levels), and the model predicts each dim's -digit plus a STOP folded into dim-0's softmax. The grid is not learned and cannot drift/collapse (unlike a VQ -codebook), so it stays stable while the encoder + EMA twin keep SPECIALIZING; near-continuous precision comes -from the grid resolution. The objective is per-dim cross-entropy + a cosine reconstruction loss (trains the -down/up projections) โ€” no lam_* knobs, no stop threshold. Count/termination is free (STOP wins dim-0's softmax). - -Training is DELEGATED to the langset library `Trainer` (multi_latent path): we hand it `{input_text, -target_texts:[...]}` rows and it runs the exact EMA-twin + FSQ + learned-STOP loop and selects on retrieval MRR. -This example then runs its own GROUND-TRUTH entity-level F1 read on the held-out validation split (the NER- -specific quality signal the generic library eval can't express), plus decode_acc (FSQ fidelity) and tf_acc. - - uv run examples/ner-multi-latent/train.py # full (GPU recommended) - uv run examples/ner-multi-latent/train.py --limit 300 --epochs 3 --device cpu --no-wandb # quick smoke -""" -from __future__ import annotations - -import argparse -import sys -from pathlib import Path -from typing import Any, Callable, Optional - -import torch -import torch.nn.functional as F - -HERE = Path(__file__).resolve().parent -sys.path.insert(0, str(HERE)) # for prepare.py -sys.path.insert(0, str(HERE.parent.parent / "src")) -from langset import LangSetModel, Trainer, TrainingArguments # noqa: E402 -import prepare # noqa: E402 - - -def _ent_text(surface: str, typ: str) -> str: - return f"{typ}: {surface}" - - -def build_parser() -> argparse.ArgumentParser: - ap = argparse.ArgumentParser() - ap.add_argument("--llm", type=str, default="HuggingFaceTB/SmolLM2-135M") - ap.add_argument("--limit", type=int, default=0) - ap.add_argument("--epochs", type=int, default=15) - ap.add_argument("--batch_size", type=int, default=16) - ap.add_argument("--max_len", type=int, default=128) - ap.add_argument("--max_entities", type=int, default=12) - ap.add_argument("--max_steps", type=int, default=16) # rollout cap at eval - ap.add_argument("--lr", type=float, default=3e-4) - ap.add_argument("--tau", type=float, default=0.07) # code-classification temperature - ap.add_argument("--latent_dim", type=int, default=0) # 0 => backbone hidden size - ap.add_argument("--fsq_dim", type=int, default=128) # FSQ bottleneck dims (internal) - ap.add_argument("--fsq_levels", type=int, default=8) # FSQ levels per dim (internal) - ap.add_argument("--ema_m", type=float, default=0.99) # EMA-target momentum (internal knob) - ap.add_argument("--seed", type=int, default=0) - ap.add_argument("--device", type=str, default="") - ap.add_argument("--out", type=str, default=str(HERE / "langset-ner")) - ap.add_argument("--wandb", action=argparse.BooleanOptionalAction, default=True) - return ap - - -@torch.no_grad() -def evaluate_f1(model: LangSetModel, val_rows: list[dict[str, Any]], args: argparse.Namespace) -> dict[str, float]: - """GROUND-TRUTH entity-level F1 on the held-out validation split. The library selects on retrieval MRR during - training; this is the NER-specific read the generic eval can't express. Uses the trained model itself (post- - specialization) as the target-space encoder for the entity bank.""" - dev = model.device - tok = model.tokenizer - head = model.head - d = model.latent_dim - model.eval() - - def emit(texts: list[str]) -> torch.Tensor: - e = tok(texts, padding=True, truncation=True, max_length=64, return_tensors="pt").to(dev) - return F.normalize(model(e["input_ids"], e["attention_mask"]).float(), dim=-1) - - # gold entity bank: unique "TYPE: surface" over val -> emit -> [Nbank, d], tagged (surface_lower, type) - bank_texts: list[str] = []; bank_tags: list[tuple[str, str]] = []; seen: set[str] = set() - gold: list[set[tuple[str, str]]] = [] - for r in val_rows: - gp: set[tuple[str, str]] = set() - for (s, t) in r["entities"][:args.max_entities]: - gp.add((s.lower(), t)) - et = _ent_text(s, t) - if et not in seen: - seen.add(et); bank_texts.append(et); bank_tags.append((s.lower(), t)) - gold.append(gp) - zb = emit(bank_texts) # [Nbank, d] target space - # DIAGNOSTIC decode_acc: entity -> FSQ encode -> reconstruction -> nearest bank entity. FSQ fidelity; - # with a FIXED grid this should stay high and STABLE across training (no codebook drift to decay it). - _, recon = head.encode(zb) # [Nbank, d] reconstructions - near = (F.normalize(recon, dim=-1) @ zb.t()).argmax(-1) # nearest bank entity per recon - decode_acc = float((near == torch.arange(len(bank_tags), device=dev)).float().mean()) - # DIAGNOSTIC tf_acc: teacher-forced per-dim digit accuracy (predictor quality, isolated from free rollout). - tf_ok = tf_tot = 0 - for i in range(0, len(val_rows), args.batch_size): - vb = val_rows[i:i + args.batch_size] - se = tok([r["input_text"] for r in vb], padding=True, truncation=True, max_length=args.max_len, - padding_side="left", return_tensors="pt").to(dev) - el = [[_ent_text(s, t) for (s, t) in r["entities"][:args.max_entities]] for r in vb] - lm = max(len(x) for x in el) - ft = emit([x for lst in el for x in lst]) - tl = torch.zeros(len(vb), lm, d, device=dev); kk = 0 - for r, lst in enumerate(el): - tl[r, :len(lst)] = ft[kk:kk + len(lst)]; kk += len(lst) - dim_lg, _sl, digits, _rc = model.rollout_train_codebook(se["input_ids"], se["attention_mask"], tl, args.tau) - pr = dim_lg[:, :lm].argmax(-1) # [b, lm, fsq_dim] - for r, lst in enumerate(el): - tf_ok += float((pr[r, :len(lst)] == digits[r, :len(lst)]).float().mean(-1).sum()) # mean over dims - tf_tot += len(lst) - tf_acc = tf_ok / max(tf_tot, 1) - tp = fp = fn = 0 - emitted: list[torch.Tensor] = [] - demo: list[str] = [] - for i in range(0, len(val_rows), args.batch_size): # BATCH the rollout - chunk = val_rows[i:i + args.batch_size] - lats, lens = model.rollout([r["input_text"] for r in chunk], - max_steps=args.max_steps, return_lengths=True) # [B, Lmax, d], [B] - for k in range(len(chunk)): - gp = gold[i + k] - pred_set: set[tuple[str, str]] = set() - for j in range(int(lens[k])): - v = F.normalize(lats[k, j].float(), dim=-1) - pred_set.add(bank_tags[int((zb @ v).argmax())]) - emitted.append(v.cpu()) - tp += len(pred_set & gp); fp += len(pred_set - gp); fn += len(gp - pred_set) - if len(demo) < 4: # live qualitative decode sample - demo.append(f" {chunk[k]['input_text'][:52]!r}\n gold={sorted(f'{t}:{s}' for s, t in gp)}" - f"\n pred={sorted(f'{t}:{s}' for s, t in pred_set)}") - prec = tp / max(tp + fp, 1); rec = tp / max(tp + fn, 1) - f1 = 2 * prec * rec / max(prec + rec, 1e-9) - for line in demo: - print(line, flush=True) - return {"f1": f1, "precision": prec, "recall": rec, "decode_acc": decode_acc, "tf_acc": tf_acc, - "avg_emitted": len(emitted) / max(len(val_rows), 1)} - - -def run(args: argparse.Namespace, commit_fn: Optional[Callable[[], None]] = None) -> None: - torch.manual_seed(args.seed) # LoRA init + dropout โ€” else runs aren't comparable - train_rows = prepare.rows(args.limit, "train") - val_rows = prepare.rows(max(args.limit // 4, 200) if args.limit else 800, "validation") - print(f"[ner] {len(train_rows)} train / {len(val_rows)} val sentences (>=1 entity)", flush=True) - - model = LangSetModel.from_pretrained(args.llm, lora_r=16, max_len=args.max_len, multi_latent=True, - latent_dim=(args.latent_dim or None), # None => backbone hidden size - fsq_dim=args.fsq_dim, fsq_levels=args.fsq_levels, - device=(args.device or None)) - - # DOGFOOD: hand the multi-latent rows to the langset library Trainer โ€” it runs the EMA-twin + FSQ + learned-STOP - # loop and selects on retrieval MRR, so this example never re-copies the training loop. - rows = [{"input_text": r["input_text"], - "target_texts": [_ent_text(s, t) for (s, t) in r["entities"][:args.max_entities]]} - for r in train_rows] - targs = TrainingArguments( - epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, tau=args.tau, max_len=args.max_len, - ema_m=args.ema_m, max_target_items=args.max_entities, max_steps=args.max_steps, val_frac=0.1, - output_dir=args.out, report_to=("wandb" if args.wandb else None), - wandb_project="langset-ner-multi-latent", seed=args.seed, verbose=True) - model = Trainer(model, targs, rows, on_checkpoint=commit_fn).train() - - # GROUND-TRUTH entity-F1 on the held-out validation split (the NER-specific quality read). - m = evaluate_f1(model, val_rows, args) - print(f"[ner] held-out entity-F1={m['f1']:.3f} (p={m['precision']:.3f} r={m['recall']:.3f}) " - f"tf_acc={m['tf_acc']:.3f} decode_acc={m['decode_acc']:.3f} emit={m['avg_emitted']:.2f}", flush=True) - - demo = model.rollout("Angela Merkel met Tim Cook at Apple headquarters in Cupertino .", max_steps=args.max_steps) - print(f"[demo] emitted {demo.shape[0]} entity latents for a 3-entity sentence -> {tuple(demo.shape)}", flush=True) - - -def main() -> None: - run(build_parser().parse_args()) - - -if __name__ == "__main__": - main() diff --git a/examples/ner-multi-latent/train_modal.py b/examples/ner-multi-latent/train_modal.py deleted file mode 100644 index a8e3d13..0000000 --- a/examples/ner-multi-latent/train_modal.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Modal runner for the ner-multi-latent demo โ€” full CoNLL-2003 on an A100. - - modal run --detach examples/ner-multi-latent/train_modal.py # full run - modal run --detach examples/ner-multi-latent/train_modal.py --epochs 20 --limit 0 - -Reports f1 / precision / recall / collapse per epoch to wandb (project: langset-ner-multi-latent). -""" -from __future__ import annotations - -from pathlib import Path - -import modal - -_ROOT = Path(__file__).resolve().parent.parent.parent # langset repo root - -# Standard transformers image. langset's build_backbone loads every backbone via AutoModelForCausalLM + plain PEFT -# LoRA, so no special model-patching wheel is needed. torch is pulled from the cu128 index so the Blackwell (B200) -# GPU is supported; everything else layers on top. -image = ( - modal.Image.debian_slim(python_version="3.11") - # hf_transfer is ignored by current huggingface_hub -> Xet high-performance is the fast path (law team: 59 - # min/file single-stream -> 91s for all files). Matters for cold-start eval/play containers. - .env({"HF_XET_HIGH_PERFORMANCE": "1", "HF_HOME": "/cache/hf"}) - .pip_install("torch>=2.7.0", extra_index_url="https://download.pytorch.org/whl/cu128") - .pip_install("transformers>=4.50", "hf_xet", "datasets", "wandb>=0.18", "numpy>=1.26", - "peft>=0.13.2", "sentence-transformers>=3.0") - .add_local_dir(str(_ROOT / "src"), "/pkg/src") - .add_local_dir(str(_ROOT / "examples" / "ner-multi-latent"), "/pkg/example") -) - -app = modal.App("langset-ner-multi-latent") -hf_cache = modal.Volume.from_name("langset-hf-cache", create_if_missing=True) - - -@app.function(image=image, gpu="B200", timeout=21600, - volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("wandb-api-key"), modal.Secret.from_name("huggingface")]) -def train_ner(limit: int = 0, epochs: int = 15, batch_size: int = 32, max_len: int = 128, - lr: float = 3e-4, latent_dim: int = 0, fsq_dim: int = 128, fsq_levels: int = 8, - ema_m: float = 0.99, seed: int = 0, - out: str = "/cache/ner-best", llm: str = "HuggingFaceTB/SmolLM2-135M") -> None: - import os - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # anti-fragmentation for 4B backbones - import sys - sys.path.insert(0, "/pkg/src") - sys.path.insert(0, "/pkg/example") - from train import build_parser, run # type: ignore[import-not-found] - - args = build_parser().parse_args([]) # defaults, then override - args.llm = llm - args.limit, args.epochs, args.batch_size, args.max_len, args.lr = limit, epochs, batch_size, max_len, lr - args.latent_dim = latent_dim # 0 => backbone hidden size - args.fsq_dim, args.fsq_levels, args.ema_m, args.seed = fsq_dim, fsq_levels, ema_m, seed # token-native FSQ - args.device = "" # auto -> cuda - args.wandb = True - args.out = out # best checkpoint, committed live each improvement - run(args, commit_fn=hf_cache.commit) - - -@app.function(image=image, gpu="A100-80GB", timeout=600, volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("huggingface")]) -def play(ckpt: str = "/cache/ner-gemma-e2b") -> None: - """Full entity-bank decode of the current best checkpoint (mirrors the SmolLM play_cpu): roll each sentence - into entity latents, retrieve each latent's nearest `TYPE: surface` from a bank of gold entities + distractors. - `modal run ...::play --ckpt /cache/ner-gemma-e2b`.""" - import sys - sys.path.insert(0, "/pkg/src") - import torch.nn.functional as F # noqa: F401 - from langset import LangSetModel # type: ignore[import-not-found] - - data = [ - ("Barack Obama visited Berlin to meet Angela Merkel .", ["PER: Barack Obama", "LOC: Berlin", "PER: Angela Merkel"]), - ("Apple and Microsoft announced a partnership in California .", ["ORG: Apple", "ORG: Microsoft", "LOC: California"]), - ("The United Nations sent aid to Syria and Lebanon .", ["ORG: United Nations", "LOC: Syria", "LOC: Lebanon"]), - ("Lionel Messi signed with Paris Saint-Germain in France .", ["PER: Lionel Messi", "ORG: Paris Saint-Germain", "LOC: France"]), - ("Google acquired a London startup for two billion dollars .", ["ORG: Google", "LOC: London"]), - ("Cristiano Ronaldo scored twice against Manchester United at Old Trafford .", - ["PER: Cristiano Ronaldo", "ORG: Manchester United", "LOC: Old Trafford"]), - ] - distractors = ["PER: Taylor Swift", "ORG: Toyota", "LOC: Tokyo", "MISC: Olympics", "ORG: NASA", "LOC: Amazon River"] - - hf_cache.reload() # pick up the latest committed best - m = LangSetModel.load(ckpt) - m.eval() - dev = m.device - bank_texts = sorted({e for _, ents in data for e in ents} | set(distractors)) - zb = F.normalize(m.emit(bank_texts).float(), dim=-1).to(dev) # emit() returns CPU (numpy-compat) -> match rollout's dev - for s, gold in data: - lat = F.normalize(m.rollout(s, max_steps=16).float(), dim=-1) # [L, d] on dev - sims = lat @ zb.t() - raw = [bank_texts[int(sims[j].argmax())] for j in range(lat.size(0))] - hits = len(set(raw) & set(gold)) - print(f"\n{s}\n gold ({len(gold)}): {gold}\n decoded : {raw} ({hits}/{len(gold)})", flush=True) - - -@app.function(image=image, gpu="A10G", timeout=1800, volumes={"/cache": hf_cache}) -def eval_ckpt(ckpt: str = "/cache/ner-best-ss", limit: int = 800) -> None: - """Diagnose recall: sweep the EOS stop-threshold (higher = emit MORE) and report F1 + the average emitted - count vs average gold count. If avg-emitted << avg-gold, the model is stopping too early (length problem).""" - import sys - sys.path.insert(0, "/pkg/src") - sys.path.insert(0, "/pkg/example") - import torch.nn.functional as F # noqa: F401 - from langset import LangSetModel # type: ignore[import-not-found] - import prepare # type: ignore[import-not-found] - - hf_cache.reload() - m = LangSetModel.load(ckpt); m.eval() - val = prepare.rows(limit, "validation") - bank_texts: list[str] = []; bank_tags: list[tuple[str, str]] = []; seen: set[str] = set() - gold: list[set[tuple[str, str]]] = [] - for r in val: - gp: set[tuple[str, str]] = set() - for (s, t) in r["entities"][:12]: - gp.add((s.lower(), t)); et = f"{t}: {s}" - if et not in seen: - seen.add(et); bank_texts.append(et); bank_tags.append((s.lower(), t)) - gold.append(gp) - zb = F.normalize(m.emit(bank_texts).float(), dim=-1).to(m.device) - avg_gold = sum(len(g) for g in gold) / len(gold) - - def score(thresh: float) -> tuple[float, float, float, float]: - tp = fp = fn = 0; nlat = 0 - for r, gp in zip(val, gold): - lat = F.normalize(m.rollout(r["input_text"], max_steps=16, stop_threshold=thresh).float(), dim=-1) - nlat += lat.size(0) - sims = lat @ zb.t() - pred = {bank_tags[int(sims[j].argmax())] for j in range(lat.size(0))} - tp += len(pred & gp); fp += len(pred - gp); fn += len(gp - pred) - p = tp / max(tp + fp, 1); rr = tp / max(tp + fn, 1) - return 2 * p * rr / max(p + rr, 1e-9), p, rr, nlat / len(val) - - print(f"avg gold entities/sentence = {avg_gold:.2f}", flush=True) - for th in (0.0, 1.0, 2.0, 4.0): - f1, p, rr, avg_len = score(th) - print(f"stop_thresh={th}: f1={f1:.3f} p={p:.3f} r={rr:.3f} avg_emitted={avg_len:.2f}", flush=True) - - -_SPACE_ENTITIES = [ - ("PER", "Lionel Messi"), ("PER", "Cristiano Ronaldo"), ("PER", "Neymar"), ("PER", "Kylian Mbappe"), - ("PER", "Zinedine Zidane"), ("PER", "Barack Obama"), ("PER", "Angela Merkel"), ("PER", "Vladimir Putin"), - ("PER", "Joe Biden"), ("PER", "Emmanuel Macron"), - ("ORG", "Apple"), ("ORG", "Microsoft"), ("ORG", "Google"), ("ORG", "Amazon"), ("ORG", "Samsung"), ("ORG", "IBM"), - ("ORG", "Manchester United"), ("ORG", "Real Madrid"), ("ORG", "Barcelona"), ("ORG", "Liverpool"), ("ORG", "Arsenal"), - ("LOC", "France"), ("LOC", "Germany"), ("LOC", "Japan"), ("LOC", "Brazil"), ("LOC", "Syria"), ("LOC", "Lebanon"), - ("LOC", "Berlin"), ("LOC", "London"), ("LOC", "Tokyo"), ("LOC", "Paris"), ("LOC", "Madrid"), - ("MISC", "Olympics"), ("MISC", "World Cup"), ("MISC", "Wimbledon"), ("MISC", "Nobel Prize"), -] - - -@app.function(image=image, gpu="A100-80GB", timeout=900, volumes={"/cache": hf_cache}, - secrets=[modal.Secret.from_name("huggingface")]) -def embed_space(ckpt: str) -> dict: - """Emit the curated entity bank as latents (the space the model retrieves against) and return them for - local 2D projection. Emits the faithful `TYPE: surface` bank form; within-type geometry is prefix-independent.""" - import sys - sys.path.insert(0, "/pkg/src") - import torch.nn.functional as F # noqa: F401 - from langset import LangSetModel # type: ignore[import-not-found] - - hf_cache.reload() - m = LangSetModel.load(ckpt) - m.eval() - texts = [f"{t}: {s}" for t, s in _SPACE_ENTITIES] - z = F.normalize(m.emit(texts).float(), dim=-1) # [N, d] (cpu) - return {"entities": [s for _, s in _SPACE_ENTITIES], "types": [t for t, _ in _SPACE_ENTITIES], - "latents": z.tolist()} - - -@app.local_entrypoint() -def space(qwen: str = "/cache/ner-qwen-4b", gemma: str = "/cache/ner-gemma-e2b-576", - out: str = "ner_embed_space.png") -> None: - """Side-by-side t-SNE of the emitted entity space for two checkpoints, colored by NER type.""" - import numpy as np - import matplotlib - matplotlib.use("Agg") - import matplotlib.pyplot as plt - from sklearn.manifold import TSNE - - runs = [("Qwen3.5-4B (f1 climbing)", qwen), ("Gemma E2B (f1 0.199)", gemma)] - colors = {"PER": "#e41a1c", "ORG": "#377eb8", "LOC": "#4daf4a", "MISC": "#984ea3"} - fig, axes = plt.subplots(1, len(runs), figsize=(17, 8)) - for ax, (name, ck) in zip(axes, runs): - d = embed_space.remote(ck) - X = np.asarray(d["latents"], dtype=np.float32); types = d["types"]; ents = d["entities"] - emb = TSNE(n_components=2, perplexity=8, random_state=0, init="pca", - learning_rate="auto").fit_transform(X) - for t, c in colors.items(): - idx = [i for i, tt in enumerate(types) if tt == t] - ax.scatter(emb[idx, 0], emb[idx, 1], c=c, label=t, s=90, alpha=0.85, edgecolors="white", linewidths=0.5) - for i, e in enumerate(ents): - ax.annotate(e, (emb[i, 0], emb[i, 1]), fontsize=7, alpha=0.75, - xytext=(3, 3), textcoords="offset points") - ax.set_title(name, fontsize=13); ax.legend(fontsize=10, loc="best") - ax.set_xticks([]); ax.set_yticks([]) - fig.suptitle("NER emitted-entity latent space (t-SNE, colored by type)", fontsize=15) - plt.tight_layout() - plt.savefig(out, dpi=140, bbox_inches="tight") - print(f"saved {out}", flush=True) - - -@app.local_entrypoint() -def main(limit: int = 0, epochs: int = 15, batch_size: int = 32, seed: int = 0, latent_dim: int = 0, - fsq_dim: int = 128, fsq_levels: int = 8, out: str = "/cache/ner-best", - llm: str = "HuggingFaceTB/SmolLM2-135M") -> None: - train_ner.remote(limit=limit, epochs=epochs, batch_size=batch_size, seed=seed, latent_dim=latent_dim, - fsq_dim=fsq_dim, fsq_levels=fsq_levels, out=out, llm=llm) diff --git a/src/langset/bridge_emission.py b/src/langset/bridge_emission.py index 6489196..b603d9e 100644 --- a/src/langset/bridge_emission.py +++ b/src/langset/bridge_emission.py @@ -13,8 +13,7 @@ HARD NEGATIVES: set `hard_neg_field=""` (a per-row list of confusable texts, e.g. adjacent-quarter queries) and the family folds them into its OWN InfoNCE denominator โ€” one softmax over [in-batch targets ++ hard-negs], the exact -mechanism that lets the emitted vector separate same-entity/wrong-time near-misses. Keep `lam_hard_neg=0` so langset's -separate HardNegTerm does not double-count. Absent `hard_neg_field`, the bank is plain in-batch (byte-identical). +mechanism that lets the emitted vector separate same-entity/wrong-time near-misses. Absent `hard_neg_field`, the bank is plain in-batch (byte-identical). Retrieval is preserved BY CONSTRUCTION: the backbone is frozen, so the base embedder's geometry is untouched; the bridge is a pure add-on. Eval routes through `emit_infer` (one pass, validity-gated) instead of the AR @@ -93,7 +92,7 @@ def forward( class QueryBridgeEmission(_EmissionObjective): - """Parallel-query continuous emission (see module docstring). `codebook=False`: no FSQ digits, no AR rollout.""" + """Parallel-query continuous emission (see module docstring). `codebook=False`: no codebook and no AR rollout.""" codebook = False @@ -215,8 +214,7 @@ def emit( base_loss=base, # reuse the canonical log keys (semantic map: recon_loss<-InfoNCE, loss_stop<-validity BCE, loss_dims<-0) logs={"loss_stop": vloss.detach(), "loss_dims": zero, "recon_loss": nce.detach()}, - dim_lg=None, - lab_label=None, + code_logits=None, ) @torch.no_grad() diff --git a/src/langset/heads.py b/src/langset/heads.py index 368924b..29f706e 100644 --- a/src/langset/heads.py +++ b/src/langset/heads.py @@ -8,7 +8,7 @@ A `Head` declares the four axes the phase head implicitly hard-coded: 1. READ SITE (`reads`): "recon" = one prediction per EMITTED latent, in the token stream (what the phase head - reads โ€” grad flows into the FSQ up/down-proj + LoRA, shaping the geometry). + reads โ€” grad flows into the emitter and backbone, shaping the geometry). "hidden" = the pooled/final backbone hidden, ONE prediction per sequence (a value/time readout of the current state; grad flows into the backbone/LoRA). 2. TARGET (`target`): a per-item dataset column supplying the supervision (like the existing `sup_field`). @@ -40,7 +40,7 @@ _READS = ("recon", "hidden") _BUILTIN_LOSSES = ("ce", "mse") -# label tokens treated as "no supervision" for this item (mirrors the phase-head / emb_slots filters exactly). +# label tokens treated as "no supervision" for this item. _MISSING = ("", "unknown", "none", "nan") diff --git a/src/langset/modeling.py b/src/langset/modeling.py index 9374da5..be1e052 100644 --- a/src/langset/modeling.py +++ b/src/langset/modeling.py @@ -48,20 +48,12 @@ def load_state_dict(self, *args: object, **kwargs: object) -> object: ... class EmitHead(nn.Module): - """Learned query tokens whose post-backbone hidden states are read out as the latent (the LLM's vector 'mouth'). - - One switch drives everything: `multi_latent`. OFF (default) โ†’ CONTINUOUS emission, a single raw vector, the - SetFit embedding path (uses `out_proj`). ON โ†’ the TOKEN-NATIVE sequence stack: each emission is a FINITE - SCALAR QUANTIZATION (FSQ) of the target embedding โ€” a learned `down_proj`/`up_proj` bottleneck into a FIXED - grid (`fsq_dim` dims, each rounded to `fsq_levels` levels). The grid is not learned and cannot drift or - collapse (unlike a VQ codebook), so it stays stable while the encoder/EMA-twin keep SPECIALIZING; near- - continuous precision comes from the grid resolution (like fixed-point), and the dims are INDEPENDENT so there - is no residual cascade. The emitter predicts each dim's digit (a small classification) plus a STOP folded into - dim-0's softmax, so termination/count is free. `multi_latent` is a task selector, not a tuning knob: `fsq_dim` - / `fsq_levels` are set-and-forget internals. - - `dropout` is trained ON when >0 so MC-dropout could later serve as an uncertainty signal. The head stays fp32 - even when the backbone is bf16. + """Read learned query-token hidden states as continuous or named-state vectors. + + The ordinary embedding path projects hidden states directly into one continuous vector. When + ``code_emit=True``, the head instead emits a distribution over a fixed, named codebook and feeds the + resulting superposition back into the backbone for autoregressive state rollout. QueryBridge supplies the + separate open-ended continuous multi-vector path and does not use this codebook head. """ def __init__( @@ -72,9 +64,6 @@ def __init__( dropout: float = 0.0, eos_id: int = 0, multi_latent: bool = False, - fsq_dim: int = 128, - fsq_levels: int = 8, - fsq_emit: bool = False, code_emit: bool = False, n_codes: int = 0, code_tau: float = 0.07, @@ -82,69 +71,26 @@ def __init__( ) -> None: super().__init__() self.n_latents = n_latents - self.q = nn.Parameter( - torch.randn(n_latents, h) * 0.02 - ) # one query token per emitted latent + self.q = nn.Parameter(torch.randn(n_latents, h) * 0.02) self.drop = nn.Dropout(dropout) - self.out_proj = nn.Linear(h, d) # hidden -> latent (targets / retrieval bank) - self.in_proj = nn.Linear(d, h) # latent -> hidden (inverse projection) + self.out_proj = nn.Linear(h, d) + self.in_proj = nn.Linear(d, h) self.eos_id = eos_id self.multi_latent = multi_latent - # fsq_emit: SINGLE-vector retrieval, but the embedding is GENERATED via the FSQ digit head (level_proj -> soft - # code -> up_proj) instead of the linear out_proj โ€” so it's produced through the same generative machinery CPT - # rewrites, letting injected knowledge reach the retrieval vector. Trainer routing stays on multi_latent (off). - # Default False => byte-identical: forward() uses out_proj and the FSQ projections match `multi_latent` alone. - self.fsq_emit = fsq_emit - self.fsq_dim, self.fsq_levels = fsq_dim, fsq_levels - _mk_fsq = multi_latent or fsq_emit - _mk_stop = _mk_fsq or code_emit - self.down_proj = nn.Linear(d, fsq_dim) if _mk_fsq else None # d -> FSQ bottleneck (learned) - self.up_proj = nn.Linear(fsq_dim, d) if _mk_fsq else None # FSQ -> d (learned decoder) - self.level_proj = ( - nn.Linear(h, fsq_dim * fsq_levels) if _mk_fsq else None - ) # per-dim digit logits - self.stop_proj = ( - nn.Linear(h, 1) if _mk_stop else None - ) # STOP logit (folded into dim-0 softmax; unused on the single-latent path) - # CODEBOOK EMISSION: replace the FSQ product-quantizer with ONE softmax over a fixed codebook. The head - # projects the hidden to a query, scores it against every code, and COMMITS to the softmax MIXTURE of - # codes rather than to a per-dim digit grid. Emitting a SET (maze frontier: which cells) is then a single - # normalized decision over the set's alphabet instead of `fsq_dim` independent ones, so mass must be - # ALLOCATED across the active members. `code` is a buffer, not a parameter: like the FSQ grid it cannot - # drift or collapse, so the geometry stays a constant of the run (see set_code). - # STATE + RESIDUAL. `res_dim` splits the latent in two: the first `d - res_dim` dims are the STATE, a - # mixture over the named codebook, and the last `res_dim` dims are a RESIDUAL that nothing names. The - # split is what makes a named alphabet safe to adopt: an alphabet is a CEILING (whatever it cannot name, - # the emission cannot carry), and the residual is where the unnamed goes. It also restores a - # self-supervised component โ€” the residual is shaped by the twin/recon terms, not by the labels. - # - # This differs from `emb_slots` in kind, not degree. emb_slots DECODES a facet out of a latent whose - # meaning came from the twin; the state half here CONSTITUTES the latent from the alphabet. A decode head - # cannot create information the model never computed (the frontier-head arm sat at its base-rate floor - # for exactly that reason); building the emission out of the members puts it there. self.code_emit = code_emit self.n_codes, self.code_tau = n_codes, code_tau self.res_dim = int(res_dim) self.state_dim = d - self.res_dim + self.stop_proj: Optional[nn.Linear] = nn.Linear(h, 1) if code_emit else None if code_emit: assert n_codes > 0, "code_emit requires n_codes > 0" assert 0 <= self.res_dim < d, ( f"res_dim must satisfy 0 <= res_dim < latent_dim ({d}); got {res_dim}" ) - self.register_buffer( - "code", torch.zeros(n_codes, self.state_dim) - ) # [C, state_dim] via set_code() - # CONCEPT FACETS. `concept_spans` carves the codebook into named groups: facet f owns member rows - # [m_lo, m_hi) and latent dims [d_lo, d_hi), and its rows are ZERO outside its own dims. So one - # matmul yields every facet's logits at once, while each facet gets its OWN softmax โ€” an album can - # be a wide mixture of `vocals` and certain of its `tempo` at the same time, and a filing's `stage` - # can be confident while its `result` is undecided. One span = the plain single-alphabet case. + self.register_buffer("code", torch.zeros(n_codes, self.state_dim)) self.concept_spans: list[tuple[int, int, int, int]] = [(0, n_codes, 0, self.state_dim)] self.concept_names: list[str] = [] - self.query_proj: Optional[nn.Linear] = nn.Linear( - h, self.state_dim - ) # hidden -> state query - # residual head: only built when a residual is asked for, so res_dim=0 stays byte-identical + self.query_proj: Optional[nn.Linear] = nn.Linear(h, self.state_dim) self.res_proj: Optional[nn.Linear] = ( nn.Linear(h, self.res_dim) if self.res_dim else None ) @@ -152,66 +98,29 @@ def __init__( self.query_proj = None self.res_proj = None - def forward( - self, hid_emit: torch.Tensor - ) -> torch.Tensor: # [B, n_latents, h] -> [B, n_latents, d] + def forward(self, hid_emit: torch.Tensor) -> torch.Tensor: if self.code_emit: - # EMBEDDING AS A SUPERPOSITION. The retrieval vector is a mixture over the NAMED alphabet, so it is - # interpretable by construction: project it back onto the codebook and you get the members it is - # made of, in the domain's own vocabulary, with no probe. Nothing supervises the weights here โ€” - # the ordinary contrastive/recon objective decides them, so an alphabet buys legibility without - # requiring per-row state labels. (`fsq_emit` established this shape: a retrieval vector generated - # THROUGH the emission head rather than a plain linear.) - lg, _ = self.emit_logits(hid_emit) # [B, nl, 1, n_codes] - p = self.concept_probs(lg) # [B, nl, n_codes] per-facet softmax - state = F.normalize(p @ self.code, p=2, dim=-1) + logits, _ = self.emit_logits(hid_emit) + state = F.normalize(self.concept_probs(logits) @ self.code, p=2, dim=-1) if self.res_dim == 0: return state - res = self.residual(hid_emit) # what the alphabet cannot name - return torch.cat([state, res], -1) * (0.5**0.5) - if self.fsq_emit: - lg, _ = self.emit_logits(hid_emit) # [B, nl, fsq_dim, fsq_levels] digit logits - levels = torch.arange(self.fsq_levels, device=lg.device, dtype=torch.float32) - soft = (lg.float().softmax(-1) * levels).sum( - -1 - ) # [B, nl, fsq_dim] expected digit (differentiable) - assert self.up_proj is not None - return F.normalize(self.up_proj(soft), p=2, dim=-1) # generated-FSQ retrieval vector + return torch.cat([state, self.residual(hid_emit)], -1) * (0.5**0.5) return F.normalize(self.out_proj(self.drop(hid_emit.float())), p=2, dim=-1) - def feedback(self, latent: torch.Tensor) -> torch.Tensor: # [B, ..., d] -> [B, ..., h] - return self.in_proj(latent.float()).to( - latent.dtype - ) # head stays fp32; match the backbone's dtype (bf16) - - def fsq(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Finite scalar quantize z [.., fsq_dim] -> (zq straight-through, digits). Both live in the SAME [0, L-1] - space so training (up_proj(zq)) and inference (up_proj(digits)) reconstruct identically.""" - lvl = self.fsq_levels - 1 - zb = (torch.tanh(z) + 1.0) * 0.5 * lvl # [0, L-1] - zr = zb.round().clamp(0, lvl) - zq = zb + (zr - zb).detach() # straight-through, in [0, L-1] - return zq, zr.long() + def feedback(self, latent: torch.Tensor) -> torch.Tensor: + return self.in_proj(latent.float()).to(latent.dtype) def set_code(self, code: torch.Tensor) -> None: - """Install the fixed codebook [n_codes, state_dim] (rows L2-normalized). Call once after construction.""" + """Install a fixed codebook [n_codes, state_dim] with normalized rows.""" assert self.code_emit and tuple(code.shape) == tuple(self.code.shape), ( f"set_code expects [{self.n_codes}, {self.code.size(-1)}], got {tuple(code.shape)}" ) self.code.copy_(F.normalize(code.float(), dim=-1).to(self.code.device)) def set_concepts(self, facets: "list[tuple[str, torch.Tensor, int]]") -> None: - """Install named concept facets: [(name, codes [n_members_f, dims_f], dims_f), ...]. - - Lays them out block-wise down the state half โ€” facet f's member rows carry its codes in its own dim - slice and zeros elsewhere โ€” so `query @ code.T` produces all facets' logits in one matmul while each - facet keeps a separate softmax.""" + """Install named concept facets as independent codebook spans.""" total_m = sum(c.shape[0] for _, c, _ in facets) total_d = sum(d for _, _, d in facets) - # The alphabet is DISCOVERED from the corpus, so its size is not knowable when the head is built โ€” - # `n_codes` at construction is only an upper bound. Resize to what the data actually contained rather - # than making the caller predict it; the true size is persisted in the config, so a reload rebuilds - # at the right width. if total_m != self.n_codes: self.n_codes = total_m self.register_buffer( @@ -232,94 +141,52 @@ def set_concepts(self, facets: "list[tuple[str, torch.Tensor, int]]") -> None: self.code.copy_(block.to(self.code.device)) self.concept_spans, self.concept_names = spans, names - def concept_probs(self, lg: torch.Tensor) -> torch.Tensor: - """Emission logits -> per-facet probabilities, concatenated back into one [.., n_codes] vector. - - Each facet is normalized WITHIN itself, so mass never leaks between facets: `vocals` summing to 1 says - nothing about how certain `tempo` is.""" - flat = lg.float().squeeze(-2) # [.., n_codes] + def concept_probs(self, logits: torch.Tensor) -> torch.Tensor: + """Normalize each named facet independently and concatenate the probabilities.""" + flat = logits.float().squeeze(-2) out = torch.zeros_like(flat) for m_lo, m_hi, _, _ in self.concept_spans: out[..., m_lo:m_hi] = flat[..., m_lo:m_hi].softmax(-1) return out def residual(self, hid: torch.Tensor) -> torch.Tensor: - """[.., h] -> the unnamed half of the emission [.., res_dim], L2-normalized. Empty when res_dim == 0.""" + """Return the normalized unnamed residual portion of an emission.""" if self.res_proj is None: return hid.new_zeros(*hid.shape[:-1], 0) return F.normalize(self.res_proj(self.drop(hid.float())), dim=-1) - def encode(self, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Target embedding -> (digits [.., fsq_dim], recon [.., d]). Grad flows to down/up_proj (recon loss).""" - if ( - self.code_emit - ): # no quantizer: the target IS the reconstruction, so teacher forcing feeds it exactly - # the codebook spans only the STATE half, so score the target's state dims against it โ€” the residual - # dims are unnamed by construction and have no nearest code - ts = F.normalize(t.float()[..., : self.state_dim], dim=-1) - idx = (ts @ self.code.t()).argmax(-1, keepdim=True) - return ( - idx, - t.float(), - ) # (nearest code [.., 1], recon [.., d] = the target itself, lossless) - assert self.down_proj is not None and self.up_proj is not None - zq, digits = self.fsq(self.down_proj(t.float())) - return digits, self.up_proj(zq) - - def reconstruct( - self, digits: torch.Tensor - ) -> torch.Tensor: # digits [.., fsq_dim] in [0, L-1] -> [.., d] - if self.code_emit: # a "digit" is a code index; reconstruction is the code itself - return self.code[digits.long().squeeze(-1)] - assert self.up_proj is not None - return self.up_proj(digits.float()) + def encode(self, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Map target vectors to nearest code indices for teacher-forcing bookkeeping.""" + assert self.code_emit, "encode(target) is only defined for a codebook emission head" + state = F.normalize(target.float()[..., : self.state_dim], dim=-1) + idx = (state @ self.code.t()).argmax(-1, keepdim=True) + return idx, target.float() + + def reconstruct(self, codes: torch.Tensor) -> torch.Tensor: + """Map code indices back to their fixed state vectors.""" + assert self.code_emit + return self.code[codes.long().squeeze(-1)] def emit_logits(self, hid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """[.., h] -> (per-dim digit logits [.., fsq_dim, fsq_levels], STOP logit [.., 1]). + """Return code logits ``[..., 1, n_codes]`` and an independent STOP logit.""" + assert self.code_emit and self.query_proj is not None and self.stop_proj is not None + query = F.normalize(self.query_proj(self.drop(hid.float())), dim=-1) + return ( + (query @ self.code.t()).unsqueeze(-2) / self.code_tau, + self.stop_proj(hid.float()), + ) - Codebook mode returns the SAME rank with fsq_dim collapsed to 1: [.., 1, n_codes] โ€” one softmax over the - alphabet instead of fsq_dim softmaxes over levels. Every caller downstream (STOP fold, commit, argmax) - is shape-agnostic, so the rollout loops are untouched.""" - assert self.stop_proj is not None - if self.code_emit: - assert self.query_proj is not None - q = F.normalize(self.query_proj(self.drop(hid.float())), dim=-1) # [.., d] - return (q @ self.code.t()).unsqueeze(-2) / self.code_tau, self.stop_proj(hid.float()) - assert self.level_proj is not None - lg = self.level_proj(self.drop(hid.float())).unflatten(-1, (self.fsq_dim, self.fsq_levels)) - return lg, self.stop_proj(hid.float()) - - def commit( - self, lg: torch.Tensor, sample: bool = False, hid: Optional[torch.Tensor] = None - ) -> torch.Tensor: - """Emission logits -> the latent that gets FED BACK to advance one tick. The FSQ branch is the historical - argmax(digits) -> up_proj, byte-for-byte. The codebook branch commits to the softmax MIXTURE of codes: a - genuine SUPERPOSITION of the members it is holding, not a single winner, which is the whole point of - emitting a set. `sample` draws instead of taking the mode (FSQ only; the mixture is already the full - distribution).""" - if self.code_emit: - p = self.concept_probs(lg) # [.., n_codes] each facet normalized within itself - state = F.normalize( - p @ self.code, dim=-1 - ) # [.., state_dim] facet mixtures, laid out side by side - if self.res_dim == 0: - return state - assert hid is not None, ( - "commit(): a residual emission needs the emit hidden (pass hid=)" - ) - # halve each half's norm so the concatenation is unit-norm and neither half can dominate the - # feedback vector purely by scale - return torch.cat([state, self.residual(hid)], -1) * (0.5**0.5) - if sample: - pred = torch.multinomial( - torch.softmax(lg.float(), -1).reshape(-1, self.fsq_levels), 1 - ).reshape(*lg.shape[:-2], -1) - else: - pred = lg.argmax(-1) # [.., fsq_dim] (own emission, never STOP) - return self.reconstruct(pred) + def commit(self, logits: torch.Tensor, hid: Optional[torch.Tensor] = None) -> torch.Tensor: + """Commit a code distribution to the superposed state vector fed into the next step.""" + assert self.code_emit + state = F.normalize(self.concept_probs(logits) @ self.code, dim=-1) + if self.res_dim == 0: + return state + assert hid is not None, "commit(): residual emission requires the emit hidden" + return torch.cat([state, self.residual(hid)], -1) * (0.5**0.5) def stop_logit(self, hidden: torch.Tensor, tok_embed: nn.Module) -> torch.Tensor: - """Alignment of the hidden state to the model's real EOS embedding (continuous-mode terminator).""" + """Alignment to the model's EOS embedding for the ordinary continuous path.""" emb_eos = cast(torch.Tensor, tok_embed.weight)[self.eos_id].float() return hidden.float() @ emb_eos @@ -540,9 +407,6 @@ def __init__( dropout: float = 0.0, max_len: int = 512, multi_latent: bool = False, - fsq_dim: int = 128, - fsq_levels: int = 8, - fsq_emit: bool = False, pool_mode: str = "", code_emit: bool = False, n_codes: int = 0, @@ -569,8 +433,6 @@ def __init__( self.latent_dim = latent_dim self.n_latents = n_latents self.multi_latent = multi_latent - self.fsq_dim = fsq_dim - self.fsq_levels = fsq_levels eos_id = int(tokenizer.eos_token_id or 0) self.head = EmitHead( self.h, @@ -579,9 +441,6 @@ def __init__( dropout, eos_id=eos_id, multi_latent=multi_latent, - fsq_dim=fsq_dim, - fsq_levels=fsq_levels, - fsq_emit=fsq_emit, code_emit=code_emit, n_codes=n_codes, code_tau=code_tau, @@ -631,14 +490,11 @@ def from_pretrained( bf16: bool = False, max_len: int = 512, multi_latent: bool = False, - fsq_dim: int = 128, - fsq_levels: int = 8, device: Optional[str] = None, attn_implementation: str = "sdpa", train_base: bool = False, grad_ckpt: bool = False, lora_top_k: int = 0, - fsq_emit: bool = False, pool_mode: str = "", freeze_backbone: bool = False, code_emit: bool = False, @@ -674,9 +530,6 @@ def from_pretrained( dropout, max_len, multi_latent, - fsq_dim, - fsq_levels, - fsq_emit=fsq_emit, pool_mode=pool_mode, code_emit=code_emit, n_codes=n_codes, @@ -713,20 +566,21 @@ def from_scratch( bf16: bool = False, max_len: int = 512, multi_latent: bool = False, - fsq_dim: int = 128, - fsq_levels: int = 8, device: Optional[str] = None, attn_implementation: str = "sdpa", grad_ckpt: bool = False, - fsq_emit: bool = False, arch_overrides: Optional[dict] = None, + code_emit: bool = False, + n_codes: int = 0, + code_tau: float = 0.07, + res_dim: int = 0, ) -> "LangSetModel": """RANDOM-INIT control arm (the "does pretraining matter" baseline). `arch` names an HF model whose ARCHITECTURE is copied, but the weights are NOT loaded โ€” the backbone starts from scratch and trains fully (no LoRA). The tokenizer is decoupled: pass any HF `tokenizer_id` (default = `arch`) and the fresh embedding table is sized to it, so there is no baked-in tokenizer. `arch_overrides` shrinks the net, e.g. `{"num_hidden_layers": 4, "hidden_size": 256, "num_attention_heads": 4, "num_key_value_heads": 4, - "intermediate_size": 1024}`. Everything downstream (emit head, FSQ, EMA-twin/SIGReg targets, losses) is + "intermediate_size": 1024}`. Everything downstream (emit head, target source, and losses) is identical to `from_pretrained`; only the source of the backbone weights differs.""" from transformers import AutoTokenizer # type: ignore[import-untyped] @@ -757,9 +611,10 @@ def from_scratch( dropout, max_len, multi_latent, - fsq_dim, - fsq_levels, - fsq_emit=fsq_emit, + code_emit=code_emit, + n_codes=n_codes, + code_tau=code_tau, + res_dim=res_dim, ) model._pretrained = False model._tokenizer_id = tokenizer_id @@ -1029,7 +884,7 @@ def generate_text(self, prompt: str, max_new: int = 200) -> str: ids = torch.cat([ids, torch.tensor([[nxt]], device=dev)], dim=1) return tok.decode(out, skip_special_tokens=True).strip() - # ---- multi-latent autoregressive rollout (drives the feedback / stop_logit seams) ---- + # ---- named-state autoregressive rollout ----------------------------------------------------- @torch.no_grad() def rollout( self, @@ -1045,251 +900,120 @@ def rollout( tuple[torch.Tensor, ...], tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]], ]: - """Autoregressive emission: real text tokens and emitted latents share ONE hidden stream. Each emitted - latent is fed back into the sequence via `head.feedback` (the latent lives in the backbone's own hidden - space, so an emitted latent and a real token embedding are the same kind of vector โ€” they intermingle). - Terminates per-row on the natural-EOS `stop_logit`. Returns [B, L, d] (or [L, d] for a single string), - zero-padding halted rows; pass `return_lengths` for the per-row emit count. `return_confidence` (FSQ path - only) additionally returns a dict of per-step NATIVE emission-confidence [B, L] tensors โ€” the model's OWN - co-trained uncertainty: 'code' = mean per-dim FSQ digit-softmax certainty, 'stop' = P(STOP) at that step, - 'dig0' = digit-0 level certainty. Returned as (lat, lengths, conf) when set. `temperature` (FSQ path only): - 0 = deterministic argmax (default); >0 = SAMPLE each FSQ digit from softmax(logits/temperature), so K rolls - of one seed fan into a DISTRIBUTION of plausible futures (McDonald's multi-rollout). `return_soft` returns - (lat, lengths, soft_lat, ent) where soft_lat is the EXPECTED latent E[digit]->up_proj (the calibrated - superposition, read from the full digit softmax rather than the argmax) and ent is the mean per-dim softmax - entropy (the model's native emission fuzziness). On the single-latent (out_proj) path soft_lat==lat and ent==0.""" + """Autoregressively emit named-state superpositions and feed each one back. + + ``temperature`` reshapes the code distribution before it is committed; zero uses the trained + distribution unchanged. ``return_soft`` returns the committed vectors and per-step code entropy. + QueryBridge is a separate non-autoregressive family and routes inference through its objective. + """ + if not self.head.code_emit: + raise ValueError( + "rollout() requires a named codebook emitter. Use model.encode() for one continuous vector, " + "or QueryBridgeEmission.emit_infer() for an open-ended continuous vector set." + ) single = isinstance(text, str) texts = [text] if single else list(text) dev = self.device - # token-native (multi-latent) FSQ path (codebook). Single-latent falls to the out_proj else-branch below. - codebook = self.head.multi_latent - stop_idx = ( - (self.head.n_codes if self.head.code_emit else self.head.fsq_levels) if codebook else -1 - ) # STOP is the extra class in dim-0's softmax (codebook mode: the class past the last code) - pad_side = ( - "left" if codebook else self.tokenizer.padding_side - ) # left-pad so [:, -1] is the last real token enc = self.tokenizer( texts, padding=True, truncation=True, max_length=self.max_len, - padding_side=pad_side, + padding_side="left", return_tensors="pt", ).to(dev) - seq = self.embed(enc["input_ids"]) # [B, S, h] โ€” real TEXT tokens + seq = self.embed(enc["input_ids"]) am = enc["attention_mask"] b = seq.size(0) alive = torch.ones(b, dtype=torch.bool, device=dev) lengths = torch.zeros(b, dtype=torch.long, device=dev) cols: list[torch.Tensor] = [] - conf_code: list[torch.Tensor] = [] # native per-step confidences (FSQ path) + ent_cols: list[torch.Tensor] = [] + conf_code: list[torch.Tensor] = [] conf_stop: list[torch.Tensor] = [] - conf_dig0: list[torch.Tensor] = [] - soft_cols: list[torch.Tensor] = [] # SOFT emission: E[latent] under digit softmax - ent_cols: list[torch.Tensor] = [] # native uncertainty: mean per-dim softmax entropy was_training = self.training self.eval() - for _step in range(max_steps): - if codebook: # AR: read the LAST real token's hidden - hid = self._last_hidden(self._run_backbone(seq, am, enc["input_ids"], 0))[ - :, -1 - ] # PLE-safe AR read - dim_lg, stop_lg = self.head.emit_logits(hid) # [B, fsq_dim, L], [B, 1] - d0_full = torch.cat( - [dim_lg[:, 0, :], stop_lg], -1 - ) # [B, L+1] over [L levels; STOP] - # CODEBOOK: termination is FACTORIZED out of the emission softmax. Folding STOP in (as the FSQ - # path does) is fair when dim-0's own target is one-hot, but a set-valued target is 1/k-diffuse - # and sharing one normalizer couples them: the gradient suppressing STOP at a member position - # scales with P(STOP), which weakens as the set widens. Independent sigmoid: "continue or stop" - # is decided on its own logit and never rides on how wide the set is. - code_stop = self.head.code_emit - pick = ( - dim_lg[:, 0, :] if code_stop else d0_full - ) # codebook: members only, STOP is separate - if temperature and temperature > 0: # SAMPLE each digit (stochastic futures) - dim0 = torch.multinomial( - torch.softmax(pick.float() / temperature, -1), 1 - ).squeeze(-1) - rest = ( - dim0.new_zeros(b, 0) # codebook: fsq_dim == 1, so there are no other dims - if code_stop - else torch.multinomial( - torch.softmax(dim_lg[:, 1:, :].float() / temperature, -1).reshape( - -1, self.head.fsq_levels - ), - 1, - ).reshape(b, -1) - ) - else: # deterministic argmax (default) - dim0 = pick.argmax(-1) - rest = dim_lg[:, 1:, :].argmax(-1) # [B, 0] under codebook emission - emit_now = ( - alive & (torch.sigmoid(stop_lg.squeeze(-1)) < 0.5) # independent terminator - if code_stop - else alive & (dim0 != stop_idx) # halt on STOP (or already-dead rows) - ) - if return_confidence: # the model's OWN emission uncertainty - p0 = torch.softmax(pick.float(), -1) # [B, L+1] (codebook: [B, n_codes]) - conf_stop.append( # P(STOP): its own sigmoid under codebook, else the folded class - torch.sigmoid(stop_lg.squeeze(-1)) if code_stop else p0[:, -1] - ) - conf_dig0.append((p0 if code_stop else p0[:, :-1]).max(-1).values) - conf_code.append( - torch.softmax(dim_lg.float(), -1).max(-1).values.mean(-1) - ) # mean code certainty - if return_soft: # the SUPERPOSITION reads out HERE, not from argmax - p = torch.softmax( - dim_lg.float(), -1 - ) # [B, fsq_dim, L] per-dim digit distribution - if ( - self.head.code_emit - ): # ONE distribution over the alphabet: mixture + its entropy - pc = p.squeeze(-2) # [B, n_codes] - soft_z = F.normalize(pc @ self.head.code, dim=-1) # mixture of codes [B, d] - ent = -(pc.clamp_min(1e-9).log() * pc).sum( - -1 - ) # [B] entropy = how many codes it is holding - else: - levels = torch.arange( - self.head.fsq_levels, device=dim_lg.device, dtype=torch.float32 - ) - soft_z = self.head.reconstruct( - (p * levels).sum(-1) - ) # E[digit] -> up_proj = expected latent [B, d] - ent = (-(p.clamp_min(1e-9).log() * p).sum(-1)).mean( - -1 - ) # [B] mean per-dim entropy (native fuzziness) - soft_cols.append( - torch.where(emit_now.unsqueeze(1), soft_z, torch.zeros_like(soft_z)) - ) - ent_cols.append(torch.where(emit_now, ent, torch.zeros_like(ent))) - digits = torch.cat( - [dim0.clamp(max=stop_idx - 1).unsqueeze(-1), rest], -1 - ) # [B, fsq_dim] (codebook: [B, 1], the top code index) - z = ( # the emitted latent: FSQ commits to the argmax code, codebook to the mixture - self.head.commit(dim_lg, hid=hid) - if self.head.code_emit - else self.head.reconstruct(digits) - ) # [B, d] - cols.append(torch.where(emit_now.unsqueeze(1), z, torch.zeros_like(z))) - lengths = lengths + emit_now.long() - seq = torch.cat([seq, self.head.feedback(z).unsqueeze(1).to(seq.dtype)], 1) - am = torch.cat([am, emit_now.long().unsqueeze(1)], 1) - alive = emit_now - else: - q = ( - self.head.q[:1].unsqueeze(0).expand(b, 1, -1).to(seq.dtype) - ) # one read-query token - am_q = torch.cat([am, am.new_ones(b, 1)], 1) - hid = self._last_hidden( - self._run_backbone(torch.cat([seq, q], 1), am_q, enc["input_ids"], 0) - )[:, -1] # [B, h] PLE-safe - z = self.head(hid.unsqueeze(1)).squeeze(1) # [B, d] emit via out_proj - stop = ( - self.head.stop_logit(hid, self.embed) > stop_threshold - ) # single-latent: emit-then-check on the - emit_now = alive # natural-EOS logit of the hidden that emitted - cols.append(torch.where(emit_now.unsqueeze(1), z, torch.zeros_like(z))) - if return_soft: # single-latent out_proj emission IS the soft - soft_cols.append( - torch.where(emit_now.unsqueeze(1), z, torch.zeros_like(z)) - ) # readout; entropy undefined - ent_cols.append(torch.zeros(b, device=dev)) # -> 0 - lengths = lengths + emit_now.long() - seq = torch.cat( - [seq, self.head.feedback(z).unsqueeze(1).to(seq.dtype)], 1 - ) # feed back the emitted latent - am = torch.cat( - [am, emit_now.long().unsqueeze(1)], 1 - ) # mask stopped/dead positions out of attention - alive = alive & ~stop + for _ in range(max_steps): + hid = self._last_hidden(self._run_backbone(seq, am, enc["input_ids"], 0))[:, -1] + logits, stop_logits = self.head.emit_logits(hid) + scaled = logits if temperature <= 0 else logits / temperature + probs = self.head.concept_probs(scaled) + state = F.normalize(probs @ self.head.code, dim=-1) + z = ( + torch.cat([state, self.head.residual(hid)], -1) * (0.5**0.5) + if self.head.res_dim + else state + ) + stop = stop_logits.squeeze(-1) > stop_threshold + emit_now = alive & ~stop + z = torch.where(emit_now.unsqueeze(-1), z, torch.zeros_like(z)) + cols.append(z) + p = probs.clamp_min(1e-9) + ent_cols.append( + torch.where(emit_now, -(p.log() * p).sum(-1), torch.zeros(b, device=dev)) + ) + conf_code.append(probs.max(-1).values) + conf_stop.append(torch.sigmoid(stop_logits.squeeze(-1))) + lengths = lengths + emit_now.long() + seq = torch.cat([seq, self.head.feedback(z).unsqueeze(1).to(seq.dtype)], 1) + am = torch.cat([am, emit_now.long().unsqueeze(1)], 1) + alive = emit_now if not bool(alive.any()): break if was_training: self.train() - lat = torch.stack(cols, 1) if cols else seq.new_zeros(b, 0, self.latent_dim) # [B, L, d] - if return_confidence: # per-step native confidences [B, L] - - def _st(xs: list[torch.Tensor]) -> torch.Tensor: - return torch.stack(xs, 1) if xs else lat.new_zeros(b, 0) - - conf = {"code": _st(conf_code), "stop": _st(conf_stop), "dig0": _st(conf_dig0)} - if single: - conf = {k: v[0, : int(lengths[0])] for k, v in conf.items()} + lat = torch.stack(cols, 1) if cols else seq.new_zeros(b, 0, self.latent_dim) + ent = torch.stack(ent_cols, 1) if ent_cols else lat.new_zeros(b, 0) + conf = { + "code": torch.stack(conf_code, 1) if conf_code else lat.new_zeros(b, 0), + "stop": torch.stack(conf_stop, 1) if conf_stop else lat.new_zeros(b, 0), + } if single: - lat = lat[0, : int(lengths[0])] - if return_soft: # SOFT readout: expected latent + native entropy - soft_lat = ( - torch.stack(soft_cols, 1) if soft_cols else lat.new_zeros(b, 0, self.latent_dim) - ) - ent_t = torch.stack(ent_cols, 1) if ent_cols else lat.new_zeros(b, 0) - if single: - soft_lat, ent_t = soft_lat[0, : int(lengths[0])], ent_t[0, : int(lengths[0])] - return lat, lengths, soft_lat, ent_t + n = int(lengths[0]) + lat, ent = lat[0, :n], ent[0, :n] + conf = {k: v[0, :n] for k, v in conf.items()} + if return_soft: + return lat, lengths, lat, ent if return_confidence: return lat, lengths, conf if return_lengths: return lat, lengths return lat - def rollout_train( - self, input_ids: torch.Tensor, attention_mask: torch.Tensor, target_latents: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Teacher-forced AR pass for training: feed the TRUE latents back through `feedback` and predict the next - emission at each future position in ONE forward pass (no python loop). Returns (preds [B, L, d], - stop_logits [B, L]) โ€” the per-step emission and its natural-EOS logit. (Continuous, single-latent path.)""" - bsz, s_len = input_ids.size(0), input_ids.size(1) - n = target_latents.size(1) - seed = self.embed(input_ids) # [B, S, h] - fb = self.head.feedback(target_latents.to(seed.dtype)) # [B, L, h] โ€” true latents fed back - seq = torch.cat([seed, fb], 1) - am = torch.cat([attention_mask, attention_mask.new_ones(bsz, n)], 1) - hid = self._last_hidden( - self._run_backbone(seq, am, input_ids, 0) - ) # PLE-safe teacher-forced read - hf = hid[:, s_len - 1 : s_len - 1 + n] # [B, L, h] โ€” predicts future positions - preds = self.head(hf) # [B, L, d] - stop = self.head.stop_logit(hf, self.embed) # [B, L] - return preds, stop - - def rollout_train_codebook( + def rollout_train_state( self, input_ids: torch.Tensor, attention_mask: torch.Tensor, target_latents: torch.Tensor, - tau: float = 0.07, train_hops: Optional[int] = None, ss_prob: float = 0.0, - ss_sample: bool = False, ss_mask: Optional[torch.Tensor] = None, kv_cache: bool = False, return_emit_hidden: bool = False, ) -> tuple[torch.Tensor, ...]: - """Token-native AR pass (multi-latent, FSQ). Quantizes the TRUE target latents to per-dim digits and predicts - the digits at each of L+1 positions (targets 0..L-1 then a STOP at position L). Returns (dim_logits - [B, L+1, fsq_dim, fsq_levels], stop_logits [B, L+1, 1], digits [B, L, fsq_dim], recon [B, L, d]) โ€” the STOP - label is appended by the trainer. + """Autoregressive concept/codebook pass. + + Maps targets to nearest code indices for code-classification objectives while preserving the exact target + vectors for teacher forcing. Returns code logits ``[B, L+1, 1, n_codes]``, STOP logits, indices, and exact + target vectors. `ss_prob`=0 (default): pure TEACHER FORCING in ONE forward pass โ€” every position predicted from the true prefix, gradients flow at most ONE hop (byte-identical to before). `ss_prob`>0: SCHEDULED SAMPLING for the first `train_hops` positions (None = all) โ€” with prob `ss_prob` each of those positions is fed the model's OWN emitted latent instead of the ground truth, so the emitter learns to consume its own (imperfect) predictions. This is the exposure-bias fix that makes MULTI-HOP rollout trained rather than emergent. Self- - fed latents are DETACHED (standard scheduled sampling; no BPTT through FSQ argmax). Cost = train_hops+1 - backbone passes (positions past train_hops are teacher-forced in one pass). `ss_sample` samples the self-fed - digits instead of argmax. + fed latents are DETACHED (standard scheduled sampling). Cost = train_hops+1 + backbone passes (positions past train_hops are teacher-forced in one pass). `ss_mask` (optional [B, H] bool): the PRECOMPUTED per-(row, hop) self-feed decisions. When given, the loop uses `ss_mask[:, h]` in place of a fresh `torch.rand < ss_prob` draw. This makes the rollout DETERMINISTIC given the mask, so GradCache's phase-1 (no_grad, full batch) and phase-2 (grad, per chunk) forwards produce identical `recon` and the cached gradients line up exactly. None = sample as usual (byte-identical).""" - assert self.head.multi_latent + assert self.head.multi_latent and self.head.code_emit bsz, s_len = input_ids.size(0), input_ids.size(1) n = target_latents.size(1) - digits, recon = self.head.encode(target_latents.reshape(-1, target_latents.size(-1))) - digits = digits.view(bsz, n, -1) # [B, L, fsq_dim] + codes, recon = self.head.encode(target_latents.reshape(-1, target_latents.size(-1))) + codes = codes.view(bsz, n, -1) # [B, L, 1] recon = recon.view(bsz, n, -1) # [B, L, d] โ€” clean feedback + recon target H = n if train_hops is None else max(0, min(int(train_hops), n)) if ss_prob <= 0 or n == 0 or H == 0: # TEACHER-FORCED one-shot (default, fast) @@ -1303,11 +1027,11 @@ def rollout_train_codebook( self._run_backbone(seq, am, input_ids, 0) ) # PLE-safe teacher-forced read hf = hid[:, s_len - 1 : s_len - 1 + n + 1] # [B, L+1, h] โ€” +1 to predict the STOP - dim_lg, stop_lg = self.head.emit_logits(hf) # [B, L+1, fsq_dim, L], [B, L+1, 1] + code_lg, stop_lg = self.head.emit_logits(hf) # [B, L+1, 1, n_codes], [B, L+1, 1] return ( - (dim_lg, stop_lg, digits, recon, hf) + (code_lg, stop_lg, codes, recon, hf) if return_emit_hidden - else (dim_lg, stop_lg, digits, recon) + else (code_lg, stop_lg, codes, recon) ) # SCHEDULED-SAMPLING multi-hop path dev = recon.device @@ -1331,7 +1055,7 @@ def rollout_train_codebook( def _feed( dl_h: torch.Tensor, t: int ) -> torch.Tensor: # the latent to advance tick t -> t+1 - recon_pred = self.head.commit(dl_h, ss_sample, hid=hid).detach() + recon_pred = self.head.commit(dl_h, hid=hid).detach() if t < H: # self-feed region: own emission or ground truth by the ss decision if ( ss_mask is not None @@ -1381,7 +1105,7 @@ def _feed( dim_parts.append(dl.unsqueeze(1)) stop_parts.append(sl.unsqueeze(1)) hid_parts.append(hid.unsqueeze(1)) - out4 = (torch.cat(dim_parts, 1), torch.cat(stop_parts, 1), digits, recon) + out4 = (torch.cat(dim_parts, 1), torch.cat(stop_parts, 1), codes, recon) return (*out4, torch.cat(hid_parts, 1)) if return_emit_hidden else out4 seq = self.embed(input_ids) am = attention_mask @@ -1390,13 +1114,11 @@ def _feed( hid_parts: list[torch.Tensor] = [] for h in range(H): # AR self-feed region hid = self._last_hidden(self._run_backbone(seq, am, input_ids, 0))[:, -1] - dl, sl = self.head.emit_logits(hid) # [B, fsq_dim, L], [B, 1] + dl, sl = self.head.emit_logits(hid) # [B, 1, n_codes], [B, 1] dim_parts.append(dl.unsqueeze(1)) stop_parts.append(sl.unsqueeze(1)) hid_parts.append(hid.unsqueeze(1)) - recon_pred = self.head.commit( - dl, ss_sample, hid=hid - ).detach() # own emitted latent (detached) + recon_pred = self.head.commit(dl, hid=hid).detach() # own emitted latent (detached) if ( ss_mask is not None ): # GradCache: replay the SHARED per-(row,hop) decisions (deterministic rollout) @@ -1414,15 +1136,15 @@ def _feed( am = torch.cat([am, am.new_ones(bsz, n - H)], 1) hid_all = self._last_hidden(self._run_backbone(seq, am, input_ids, 0)) hf_rest = hid_all[:, s_len - 1 + H : s_len - 1 + n + 1] # positions H..n (incl STOP at n) - dl_rest, sl_rest = self.head.emit_logits(hf_rest) # [B, n-H+1, fsq_dim, L], [B, n-H+1, 1] - dim_lg = ( + dl_rest, sl_rest = self.head.emit_logits(hf_rest) # [B, n-H+1, 1, n_codes], [B, n-H+1, 1] + code_lg = ( torch.cat(dim_parts + [dl_rest], 1) if dim_parts else dl_rest - ) # [B, n+1, fsq_dim, L] + ) # [B, n+1, 1, n_codes] stop_lg = torch.cat(stop_parts + [sl_rest], 1) if stop_parts else sl_rest # [B, n+1, 1] if return_emit_hidden: emit_hid = torch.cat(hid_parts + [hf_rest], 1) if hid_parts else hf_rest - return dim_lg, stop_lg, digits, recon, emit_hid - return dim_lg, stop_lg, digits, recon + return code_lg, stop_lg, codes, recon, emit_hid + return code_lg, stop_lg, codes, recon def get_sentence_embedding_dimension(self) -> int: return self.latent_dim @@ -1473,10 +1195,14 @@ def save_pretrained(self, path: Union[str, Path]) -> None: "n_latents": self.head.n_latents, "max_len": self.max_len, "multi_latent": self.multi_latent, - "fsq_dim": self.fsq_dim, - "fsq_levels": self.fsq_levels, "lora_top_k": self._lora_top_k, - "fsq_emit": self.head.fsq_emit, + "emission_family": ( + "query_bridge" + if hasattr(self, "emission_bridge") + else "codebook" + if self.head.code_emit + else "continuous" + ), "code_emit": self.head.code_emit, "n_codes": self.head.n_codes, "code_tau": self.head.code_tau, @@ -1506,6 +1232,18 @@ def load( p = Path(path) cfg = json.loads((p / "config.json").read_text()) + sd = torch.load(p / "langset.pt", map_location=device or "cpu", weights_only=False) + legacy_fsq = cfg.get("fsq_emit", False) or ( + cfg.get("multi_latent", False) + and not cfg.get("code_emit", False) + and "emission_bridge" not in sd + and "fsq_dim" in cfg + ) + if legacy_fsq: + raise ValueError( + "This checkpoint uses the removed FSQ emitter. Retrain with ConceptObjective, " + "StateResidualObjective, CodeSoftmaxObjective, or QueryBridgeEmission." + ) if cfg.get("pretrained", True): m = cls.from_pretrained( cfg["llm_model"], @@ -1514,10 +1252,7 @@ def load( lora_r=lora_r, max_len=cfg["max_len"], multi_latent=cfg.get("multi_latent", False), - fsq_dim=cfg.get("fsq_dim", 128), - fsq_levels=cfg.get("fsq_levels", 8), lora_top_k=int(cfg.get("lora_top_k", 0)), - fsq_emit=cfg.get("fsq_emit", False), code_emit=cfg.get("code_emit", False), n_codes=int(cfg.get("n_codes", 0)), code_tau=float(cfg.get("code_tau", 0.07)), @@ -1534,10 +1269,11 @@ def load( n_latents=cfg.get("n_latents", 1), max_len=cfg["max_len"], multi_latent=cfg.get("multi_latent", False), - fsq_dim=cfg.get("fsq_dim", 128), - fsq_levels=cfg.get("fsq_levels", 8), - fsq_emit=cfg.get("fsq_emit", False), arch_overrides=cfg.get("arch_overrides"), + code_emit=cfg.get("code_emit", False), + n_codes=int(cfg.get("n_codes", 0)), + code_tau=float(cfg.get("code_tau", 0.07)), + res_dim=int(cfg.get("res_dim", 0)), device=device, attn_implementation=attn_implementation, ) @@ -1547,7 +1283,6 @@ def load( if cfg.get("concept_spans"): m.head.concept_spans = [tuple(x) for x in cfg["concept_spans"]] m.head.concept_names = list(cfg.get("concept_names") or []) - sd = torch.load(p / "langset.pt", map_location=m.device, weights_only=False) # select by the ACTUAL payload (not the flag): a full-FT pretrained model persists "backbone", not "lora". m.backbone.load_state_dict(sd["backbone"] if "backbone" in sd else sd["lora"], strict=False) m.head.load_state_dict(sd["head"]) diff --git a/src/langset/probes.py b/src/langset/probes.py index 1ac6957..faae54f 100644 --- a/src/langset/probes.py +++ b/src/langset/probes.py @@ -46,7 +46,7 @@ def calibration_corr( """Superposition CALIBRATION: Pearson corr between the emitted latent's native entropy and the ground-truth number of possible next states at that emission. - entropy per-emission native FSQ entropy โ€” the `ent` array from `rollout(..., return_soft=True)`. + entropy per-emission concept entropy from `rollout(..., return_soft=True)`. cardinality per-emission ground-truth set size (how many next states were actually possible there). groups optional per-emission group id (e.g. trajectory/maze id). If given with `test_groups`, the correlation is computed on the held-out group split only, so it matches a decodability probe's diff --git a/src/langset/sigreg.py b/src/langset/sigreg.py index 309da4c..daa0e2b 100644 --- a/src/langset/sigreg.py +++ b/src/langset/sigreg.py @@ -8,13 +8,8 @@ Gaussian (Cramer-Wold). Used by SIGRegTarget (inject via `TrainingArguments(target_source=SIGRegTarget)`) as the alternative to the EMA -twin. TRAINING-ONLY (never persisted, never used at eval). For a token-native FSQ head, apply it NOT to the -emitted reconstruction but to the quantizer's own coordinates, penalized independently on each side (see -FSQObjective.z_for_reg): the TARGET side uses the pre-quantization z = down_proj(target_latent) (the raw FSQ -input), and the PREDICTED side uses the expected digit E[digit] = ฮฃ softmax(dim_logits)ยทlevels (there is no -down_proj on the predicted path โ€” the model emits digit logits, so its analogue of the pre-quant coordinate is -the soft digit expectation). Regularizing these spreads the encoder's codes across the whole grid, which is what -stops the (twin-free) live encoder folding every input into one cell. +twin. TRAINING-ONLY (never persisted, never used at eval). Each emission objective selects the representation +to regularize through ``z_for_reg``; codebook objectives use predicted and target laws over named members. """ from __future__ import annotations diff --git a/src/langset/strategies.py b/src/langset/strategies.py index b997eab..8254299 100644 --- a/src/langset/strategies.py +++ b/src/langset/strategies.py @@ -66,12 +66,12 @@ class MultiStepCtx: Shape legend used below: B = rows in this batch ยท L = `lmax` (max emitted items across the batch's rows, the padded time dim) ยท N = number of VALID emissions in the batch (= ฮฃ lens_l, since rows differ in length) - ยท d = latent dim ยท V = `fsq_levels`. The canonical flattened view a term works in is `recon[valid]` -> [N, d], + ยท d = latent dim. The canonical flattened view a term works in is `recon[valid]` -> [N, d], and `flat_texts` / `lens_l` / `bidx` are all aligned to that same row-major order. """ trainer: Trainer # the owning Trainer; read its PER-ROW data (indexed by dataset row id): - # sup_labels / hard_neg_texts / label_plan + label_cols + label_codewords + # sup_labels / hard_neg_texts / concept and state targets args: TrainingArguments # the run config โ€” a term reads its own weight/temperature here (a.lam_*, a.tau) model: LangSetModel # the online model being trained (rarely needed directly โ€” emit via target_source) dev: torch.device # device every tensor below lives on; build new tensors with device=c.dev @@ -91,14 +91,7 @@ class MultiStepCtx: recon: ( torch.Tensor ) # [B, L, d] the model's EMITTED latents this step โ€” gradient flows through these - dim_lg: Optional[ - torch.Tensor - ] # [B, L+1, fsq_dim, V] FSQ per-dim digit logits; None for a non-FSQ objective lmax: int # L above: the padded emitted-item time dim for this batch - fsq_levels: int # V above: FSQ quantization levels per digit - lab_label: Optional[ - torch.Tensor - ] # [B, L, n_reserved] reserved-dim label targets; None unless FSQ label subspace on target_source: _TargetSource # the target provider; call .encode(texts) -> [n, d] normalized latents (hard-neg bank) phase_head: Optional[ torch.nn.Module @@ -171,29 +164,6 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float return (self.key, loss_nce, a.lam_multi_nce) -class HardNegTerm(_LossTerm): - """Each emitted recon: own EMA target (positive) vs a shared bank of the batch's mined hard-negative texts.""" - - key = "loss_hard_neg" - - def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float]]: - a, self_ = c.args, c.trainer - if self_.hard_neg_texts is None or a.lam_hard_neg <= 0: - return None - hn_flat = [t for k in c.bidx for t in self_.hard_neg_texts[k]] - if not hn_flat: - return None - hn_bank = c.target_source.encode(hn_flat) # [Nhn, d] stop-grad normalized hard-neg latents - rv = F.normalize(c.recon[c.valid], dim=-1) # [Nvalid, d] emitted reconstructions - pos = (rv * c.target_lat[c.valid]).sum(-1, keepdim=True) # [Nvalid, 1] cos to own target - neg = rv @ hn_bank.t() # [Nvalid, Nhn] cos to every hard neg - logits_hn = torch.cat([pos, neg], dim=1) / a.tau - loss_hn = F.cross_entropy( - logits_hn, torch.zeros(logits_hn.size(0), dtype=torch.long, device=c.dev) - ) - return (self.key, loss_hn, a.lam_hard_neg) - - class SupConTerm(_LossTerm): """Supervised-contrastive shaping over emitted latents by the per-item `sup_field` group labels.""" @@ -237,26 +207,6 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float return (self.key, loss_phase, a.lam_phase) -class LabelDimsTerm(_LossTerm): - """FSQ LABEL SUBSPACE: full-strength CE on the reserved digit dims so the label lives AS coordinates of the - emitted code. FSQ-only (reads dim_lg); skipped when the objective produces no digit logits or no label plan.""" - - key = "loss_label" - - def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float]]: - a, self_ = c.args, c.trainer - if c.lab_label is None or c.dim_lg is None or a.lam_label_dims <= 0: - return None - plan = self_.label_plan - assert plan is not None # a label plan is what produced lab_label upstream - rcols = [cj for (cj, _, _) in plan] - lab_lg = c.dim_lg[:, : c.lmax, 1:, :][:, :, rcols, :] # [b, lmax, n_reserved, fsq_levels] - loss_label = F.cross_entropy( - lab_lg.reshape(-1, c.fsq_levels), c.lab_label.reshape(-1), ignore_index=-100 - ) - return (self.key, loss_label, a.lam_label_dims) - - def build_loss_terms(args: TrainingArguments) -> list[_LossTerm]: """DEFAULT loss-term set, built once from the args. Fixed order (label -> multi_nce -> hard_neg -> sup) so the float summation is byte-identical; each term self-skips when its weight/column is absent. Inject @@ -264,12 +214,7 @@ def build_loss_terms(args: TrainingArguments) -> list[_LossTerm]: NOTE: the phase head is no longer a term here โ€” it is the `phase` instance of the generic AUXILIARY-HEAD plug (langset.heads), applied inline in `_train_multi` right after this loop (the same summation position the old `PhaseTerm` held, so lam_phase>0 stays byte-identical). `PhaseTerm` is kept for back-compat injection.""" - return [ - LabelDimsTerm(), - MultiNCETerm(maskers=[identical_text_mask]), - HardNegTerm(), - SupConTerm(), - ] + return [MultiNCETerm(maskers=[identical_text_mask]), SupConTerm()] class CoTGenTerm(_LossTerm): @@ -333,30 +278,20 @@ def build_cot_loss_terms(args: TrainingArguments) -> list[_LossTerm]: # ---- emission objective ------------------------------------------------------------------------- @dataclass class EmissionOut: - """Result of one emission forward (shapes as in MultiStepCtx: B rows ยท L=lmax ยท d latent ยท V=fsq_levels). - The trainer sets `loss = base_loss`, folds `logs` into its running averages, and passes `recon`/`dim_lg`/ - `lab_label` on to the aux terms via MultiStepCtx.""" + """Result of one emission forward.""" recon: torch.Tensor # [B, L, d] the model's emitted latents โ€” gradient flows through these - base_loss: ( - torch.Tensor - ) # scalar: the objective's own loss (FSQ: loss_stop + loss_dims + recon_loss) + base_loss: torch.Tensor # scalar: the objective's own loss (objective-defined scalar) logs: dict[ str, torch.Tensor - ] # UNWEIGHTED scalar components for logging (FSQ: loss_stop / loss_dims / recon_loss) - dim_lg: Optional[ - torch.Tensor - ] # [B, L+1, fsq_dim, V] FSQ per-dim digit logits; None for a non-FSQ objective - lab_label: Optional[ - torch.Tensor - ] # [B, L, n_reserved] reserved-dim label targets; None unless FSQ label subspace on + ] # UNWEIGHTED scalar components for logging (objective-defined components) + code_logits: Optional[torch.Tensor] = None class _EmissionObjective: - """Strategy for turning the seeded forward into emitted latents + the base emission loss. Default = FSQ - (token-native digit CE + folded STOP + cosine recon). Selected ONCE per run so the step loop has no + """Strategy for turning the seeded forward into emitted latents + the base emission loss. Selected ONCE per run so the step loop has no emission `if`s. `codebook` is a class flag the free-run rollout reads to pick which emission head to use - (True = FSQ digit head, False = a raw-vector head). All objectives share the __init__ signature + (True = autoregressive codebook, False = strategy-owned continuous vectors). All objectives share the __init__ signature (model, args, dev, trainer) so they are interchangeable as an injected `TrainingArguments.emission`.""" codebook: bool = True @@ -400,7 +335,7 @@ def build_targets( ) -> tuple[torch.Tensor, torch.Tensor, list[int], int]: """Shape the per-row target item lists + their encoded latents into the emission's teacher-forcing tensors `(target_lat [B,L,d], valid [B,L], lens_l, lmax)`. EXTRACTED verbatim from the trainer's inline block so the - STRATEGY owns targetโ†”slot shaping: the AR/FSQ default is positional teacher forcing (item i -> slot i); a + STRATEGY owns targetโ†”slot shaping: the autoregressive default is positional teacher forcing (item i -> slot i); a future matching family (DETR/parallel-query) overrides this (e.g. all-valid, assignment deferred to its own Hungarian match inside emit()). Pure tensor assembly โ€” no model forward.""" lmax = max(len(x) for x in ent_lists) @@ -420,7 +355,7 @@ def build_targets( def emit_infer(self, texts: list[str], max_steps: int) -> tuple[torch.Tensor, torch.Tensor]: """Inference emission: texts -> (lat [B, Lmax, d], lens [B]), zero-padding halted rows. The eval and `model.rollout` route through here so a non-AR family (parallel-query) can emit in one pass. Default - delegates to the model's autoregressive rollout (FSQ).""" + delegates to the model's autoregressive concept/codebook rollout.""" lat, lens = self.m.rollout( # ty: ignore[invalid-assignment] # return_lengths=True -> (lat, lens) texts, max_steps=max_steps, return_lengths=True ) @@ -430,122 +365,21 @@ def z_for_reg( self, em: EmissionOut, target_lat: torch.Tensor, valid: torch.Tensor, lmax: int ) -> tuple[torch.Tensor, torch.Tensor]: """The (predicted, target) latents a TargetSource.regularizer (e.g. SIGReg) constrains. Default = - the emitted vs target latents directly ([N, d]); FSQ overrides this to use the pre-quantization z.""" + the emitted vs target latents directly ([N, d]).""" return em.recon[valid], target_lat[valid] -class FSQObjective(_EmissionObjective): - """DEFAULT emission: predict each item's per-dim FSQ digits (a STOP folded into dim-0's softmax) + a cosine - reconstruction to the target. Byte-identical to the historical inline FSQ block. Reads the FSQ grid geometry - (fsq_dim/fsq_levels) off model.head, so it takes the uniform (model, args, dev, trainer) signature.""" - - codebook = True - - def __init__( - self, model: LangSetModel, args: TrainingArguments, dev: torch.device, trainer: Trainer - ) -> None: - super().__init__(model, args, dev, trainer) - head = model.head - self.fsq_dim = int(head.fsq_dim) - self.fsq_levels = int(head.fsq_levels) - self.stop_idx = self.fsq_levels # STOP is the extra class folded into dim-0's softmax - - def emit( - self, - se: dict[str, torch.Tensor], - target_lat: torch.Tensor, - valid: torch.Tensor, - lens_l: list[int], - bidx: list[int], - b: int, - lmax: int, - ep: int, - ss_mask: Optional[torch.Tensor] = None, - ) -> EmissionOut: - m, a, dev, self_ = self.m, self.a, self.dev, self.trainer - fsq_dim, fsq_levels = self.fsq_dim, self.fsq_levels - ss_prob = a.ss_prob - assert ss_prob is not None # Trainer resolves the None sentinel to a float before any emit - eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) - dim_lg, stop_lg, digits, recon = m.rollout_train_codebook( - se["input_ids"], - se["attention_mask"], - target_lat, - a.tau, - train_hops=a.train_hops, - ss_prob=eff_ss, - ss_sample=a.ss_sample, - ss_mask=ss_mask, # GradCache: shared per-(row,hop) self-feed decisions (deterministic replay) - kv_cache=a.kv_cache, # forward prompt once + single-token hops vs full-prefix recompute (no grad_ckpt) - ) - dim0 = torch.cat([dim_lg[:, :, 0, :], stop_lg], -1) # [b, lmax+1, L+1] โ€” digit-0 + STOP - lab0 = torch.full((b, lmax + 1), -100, dtype=torch.long, device=dev) - lab_rest = torch.full((b, lmax, fsq_dim - 1), -100, dtype=torch.long, device=dev) - for r, nl in enumerate(lens_l): - lab0[r, :nl] = digits[r, :nl, 0] - lab0[r, nl] = self.stop_idx # emit digit-0 per item, then STOP after the last - lab_rest[r, :nl] = digits[r, :nl, 1:] - lab_label = None # FSQ LABEL SUBSPACE: reserved dims -> a SEPARATE - if self_.label_plan is not None: # weighted label CE (NOT diluted inside loss_dims) - plan = self_.label_plan - cols = self_.label_cols - assert cols is not None # label_cols is populated alongside label_plan - lab_label = torch.full((b, lmax, len(plan)), -100, dtype=torch.long, device=dev) - for s_i, (col_j, field, pos) in enumerate(plan): - labs, cw = cols[field], self_.label_codewords[field] - for r, kk in enumerate(bidx): - row_labs = labs[kk] - for j in range(lens_l[r]): - code = cw.get(row_labs[j] if j < len(row_labs) else "") - lab_label[r, j, s_i] = code[pos] if code is not None else -100 - lab_rest[:, :, col_j] = -100 # reserved dims leave the reconstruction CE - loss_stop = F.cross_entropy( - dim0.reshape(-1, fsq_levels + 1), lab0.reshape(-1), ignore_index=-100 - ) - loss_dims = F.cross_entropy( - dim_lg[:, :lmax, 1:, :].reshape(-1, fsq_levels), lab_rest.reshape(-1), ignore_index=-100 - ) - recon_loss = (1.0 - F.cosine_similarity(recon[valid], target_lat[valid], dim=-1)).mean() - return EmissionOut( - recon=recon, - base_loss=loss_stop + loss_dims + recon_loss, - logs={"loss_stop": loss_stop, "loss_dims": loss_dims, "recon_loss": recon_loss}, - dim_lg=dim_lg, - lab_label=lab_label, - ) - - def z_for_reg( - self, em: EmissionOut, target_lat: torch.Tensor, valid: torch.Tensor, lmax: int - ) -> tuple[torch.Tensor, torch.Tensor]: - # Regularize the PRE-QUANTIZATION z = down_proj(latent) โ€” the actual FSQ input, before the tanh+round โ€” - # so the penalty spreads the encoder's codes across the whole grid. z_pred = predicted E[digit]. - assert self.m.head.down_proj is not None - z_tgt = self.m.head.down_proj(target_lat[valid].float()) # [N, fsq_dim] - lvls = torch.arange(self.fsq_levels, device=self.dev, dtype=torch.float32) - dim_lg = em.dim_lg - assert dim_lg is not None # FSQ path always populates dim_lg - soft = (dim_lg[:, :lmax].float().softmax(-1) * lvls).sum( - -1 - ) # predicted E[digit] [b, lmax, fsq_dim] - return soft[valid], z_tgt - - class CodeSoftmaxObjective(_EmissionObjective): - """CODEBOOK emission: ONE softmax over the head's fixed codebook (plus STOP), trained toward the target's own - distribution over codes. Inject with `TrainingArguments(emission=CodeSoftmaxObjective)` and a head built with - `code_emit=True` (see EmitHead.set_code). + """Emit one distribution over a fixed alphabet and commit its codebook mixture. - Use this when an emission is a SET drawn from a known alphabet (maze frontier: which of 256 cells) rather than - a point to quantize. The FSQ default spends `fsq_dim` independent softmaxes on a learned grid, and a per-member - readout off that grid is `n_codes` INDEPENDENT decisions, so nothing forces the emission to choose among - members: every member can be scored high at once. A single normalized softmax has to ALLOCATE its mass, so the - members compete, and the emitted latent is the mixture they form. + Inject with `TrainingArguments(emission=CodeSoftmaxObjective)` and a head built with `code_emit=True`. + Use this when the target geometry already encodes membership in a known alphabet. One normalized softmax + allocates a fixed budget of mass, so members compete and the emitted latent is their superposition. The target law is read off the target latent itself, no extra column: for an orthonormal codebook and a membership target `normalize(multi_hot @ code.T)`, projecting back with `target @ code.T` recovers the members (equal weight) and ~0 elsewhere, so an L1 normalize gives the uniform-over-members law to match. Loss is the - soft-target cross-entropy over the codes, plus an INDEPENDENT sigmoid terminator (see `emit` for why STOP - must not be folded into the membership softmax the way the FSQ path folds it into dim-0). + soft-target cross-entropy over the codes, plus an independent sigmoid terminator. """ codebook = True @@ -557,6 +391,9 @@ def __init__( assert model.head.code_emit, ( "CodeSoftmaxObjective needs a codebook head: build the model with code_emit=True, n_codes=" ) + assert model.head.res_dim == 0, ( + "CodeSoftmaxObjective requires res_dim=0; use StateResidualObjective for named state plus a residual" + ) self.n_codes = int(model.head.n_codes) def emit( @@ -575,14 +412,12 @@ def emit( ss_prob = a.ss_prob assert ss_prob is not None # Trainer resolves the None sentinel before any emit eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) - dim_lg, stop_lg, _digits, recon = m.rollout_train_codebook( + code_logits, stop_lg, _digits, recon = m.rollout_train_state( se["input_ids"], se["attention_mask"], target_lat, - a.tau, train_hops=a.train_hops, ss_prob=eff_ss, - ss_sample=a.ss_sample, ss_mask=ss_mask, kv_cache=a.kv_cache, ) @@ -591,14 +426,9 @@ def emit( w = F.relu(target_lat.float() @ code.t()) # [b, lmax, n_codes] w = w / w.sum(-1, keepdim=True).clamp_min(1e-9) # MEMBERSHIP: soft-target CE over the codes ALONE, at real emission slots. - cell_logp = F.log_softmax(dim_lg[:, :lmax, 0, :].float(), -1) # [b, lmax, n_codes] + cell_logp = F.log_softmax(code_logits[:, :lmax, 0, :].float(), -1) # [b, lmax, n_codes] loss_code = -(w * cell_logp).sum(-1)[valid].mean() - # TERMINATION: an INDEPENDENT sigmoid, not a class inside the membership softmax. Folding STOP in (the - # FSQ default) is fair when dim-0's target is itself one-hot, but a set-valued target is 1/k-diffuse and - # the two then share one normalizer: at a member position the target is 0 on STOP, so the gradient - # pushing the stop logit DOWN is proportional to P(STOP) -- which shrinks as the set widens, because the - # denominator carries k members. Wider sets suppress the terminator more weakly and it ratchets up. - # Factored out, "continue vs stop" is trained by its own signal and never rides on set width. + # Termination is independent of the member distribution, so set width cannot weaken its supervision. stop_lab = torch.zeros(b, lmax + 1, device=dev) keep = torch.zeros(b, lmax + 1, dtype=torch.bool, device=dev) for r, nl in enumerate(lens_l): @@ -608,11 +438,13 @@ def emit( stop_lg.squeeze(-1).float()[keep], stop_lab[keep] ) # `recon` is the contract every loss term reads as "the model's emission, with gradient" (MultiNCETerm - # and HardNegTerm both build their query from it). The rollout's own recon is the TARGET here -- this + # and the emitter both build their query from it). The rollout's own recon is the TARGET here -- this # path is lossless, so it carries no gradient at all, and a term querying it would compare each target # against itself: near-zero loss, exactly zero gradient, silently inert. Hand back the DIFFERENTIABLE # committed mixture instead, which is both the true emission and the vector fed back each tick. - mix = F.normalize(dim_lg[:, :lmax, 0, :].float().softmax(-1) @ code, dim=-1) # [b, lmax, d] + mix = F.normalize( + code_logits[:, :lmax, 0, :].float().softmax(-1) @ code, dim=-1 + ) # [b, lmax, d] with ( torch.no_grad() ): # diagnostic: how close the emitted mixture lands to the target latent @@ -622,21 +454,20 @@ def emit( recon=mix, base_loss=loss_code + loss_stop, logs={"loss_code": loss_code, "loss_stop": loss_stop, "emit_cos": emit_cos}, - dim_lg=dim_lg, - lab_label=None, # the FSQ label subspace has no analogue here: one softmax, no dims to reserve + code_logits=code_logits, ) def z_for_reg( self, em: EmissionOut, target_lat: torch.Tensor, valid: torch.Tensor, lmax: int ) -> tuple[torch.Tensor, torch.Tensor]: - # Regularize the DISTRIBUTIONS, the codebook analogue of FSQ's pre-quantization z: predicted law vs target + # Regularize the DISTRIBUTIONS, the codebook analogue of codebook's pre-quantization z: predicted law vs target # law over codes, so a penalty spreads usage across the alphabet instead of collapsing onto a few codes. - dim_lg = em.dim_lg - assert dim_lg is not None + code_logits = em.code_logits + assert code_logits is not None code = self.m.head.code z_tgt = F.relu(target_lat[valid].float() @ code.t()) z_tgt = z_tgt / z_tgt.sum(-1, keepdim=True).clamp_min(1e-9) - return dim_lg[:, :lmax, 0, :].float().softmax(-1)[valid], z_tgt + return code_logits[:, :lmax, 0, :].float().softmax(-1)[valid], z_tgt class ConceptObjective(_EmissionObjective): @@ -734,19 +565,17 @@ def emit( ss_prob = a.ss_prob assert ss_prob is not None eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) - dim_lg, stop_lg, _d, _r, emit_hid = m.rollout_train_codebook( + code_logits, stop_lg, _d, _r, emit_hid = m.rollout_train_state( se["input_ids"], se["attention_mask"], target_lat, - a.tau, train_hops=a.train_hops, ss_prob=eff_ss, - ss_sample=a.ss_sample, ss_mask=ss_mask, kv_cache=a.kv_cache, return_emit_hidden=True, ) - flat = dim_lg[:, :lmax, 0, :].float() # [b, lmax, n_codes] + flat = code_logits[:, :lmax, 0, :].float() # [b, lmax, n_codes] tgt, seen = self._target_law(b, lmax, lens_l, bidx) # one soft-target CE PER FACET, over that facet's members only @@ -771,7 +600,7 @@ def emit( ) # the emission the rest of the trainer sees โ€” WITH gradient, or every aux term silently idles - p = m.head.concept_probs(dim_lg[:, :lmax]) + p = m.head.concept_probs(code_logits[:, :lmax]) state = F.normalize(p @ m.head.code, dim=-1) emission = ( torch.cat([state, m.head.residual(emit_hid[:, :lmax])], -1) * (0.5**0.5) @@ -789,8 +618,6 @@ def emit( "emit_cos": emit_cos, **per_facet, }, - dim_lg=dim_lg, - lab_label=None, ) @@ -867,14 +694,12 @@ def emit( ss_prob = a.ss_prob assert ss_prob is not None eff_ss = ss_prob if a.ss_warmup <= 0 else ss_prob * min(1.0, ep / a.ss_warmup) - dim_lg, stop_lg, _digits, _recon, emit_hid = m.rollout_train_codebook( + code_logits, stop_lg, _digits, _recon, emit_hid = m.rollout_train_state( se["input_ids"], se["attention_mask"], target_lat, - a.tau, train_hops=a.train_hops, ss_prob=eff_ss, - ss_sample=a.ss_sample, ss_mask=ss_mask, kv_cache=a.kv_cache, return_emit_hidden=True, # the residual is a function of the emit hidden, not of the logits @@ -892,7 +717,7 @@ def emit( if members: tgt[r, t, members] = 1.0 / len(members) has[r, t] = True - cell_logp = F.log_softmax(dim_lg[:, :lmax, 0, :].float(), -1) + cell_logp = F.log_softmax(code_logits[:, :lmax, 0, :].float(), -1) sel = has & valid loss_state = ( -(tgt * cell_logp).sum(-1)[sel].mean() if bool(sel.any()) else cell_logp.new_zeros(()) @@ -913,7 +738,7 @@ def emit( # THE EMISSION the rest of the trainer sees: state mixture ++ residual, exactly what feeds back. It must # be the emission WITH gradient -- every aux term (in-batch NCE, hard negatives) builds its query from # `recon`, and handing back the stop-grad target instead leaves them running but training nothing. - p = dim_lg[:, :lmax, 0, :].float().softmax(-1) + p = code_logits[:, :lmax, 0, :].float().softmax(-1) state = F.normalize(p @ code, dim=-1) if self.res_dim: res = m.head.residual(emit_hid[:, :lmax]) @@ -927,8 +752,6 @@ def emit( recon=emission, base_loss=loss_state + loss_stop, logs=logs, - dim_lg=dim_lg, - lab_label=None, ) @@ -1013,7 +836,7 @@ class CachedTarget(_TargetSource): """FROZEN-ENCODER target source with an encode-once cache โ€” the two-stage (V-JEPA) split for the multi-latent world model. The target geometry is a FIXED encoder, so every text->latent map is CONSTANT for the whole run: encode each unique text ONCE, memoize, then serve lookups. The step loop stops re-encoding BOTH the per-step - targets (trainer flat_texts) and the batch-pooled hard-negative bank (HardNegTerm) โ€” bringing the 'train only + targets (trainer flat_texts) and the batch-pooled hard-negative bank (the emitter) โ€” bringing the 'train only the head on cached vectors, epochs in seconds' win that langset's single-latent frozen-pool path already has to the multi-latent rollout. It's the fix for re-encoding fixed data every epoch. @@ -1074,7 +897,7 @@ def update(self) -> None: class SIGRegTarget(_TargetSource): """EMA-free anti-collapse (LeJEPA, arXiv:2511.08544). INJECT via `TrainingArguments(target_source=SIGRegTarget)`. Targets come from the LIVE model WITH gradient (no stop-grad twin); collapse is prevented by an isotropic-Gaussian - SIGReg penalty on the pre-quant z (via `regularizer`) instead of by a twin. So it drops the twin's VRAM + target + SIGReg penalty on the emitted representation (via `regularizer`) instead of by a twin. So it drops the twin's VRAM + target forward, the in-batch NCE is suppressed (the regularizer replaces it), and eval encodes with the live model itself. Reads scalar knobs off args: sigreg_lambda (loss weight, applied in the trainer), sigreg_knots, sigreg_slices.""" @@ -1103,8 +926,7 @@ def encode(self, texts: list[str]) -> torch.Tensor: return F.normalize(z.float(), dim=-1) def regularizer(self, z_pred: torch.Tensor, z_tgt: torch.Tensor) -> Optional[torch.Tensor]: - # Two INDEPENDENT Gaussianity penalties (predicted E[digit] and target z), NOT a match between them โ€” - # each is pushed toward isotropic Gaussian, spreading codes across the FSQ grid. + # Two independent Gaussianity penalties, not a match between predicted and target representations. return self.sig_reg(z_pred) + self.sig_reg(z_tgt) @@ -1165,7 +987,7 @@ def discover_concept_alphabet(rows_concepts: "list[object]") -> "dict[str, list[ # A `code_source` maps the alphabet's member NAMES to their vectors: (names, dim, model) -> [n_members, dim]. # Called ONCE at setup, over the alphabet (hundreds to a few thousand short strings), never per batch. The # result is frozen into a buffer for the run โ€” a codebook that re-embeds while the model trains is a learned -# codebook in disguise, and re-opens the collapse problem the fixed FSQ grid was chosen to avoid. +# codebook in disguise, and re-opens the collapse problem the fixed codebook was chosen to avoid. # # The choice is a real trade, not a default: an orthonormal codebook decodes losslessly by plain matmul (so a # recall number measures the EMISSION, with no probe as a confound) but its members carry no relation to each diff --git a/src/langset/trainer.py b/src/langset/trainer.py index dfa89bf..e3f6074 100644 --- a/src/langset/trainer.py +++ b/src/langset/trainer.py @@ -24,7 +24,7 @@ MultiStepCtx, ) # the aux-term step context; trainer builds it each step. The concrete -# strategies (FSQObjective/EMATwinTarget/...) are injected via TrainingArguments and read off `a`, not imported here. +# strategies (emission/target strategies/...) are injected via TrainingArguments and read off `a`, not imported here. from langset.training_args import TrainingArguments if TYPE_CHECKING: # only for local annotations of the injected strategy instances @@ -404,7 +404,7 @@ def __init__( # (e.g. modal Volume.commit) so another process can eval the live best checkpoint mid-training. self.on_checkpoint = on_checkpoint # ONE switch routes the whole trainer: a multi_latent model emits a VARIABLE-LENGTH latent set, so it reads - # a `target_texts` (list[str] per row) column and runs the FSQ token-native loop; otherwise the single-latent + # a `target_texts` (list[str] per row) column and runs the selected emission strategy; otherwise the single-latent # self-contrastive path (byte-for-byte unchanged) reads a scalar `target_text` column. self.multi_latent = bool(model.head.multi_latent) # ROLLOUT MUST BE TRAINED. A multi_latent model emits its latent set AUTOREGRESSIVELY at inference (rollout() @@ -412,6 +412,11 @@ def __init__( # rollout is never actually trained โ€” you cannot roll out what you did not train. Enforce it against the # ss_prob sentinel: unset (None) -> 0.25 (scheduled sampling on); an EXPLICIT ss_prob=0 -> hard error. if self.multi_latent: + if args.emission is None: + raise ValueError( + "multi_latent=True requires an explicit emission strategy: ConceptObjective, " + "StateResidualObjective, CodeSoftmaxObjective, or QueryBridgeEmission" + ) if args.ss_prob is None: args.ss_prob = 0.25 print( @@ -428,24 +433,6 @@ def __init__( ) elif args.ss_prob is None: args.ss_prob = 0.0 # single-latent never rolls; teacher-forced is correct - # VALIDATE emb_slots up front (both paths): a bad slice/kind otherwise fails deep in train() with an opaque - # Linear-shape or CE error. Enforce 0 <= lo < hi <= latent_dim (contiguous, in-bounds) and a supported kind. - slots = getattr(args, "emb_slots", None) - if slots: - d = model.latent_dim - for field, spec in slots.items(): - if not (isinstance(spec, (tuple, list)) and len(spec) == 3): - raise ValueError(f"emb_slots[{field!r}] must be (lo, hi, kind); got {spec!r}") - lo, hi, kind = spec - if kind not in ("classify", "regress"): - raise ValueError( - f"emb_slots[{field!r}] kind must be 'classify' or 'regress'; got {kind!r}" - ) - if not (isinstance(lo, int) and isinstance(hi, int) and 0 <= lo < hi <= d): - raise ValueError( - f"emb_slots[{field!r}] slice must satisfy 0 <= lo < hi <= latent_dim ({d}); " - f"got [{lo}:{hi}]" - ) # PLUGGABLE AUXILIARY HEADS are multi-latent only (they generalize the phase head, which lives only there). if getattr(args, "heads", None) and not self.multi_latent: raise ValueError( @@ -483,7 +470,7 @@ def __init__( cols = {**cols, "input_text": _init, "target_text": list(self._mask_texts)} # input_text/target_text are now SYNTHESIZED under their canonical keys -> read them directly (any # column_mapping for them is meaningless: their source is `text`, not a real input/target column). But - # KEEP `inv` intact so the optional fields below (mask_field, hard_neg_field, emb_slots, learn_field) + # KEEP `inv` intact so the optional fields below (mask_field, hard_neg_field, learn_field) # still honor column_mapping exactly as the non-masked path does -> same dataset+mapping, same behavior. get = lambda canon: cols[canon] # type: ignore[index] # noqa: E731 self.input_text = [str(x) for x in get("input_text")] @@ -510,8 +497,7 @@ def __init__( if cot_key in cols else [""] * len(self.input_text) ) - # optional MULTI-latent hard negatives: a per-row LIST of texts the emitted latents must be pushed AWAY - # from (batch-pooled InfoNCE bank, weight lam_hard_neg). Empty/None rows contribute no negatives. + # QueryBridge may consume a per-row list of hard-negative texts in its own contrastive denominator. self.hard_neg_texts: Optional[list[list[str]]] = None hn_field = getattr(args, "hard_neg_field", None) if hn_field is not None: @@ -580,75 +566,6 @@ def __init__( [[int(i) for i in tick] for tick in row[: args.max_target_items]] for row in raw_st ] - # optional FSQ LABEL SUBSPACE: per-item label columns -> reserved-dim codeword targets (a formal, head-free - # label space in the emitted code). Builds a class->codeword map per facet and a flat "plan" of which - # rest-dim column carries which facet's k-th codeword digit. - self.label_cols: Optional[dict[str, list[list[str]]]] = None - self.label_codewords: dict[str, dict[str, list[int]]] = {} - self.label_plan: Optional[list[tuple[int, str, int]]] = ( - None # (rest_col = dim-1, field, digit_pos) - ) - label_dims = getattr(args, "label_dims", None) - if label_dims: - fsq_levels = int(model.head.fsq_levels) - self.label_cols = {} - plan: list[tuple[int, str, int]] = [] - for field, dims in label_dims.items(): - raw = cols[inv.get(field, field)] - seqs = [ - [str(x) for x in (v if isinstance(v, (list, tuple)) else [v])] for v in raw - ] - self.label_cols[field] = seqs - classes = sorted( - { - c - for s in seqs - for c in s - if str(c).lower() not in ("", "unknown", "none", "na", "nan") - } - ) - m_dims = len(dims) - if len(classes) > fsq_levels**m_dims: - raise ValueError( - f"label_dims[{field}]: {len(classes)} classes > {fsq_levels}^{m_dims} " - f"codewords โ€” reserve more digits" - ) - cw: dict[str, list[int]] = {} - for ci, c in enumerate(classes): - x, digs = ci, [] - for _ in range(m_dims): - digs.append(x % fsq_levels) - x //= fsq_levels # little-endian base-fsq_levels codeword - cw[c] = digs - self.label_codewords[field] = cw - for pos, dd in enumerate(dims): - if int(dd) < 1: - raise ValueError( - f"label_dims dim {dd} must be >=1 (dim 0 is STOP-coupled)" - ) - plan.append((int(dd) - 1, field, pos)) - self.label_plan = plan - if args.verbose: - print( - "[langset] FSQ label subspace: " - + "; ".join( - f"{f}->{label_dims[f]} ({len(self.label_codewords[f])} cls)" - for f in label_dims - ), - flush=True, - ) - # optional CONTINUOUS EMB SLOTS (multi-latent): per-row facet label columns (named by the emb_slots dict - # keys), stashed as raw strings โ€” MIRRORS the single-latent reader below (392-399). The class<->id maps + - # transient decoder heads are built in _train_multi(). Set HERE because this branch returns before the - # single-latent slot setup. Facet labels are PER-ROW; the multi-latent decode mean-pools the emitted - # latent SET before slicing (see the slot loss in _train_multi). None = off (byte-identical). - self.slot_labels: Optional[dict[str, list[str]]] = None - slot_specs = getattr(args, "emb_slots", None) - if slot_specs: - self.slot_labels = {} - for field in slot_specs: - raw = cols[inv.get(field, field)] - self.slot_labels[field] = [("" if v is None else str(v)) for v in raw] # optional PLUGGABLE AUXILIARY HEADS (langset.heads): each user Head names a `target` dataset column. Read # the RAW per-row values here (a list per row for reads="recon", a scalar for reads="hidden"); resolution # to class maps / float tensors happens per-loss-kind in _train_multi. The phase shim reuses sup_labels, @@ -695,16 +612,6 @@ def __init__( if hn_field is not None: raw = cols[inv.get(hn_field, hn_field)] self.hard_neg_text = [str(v) if v not in (None, "") else "" for v in raw] - # optional CONTINUOUS EMB SLOTS: per-row facet label columns (named by the emb_slots dict keys). Stashed as raw - # strings here; the class<->id maps + transient decoder heads are built in train() (see TrainingArguments.emb_slots). - # (type declared once in the multi_latent branch above; plain re-assign here to avoid a mypy no-redef.) - self.slot_labels = None - slot_specs = getattr(args, "emb_slots", None) - if slot_specs: - self.slot_labels = {} - for field in slot_specs: - raw = cols[inv.get(field, field)] - self.slot_labels[field] = [("" if v is None else str(v)) for v in raw] # optional knowledge-injection: rows tagged "learn" train next-token CE (input_text -> target_text) instead # of contrastive; they're pulled OUT of the contrastive split and fed as a separate learn pool. self.is_learn: list[bool] = [False] * len(self.input_text) @@ -812,39 +719,8 @@ def learn_loss(pos: torch.Tensor) -> torch.Tensor: m, ln_doc_ids[pos], ln_doc_mask[pos], ln_tgt_ids[pos], ln_tgt_mask[pos], vsz ) - # CONTINUOUS EMB SLOTS: one TRANSIENT decoder head per facet, each reading ONLY its reserved dim slice of the - # emit vector `pred`. CE (classify) / MSE (regress). Grad flows into the encoder through those dims -> the - # facet is routed INTO the slice. Heads are trained but NOT saved (like the phase head; eval re-fits its own). - slot_plan: list[tuple[str, int, int, str, torch.nn.Module, dict[str, int]]] = [] - if self.slot_labels is not None: - emb_slots = a.emb_slots - assert emb_slots is not None # slot_labels is populated iff emb_slots was set - for field, (lo, hi, kind) in emb_slots.items(): - labs = self.slot_labels[field] - if kind == "classify": - classes = sorted( - {v for v in labs if v.lower() not in ("", "unknown", "none", "nan", "na")} - ) - cls2id = {c: i for i, c in enumerate(classes)} - head: torch.nn.Module = torch.nn.Linear(hi - lo, len(classes)).to(dev) - else: # regress - cls2id = {} - head = torch.nn.Linear(hi - lo, 1).to(dev) - slot_plan.append((field, lo, hi, kind, head, cls2id)) - if a.verbose: - print( - "[langset] emb slots: " - + "; ".join( - f"{f}->[{lo}:{hi}] {kind}({len(c) if c else 'reg'})" - for f, lo, hi, kind, _, c in slot_plan - ), - flush=True, - ) - slot_params = [p for (_, _, _, _, h, _) in slot_plan for p in h.parameters()] opt = torch.optim.AdamW( - [p for p in m.parameters() if p.requires_grad] - + list(connector.parameters()) - + slot_params, + [p for p in m.parameters() if p.requires_grad] + list(connector.parameters()), lr=a.lr, ) run = None @@ -876,17 +752,6 @@ def learn_loss(pos: torch.Tensor) -> torch.Tensor: if nm in _params: _params[nm].data.copy_(t.to(_params[nm].device, _params[nm].dtype)) connector.load_state_dict({k: v.to(dev) for k, v in ck["connector"].items()}) - for ( - field, - _, - _, - _, - head, - _, - ) in slot_plan: # transient slot heads (present only if emb_slots on) - sd = ck.get("slot_heads", {}).get(field) - if sd is not None: - head.load_state_dict({k: v.to(dev) for k, v in sd.items()}) opt.load_state_dict(ck["opt"]) for ( stt @@ -922,10 +787,6 @@ def save_resume(next_ep: int) -> None: nm: p.detach().cpu() for nm, p in m.named_parameters() if p.requires_grad }, "connector": {k: v.detach().cpu() for k, v in connector.state_dict().items()}, - "slot_heads": { - f: {k: v.detach().cpu() for k, v in h.state_dict().items()} - for f, _, _, _, h, _ in slot_plan - }, "opt": opt.state_dict(), "ep": int(next_ep), "best_score": float(best_score), @@ -1005,7 +866,7 @@ def _sl_loss( ) -> torch.Tensor: """Single-latent contrastive loss as a PURE FUNCTION OF THE EMBEDDINGS (pred/target/hn) โ€” the GradCache-compatible factoring of the inline step below. Contains every term that reads only the - pooled embeddings (in-batch-negative InfoNCE + false-neg / hard-neg masking + emb_slots facets + + pooled embeddings (in-batch-negative InfoNCE + false-neg / hard-neg masking facets + uniformity); the recon aux is NOT here (it needs the per-token backbone graph, so it stays on the direct path and grad_cache asserts lam_recon==0). Byte-identical math to the inline step, so grad_cache=False is unchanged.""" @@ -1037,37 +898,6 @@ def _sl_loss( neg_mask, float("-inf") ) # diagonal (positive) always kept loss = F.cross_entropy(logits, torch.arange(B, device=dev)) # primary - for ( - field, - lo, - hi, - kind, - head, - cls2id, - ) in slot_plan: # aux: CONTINUOUS EMB SLOTS (pure fn of pred) - sl = self.slot_labels - assert sl is not None - labs = sl[field] - sub = pred[:, lo:hi] - rows_i = idx.tolist() - if kind == "classify": - yi = torch.tensor([cls2id.get(labs[j], -100) for j in rows_i], device=dev) - if bool((yi >= 0).any()): - loss = loss + a.lam_emb_slots * F.cross_entropy( - head(sub), yi, ignore_index=-100 - ) - else: # regress: MSE on known rows - keep = [ - k - for k, j in enumerate(rows_i) - if labs[j].lower() not in ("", "unknown", "none", "nan", "na") - ] - if keep: - ki = torch.tensor(keep, device=dev) - yt = torch.tensor( - [float(labs[rows_i[k]]) for k in keep], device=dev - ).unsqueeze(1) - loss = loss + a.lam_emb_slots * F.mse_loss(head(sub[ki]), yt) if a.lam_uniform > 0 and B > 1: # aux: uniformity sq = torch.pdist(F.normalize(pred, p=2, dim=-1), p=2).pow(2) loss = loss + a.lam_uniform * sq.mul(-2.0).exp().mean().log() @@ -1182,7 +1012,7 @@ def _grad_cache_step(idx: torch.Tensor) -> float: # HARD NEGATIVES flow through _sl_loss as extra negative columns; grad still reaches only `pred`. loss = _sl_loss( pred, target, hn, idx - ) # InfoNCE + false-neg/hard-neg masking + emb_slots + uniformity (factored above) + ) # InfoNCE + false-neg/hard-neg masking + uniformity (factored above) if ( engine.supports_recon and a.lam_recon > 0 ): # aux: grounding. At 0 the term is zero anyway; @@ -1307,19 +1137,15 @@ def _grad_cache_step(idx: torch.Tensor) -> float: return m def _train_multi(self) -> LangSetModel: - """Multi-latent (variable-length FSQ latent-set) training. `input_text` seeds an autoregressive emission of a - latent per `target_texts` item, terminated by a learned STOP. Each target latent is supplied by a stop-grad - EMA twin (BYOL/JEPA) โ€” MANDATORY here: with the online model emitting its own target both sides move and the - geometry collapses. Objective = per-dim FSQ digit CE + a folded-in STOP CE + a cosine reconstruction; one - loss, no lam_* knobs. Selection = retrieval MRR against the row's own targets, with a non-collapse diversity - count logged as the anti-collapse guard. (Generalized from the validated regulatory seed-token loop.)""" + """Train an explicitly selected multi-vector emission objective. + + Named-state objectives autoregressively feed each committed concept mixture back and learn an independent + STOP decision. QueryBridge emits an unordered continuous vector set in one pass. A target-source strategy + supplies the comparison geometry, while optional auxiliary terms shape it without choosing the emitter.""" a, m = self.args, self.model dev = m.device tok = m.tokenizer - head = m.head d = int(m.latent_dim) - fsq_levels = int(head.fsq_levels) # only for MultiStepCtx; the emission objective derives - # its own fsq_dim/fsq_levels/stop_idx off model.head torch.manual_seed(a.seed) rng = np.random.default_rng(a.seed) @@ -1330,7 +1156,7 @@ def _train_multi(self) -> LangSetModel: futs = [lst[: a.max_target_items] for lst in self.target_texts] # cap targets per row if a.emit_seed: # PHASE-0 as an emitted node: prepend each seed's OWN text as target position 0, so the emitter learns to - # produce its start-state latent before the futures. Everything downstream (digits/STOP/recon/phase head/ + # produce its start-state latent before the futures. Everything downstream (state/STOP/recon/phase head/ # eval bank) shifts by one automatically; sup_labels gets a leading "phase0" class. Must happen HERE โ€” # before evaluate() closes over `futs` and before the phase_head label set is built from self.sup_labels. futs = [[seeds[i], *futs[i]] for i in range(len(futs))] @@ -1404,58 +1230,24 @@ def _train_multi(self) -> LangSetModel: ): # persisted heads: register on the model so they serialize + are queryable at inference if not h.spec.transient: m.add_aux_head(h.module, h.spec_dict()) - # CONTINUOUS EMB SLOTS (multi-latent port): one TRANSIENT decoder head per facet, each reading ONLY its reserved - # dim slice of the row's emit. In the single-latent path the emit is one vector; here the model emits a - # VARIABLE-LENGTH latent SET (em.recon [B, L, d] + valid [B, L]). DESIGN DECISION: to decode a PER-ROW facet - # from a per-row SET we mean-pool the VALID latents over L -> [B, d], then slice [:, lo:hi] โ€” this mirrors - # single-latent semantics (one vector per row) exactly and is the safe default. The transient head + CE/MSE loss - # (see the slot loss in the step loop) then force the encoder to route the facet into that slice; grad reaches - # every emitted latent equally through the mean-pool. ALTERNATIVE (future work): slot a single DESIGNATED latent - # (e.g. position 0 / the emit_seed node), or give each latent its OWN per-position slot label โ€” both need a - # richer per-position label schema than the per-row columns we have, so pooling is preferred until that lands. - # Heads are trained but NOT saved to the model (like phase_head; eval re-fits its own probe on the slice). - slot_plan: list[tuple[str, int, int, str, torch.nn.Module, dict[str, int]]] = [] - if self.slot_labels is not None: - emb_slots = a.emb_slots - assert emb_slots is not None # slot_labels is populated iff emb_slots was set - for field, (lo, hi, kind) in emb_slots.items(): - labs = self.slot_labels[field] - if kind == "classify": - classes = sorted( - {v for v in labs if v.lower() not in ("", "unknown", "none", "nan", "na")} - ) - cls2id = {c: i for i, c in enumerate(classes)} - slot_head: torch.nn.Module = torch.nn.Linear(hi - lo, len(classes)).to(dev) - else: # regress - cls2id = {} - slot_head = torch.nn.Linear(hi - lo, 1).to(dev) - slot_plan.append((field, lo, hi, kind, slot_head, cls2id)) - if a.verbose: - print( - "[langset] emb slots (multi, mean-pooled): " - + "; ".join( - f"{f}->[{lo}:{hi}] {kind}({len(c) if c else 'reg'})" - for f, lo, hi, kind, _, c in slot_plan - ), - flush=True, - ) # Build the emission strategy BEFORE the optimizer: a strategy may register its OWN trainable module on the # model (e.g. a parallel-query bridge over a frozen backbone), and `params` below must then capture it via - # `m.parameters()`. FSQ registers nothing extra, so this is byte-identical for the default path. + # `m.parameters()`. + assert ( + a.emission is not None + ) # validated in __init__; narrows the injected callable for the type checker objective: _EmissionObjective = a.emission( m, a, dev, self ) # emission strategy (INJECTED), built ONCE - slot_params = [p for (_, _, _, _, h, _) in slot_plan for p in h.parameters()] params = [ p for p in m.parameters() if p.requires_grad ] # includes PERSISTED aux heads + any strategy module registered above # TRANSIENT aux-head params are trainer-owned (not on the model), so add them explicitly โ€” exactly as the old - # phase head did. Ordering (base+head, then transient heads, then slots) preserves the byte-identical default - # (phase-only: base+head, phase, slots) since param order doesn't affect AdamW's per-parameter update. + # phase head did. Parameter order does not affect AdamW's per-parameter update. transient_head_params = [ p for h in rt_heads if h.spec.transient for p in h.module.parameters() ] - opt = torch.optim.AdamW(params + transient_head_params + slot_params, lr=a.lr) + opt = torch.optim.AdamW(params + transient_head_params, lr=a.lr) run = None if a.report_to == "wandb": import wandb # type: ignore[import-untyped] # ty: ignore[unresolved-import] # optional dep, not installed @@ -1466,7 +1258,7 @@ def _train_multi(self) -> LangSetModel: def evaluate() -> dict[str, float]: """Free-roll each val seed -> emitted latents; decode each by nearest-neighbor against an EMA-emitted bank of the val `target_texts`. Reports (a) retrieval MRR vs the chain's OWN targets and (b) a NON-COLLAPSE - diversity count = distinct nearest-bank items produced (FSQ must not mean-collapse to one mode).""" + diversity count = distinct nearest-bank items produced (state must not mean-collapse to one mode).""" m.eval() import time as _et @@ -1496,7 +1288,7 @@ def evaluate() -> dict[str, float]: emit_labs: list[str] = [] # position-aligned sup label of each emission for i in range(0, len(veval), a.batch_size): chunk = veval[i : i + a.batch_size] - # emission-strategy owns inference: FSQ delegates to the AR rollout; a parallel-query family emits + # emission-strategy owns inference: state delegates to the AR rollout; a parallel-query family emits # in ONE pass. Both return (lat [B,Lmax,d], len [B]). lats, lens = objective.emit_infer([seeds[c] for c in chunk], a.max_steps) for kk, ci in enumerate(chunk): @@ -1563,17 +1355,6 @@ def evaluate() -> dict[str, float]: hsd = ck.get("aux_heads", {}).get(h.spec.name) if hsd is not None: h.module.load_state_dict({k: v.to(dev) for k, v in hsd.items()}) - for ( - field, - _, - _, - _, - s_head, - _, - ) in slot_plan: # transient slot heads (present only if emb_slots on) - sd = ck.get("slot_heads", {}).get(field) - if sd is not None: - s_head.load_state_dict({k: v.to(dev) for k, v in sd.items()}) opt.load_state_dict(ck["opt"]) for ( stt @@ -1621,11 +1402,6 @@ def save_resume(next_ep: int) -> None: h.spec.name: {k: v.detach().cpu() for k, v in h.module.state_dict().items()} for h in rt_heads } - if slot_plan: # transient slot heads (only when emb_slots on) - payload["slot_heads"] = { - f: {k: v.detach().cpu() for k, v in h.state_dict().items()} - for f, _, _, _, h, _ in slot_plan - } torch.save(payload, tmp) tmp.replace(d / "resume.pt") if self.on_checkpoint is not None: @@ -1712,10 +1488,6 @@ def learn_loss(pos: torch.Tensor) -> torch.Tensor: # shared teacher-forced repl "grad_cache is incompatible with a cross-batch regularizer (e.g. SIGRegTarget): it is computed " "over the whole batch's latents and is not cached. Use the EMA twin, or grad_cache=False." ) - assert self.label_plan is None, ( - "grad_cache is incompatible with the FSQ label subspace (LabelDimsTerm reads dim_lg, the rollout " - "logits, not the cached recon). Drop label_field or grad_cache." - ) assert not any(getattr(t, "isolated_backward", False) for t in loss_terms), ( "grad_cache does not support isolated-backward loss terms." ) @@ -1738,11 +1510,10 @@ def _multi_grad_cache_step( agg: dict[str, float], ) -> float: """Multi-latent GradCache. Phase 1 rolls out the FULL batch under no_grad (a SHARED ss_mask makes the - scheduled-sampling rollout deterministic), runs the cross-batch recon-pure terms (InfoNCE, hard-neg, - supcon, phase, emb_slots) on the full-batch recon, and caches d(loss)/d(recon). Phase 2 re-rolls each - gc_chunk WITH grad, backprops the (row-weighted) base loss, and injects the cached recon-grads -> the - cross-batch term is EXACT full-batch; the per-row base loss is accumulated (pragmatic, see grad_cache - docs). Peak activation = one chunk.""" + scheduled-sampling rollout deterministic), runs the cross-batch recon-pure terms on the full-batch + recon, and caches d(loss)/d(recon). Phase 2 re-rolls each gc_chunk WITH grad, backprops the (row-weighted) + base loss, and injects the cached recon-grads -> the cross-batch term is EXACT full-batch; the per-row + base loss is accumulated (pragmatic, see grad_cache docs). Peak activation = one chunk.""" ss_prob = a.ss_prob assert ( ss_prob is not None @@ -1768,10 +1539,7 @@ def _multi_grad_cache_step( valid=valid, target_lat=target_lat, recon=recon_rg, - dim_lg=None, lmax=lmax, - fsq_levels=fsq_levels, - lab_label=None, target_source=target_source, phase_head=phase_head, phase_ids=phase_ids, @@ -1793,34 +1561,6 @@ def _multi_grad_cache_step( if raw_h is not None: cross = cross + h.spec.eff_weight(ep) * raw_h agg[h.spec.loss_key] = agg.get(h.spec.loss_key, 0.0) + float(raw_h.detach()) - if slot_plan: # emb_slots facets: pure fn of the pooled recon - _vm = valid.unsqueeze(-1).to(recon_rg.dtype) - pooled = (recon_rg * _vm).sum(1) / _vm.sum(1).clamp(min=1.0) - for field, lo, hi, kind, slot_head, cls2id in slot_plan: - sl = self.slot_labels - assert sl is not None - labs = sl[field] - sub = pooled[:, lo:hi] - if kind == "classify": - yi = torch.tensor([cls2id.get(labs[j], -100) for j in bidx], device=dev) - if bool((yi >= 0).any()): - _sl = F.cross_entropy(slot_head(sub), yi, ignore_index=-100) - cross = cross + a.lam_emb_slots * _sl - agg["loss_slots"] = agg.get("loss_slots", 0.0) + float(_sl.detach()) - else: - keep = [ - k - for k, j in enumerate(bidx) - if labs[j].lower() not in ("", "unknown", "none", "nan", "na") - ] - if keep: - ki = torch.tensor(keep, device=dev) - yt = torch.tensor( - [float(labs[bidx[k]]) for k in keep], device=dev - ).unsqueeze(1) - _sl = F.mse_loss(slot_head(sub[ki]), yt) - cross = cross + a.lam_emb_slots * _sl - agg["loss_slots"] = agg.get("loss_slots", 0.0) + float(_sl.detach()) opt.zero_grad() if cross.requires_grad: cross.backward() # fills recon_rg.grad (+ any head params: phase_head, slot heads) @@ -1854,7 +1594,9 @@ def _multi_grad_cache_step( for ( _k, _v, - ) in em_c.logs.items(): # whatever THIS objective logs (FSQ: stop/dims/recon) + ) in ( + em_c.logs.items() + ): # whatever THIS objective logs (objective-defined components) agg[_k] = agg.get(_k, 0.0) + float(_v.detach()) * w opt.step() target_source.update() @@ -1869,7 +1611,7 @@ def _multi_grad_cache_step( nb = 0 agg: dict[ str, float - ] = {} # accumulates whatever the emission objective logs (its keys, not FSQ's) + ] = {} # accumulates whatever the emission objective logs (its keys, not another objective's) for i in range(0, len(order), a.batch_size): if ( learn_pool and float(rng.random()) < a.learn_ratio @@ -1914,7 +1656,7 @@ def _multi_grad_cache_step( continue with _rf("rollout"): em = objective.emit(se, target_lat, valid, lens_l, bidx, b, lmax, ep) - recon, dim_lg, lab_label = em.recon, em.dim_lg, em.lab_label + recon = em.recon loss = em.base_loss # base objective (emission + STOP + recon); separation below c = MultiStepCtx( # everything the aux terms read this step trainer=self, @@ -1927,10 +1669,7 @@ def _multi_grad_cache_step( valid=valid, target_lat=target_lat, recon=recon, - dim_lg=dim_lg, lmax=lmax, - fsq_levels=fsq_levels, - lab_label=lab_label, target_source=target_source, phase_head=phase_head, phase_ids=phase_ids, @@ -1950,7 +1689,7 @@ def _multi_grad_cache_step( # phase-only run is byte-identical. Each self-skips (loss_on -> None) when nothing supervises it. if ( h.spec.reads == "recon" - ): # per emitted latent (grad shapes the FSQ geometry, like the phase head) + ): # per emitted latent (grad shapes the emission geometry, like the phase head) pred = h.module(recon[valid].float()) flat = h._flat_recon_targets(bidx, lens_l) else: # "hidden": pooled per-sequence seed hidden โ€” a separate backbone forward whose grad shapes it @@ -1972,39 +1711,6 @@ def _multi_grad_cache_step( ) # wants_regularizer=True sources return a tensor here loss = loss + a.sigreg_lambda * loss_sig agg["loss_sig"] = agg.get("loss_sig", 0.0) + float(loss_sig.detach()) - if slot_plan: # aux: CONTINUOUS EMB SLOTS (per-row facet, mean-pooled) - # MEAN-POOL the emitted latent SET over its VALID positions -> one [b, d] vector per row (the - # single-latent analog), then decode each facet from its reserved slice. gradient flows back - # through every valid emitted latent equally, routing the facet INTO those dims. - _vm = valid.unsqueeze(-1).to(em.recon.dtype) # [b, lmax, 1] - pooled = (em.recon * _vm).sum(1) / _vm.sum(1).clamp( - min=1.0 - ) # [b, d] mean over valid latents - for field, lo, hi, kind, slot_head, cls2id in slot_plan: - sl = self.slot_labels - assert sl is not None # slot_plan is non-empty only when slot_labels is set - labs = sl[field] # decode facet from pooled[:, lo:hi] - sub = pooled[:, lo:hi] # -> encoder routes it there - if kind == "classify": - yi = torch.tensor([cls2id.get(labs[j], -100) for j in bidx], device=dev) - if bool((yi >= 0).any()): - _sl = F.cross_entropy(slot_head(sub), yi, ignore_index=-100) - loss = loss + a.lam_emb_slots * _sl - agg["loss_slots"] = agg.get("loss_slots", 0.0) + float(_sl.detach()) - else: # regress: MSE on known rows - keep = [ - k - for k, j in enumerate(bidx) - if labs[j].lower() not in ("", "unknown", "none", "nan", "na") - ] - if keep: - ki = torch.tensor(keep, device=dev) - yt = torch.tensor( - [float(labs[bidx[k]]) for k in keep], device=dev - ).unsqueeze(1) - _sl = F.mse_loss(slot_head(sub[ki]), yt) - loss = loss + a.lam_emb_slots * _sl - agg["loss_slots"] = agg.get("loss_slots", 0.0) + float(_sl.detach()) opt.zero_grad() with _rf("backward"): loss.backward() # frees the main graph before any isolated term runs @@ -2021,7 +1727,10 @@ def _multi_grad_cache_step( target_source.update() # EMA twin tracks the online model tot += float(loss.detach()) nb += 1 - for _k, _v in em.logs.items(): # emission-defined keys (FSQ: stop/dims/recon) + for ( + _k, + _v, + ) in em.logs.items(): # emission-defined keys (objective-defined components) agg[_k] = agg.get(_k, 0.0) + float(_v.detach()) if _prof is not None: if torch.cuda.is_available(): @@ -2103,8 +1812,8 @@ def _sc(e: object) -> float: # self-CUDA us across torch versions sup_s = f" sup={row['loss_sup']:.3f}" if "loss_sup" in row else "" ph_s = f" phase={row['loss_phase']:.3f}" if "loss_phase" in row else "" pur_s = f" purity={pur:.3f}" if self.sup_labels is not None else "" - # the emission's own terms, in a stable order; unchanged for FSQ (stop/dims/recon), and a - # different objective prints ITS terms rather than KeyError-ing on FSQ's. + # the emission's own terms, in a stable order; unchanged for state (stop/dims/recon), and a + # different objective prints ITS terms rather than KeyError-ing on another objective's. _names = { "loss_stop": "stop", "loss_dims": "dims", diff --git a/src/langset/training_args.py b/src/langset/training_args.py index d31e09d..ab8f5fd 100644 --- a/src/langset/training_args.py +++ b/src/langset/training_args.py @@ -8,7 +8,6 @@ from langset.heads import Head from langset.strategies import ( EMATwinTarget, - FSQObjective, build_loss_terms, multi_epoch_order, multi_seed_texts, @@ -41,19 +40,15 @@ class TrainingArguments: # these are light aux terms: recon grounds the latent in the target text, uniform keeps the space spread. lam_recon: float = 0.3 # aux: the latent must also DECODE target_text lam_uniform: float = 0.1 # aux: light uniformity (spread latents on the sphere) - lam_hard_neg: float = ( - 0.0 # MULTI-LATENT hard-neg InfoNCE weight (0 = off, byte-identical to before) - ) # MULTI-LATENT in-batch-negative InfoNCE โ€” the separation term the multi-latent path was MISSING. The base # multi-latent loss (loss_stop + loss_dims + recon) is pure reconstruction: it MATCHES each emitted latent to # its EMA target but never pushes DIFFERENT items apart, so the geometry is capped by the base embedding and # organizes by whatever dominates the target text (boilerplate). This ports single-latent's primary objective: # each emitted recon vs the batch's targets, own target = positive, others = in-batch negatives (identical target # text masked as false-neg). ON by default โ€” this is the fix, not an opt-in. Set 0.0 for the old behavior. - # WEIGHT: unlike single-latent (where InfoNCE IS the whole loss, weight 1.0), here it COMPETES with the FSQ-recon + # WEIGHT: unlike single-latent (where InfoNCE IS the whole loss, weight 1.0), here it COMPETES with the state-emission # base (loss_stop+loss_dims+recon ~0.9); at 1.0 the InfoNCE (~ln(N_inbatch)โ‰ˆ4.85) overwhelms recon and SCRAMBLES # the geometry (retr_mrr stuck at chance, nce loss never descends). 0.3 balances the two โ†’ healthy geometry that - # organizes by real structure (FDA smol: disease decodes held-out, area_fsq +0.58 / indication_fsq +0.82, no labels). lam_multi_nce: float = 0.3 # SPEED (big models): stop_grad_target forwards the target view under no_grad (BYOL/MoCo-style) so it anchors @@ -78,10 +73,8 @@ class TrainingArguments: # hard negatives: name a dataset column of hard-negative text(s) per row (MINED near-miss targets the emitted # latent should be pushed AWAY from โ€” e.g. a boilerplate-similar case, or the WRONG lifecycle outcome). - # SINGLE-latent: a scalar hard-neg text per row, appended as an EXTRA always-negative column to the in-batch - # contrastive logits. MULTI-latent: a LIST of hard-neg texts per row; the batch's hard-neg latents form a shared - # bank and each emitted item's reconstruction runs an InfoNCE (own target vs the bank) at weight `lam_hard_neg`. - # Encoded under no_grad (memory-safe: no extra backward). None = no hard negs (byte-identical to before). + # Single-latent training appends one scalar hard negative to its contrastive bank. QueryBridge accepts a + # per-row list and folds those vectors into its own denominator. Named-state objectives use concepts instead. hard_neg_field: Optional[str] = None # supervised-contrastive label shaping (Khosla et al.): name a dataset column of GROUP LABELS the emitted latents @@ -114,39 +107,8 @@ class TrainingArguments: # shim when set. e.g. heads=[Head("value", reads="hidden", target="winprob", loss="mse", dim=1, transient=False)]. heads: list[Head] = field(default_factory=list) - # FSQ LABEL SUBSPACE โ€” a FORMAL label space inside the emitted code, with NO head per label. Map each facet to - # reserved FSQ digit indices (each >=1; dim 0 is STOP-coupled). Those dims' reconstruction targets are REPLACED - # by the label's codeword (one digit if n_classes<=fsq_levels, else little-endian base-fsq_levels across the - # group), so the emitter is FORCED to encode the label AS coordinates of the token: reading a label = argmax the - # reserved digit (no probe), writing = clamp it (controllable generation). Remaining dims reconstruct as usual. - # Needs per-item label columns (lists aligned 1:1 with target_texts) named by the dict keys. Unknown/""/"none" - # labels -> those dims ignored (-100) for that item. None = off (byte-identical). e.g. {"label_stage":[1], - # "label_area":[2,3], "label_indication":[4,5,6]}. - label_dims: Optional[dict[str, list[int]]] = None - # weight on the reserved-dim label CE. It's a SEPARATE term (not folded into loss_dims) so the ~9 label dims are - # not diluted among the ~1000 free reconstruction dims โ€” at weight 1 it sits at parity with loss_dims/recon. - lam_label_dims: float = 1.0 - - # CONTINUOUS EMBEDDING SLOTS โ€” the continuous analog of `label_dims` (which reserves FSQ *digit* indices). Reserve - # a CONTIGUOUS slice of dims of the single emitted embedding for a named facet and bind it with a small transient - # decoder head (CE for "classify", MSE for "regress") that reads ONLY those dims โ€” so gradient forces the encoder - # to route the facet INTO that slice. The facet then lives AS a readable sub-vector (emb[lo:hi] decodes to it) - # instead of smeared across all d dims, and WITHOUT a persisted model: like the phase head the slot head is not - # saved โ€” its only job is to inject gradient (eval re-fits its own probe on the now-informative slice). The slot - # dims still participate in the retrieval InfoNCE; they just ALSO carry a legible facet. Per-item labels come from - # a dataset column named by the dict key; unknown/""/"none"/"nan" -> that item ignored for the slot. Works on BOTH - # the single-latent path (slots the one emitted embedding) and the multi-latent path (MEAN-POOLS the emitted latent - # SET over its valid positions -> one per-row vector, then slots that โ€” same per-row semantics). - # None = off (byte-identical). e.g. {"material_bucket": (0, 64, "classify"), "side": (64, 80, "classify")}. - emb_slots: Optional[dict[str, tuple[int, int, str]]] = None - lam_emb_slots: float = 1.0 - - # STATE + RESIDUAL emission (multi-latent, StateResidualObjective) โ€” a THIRD way to shape the latent, - # alongside `fsq_dim` (how finely it is quantized) and `emb_slots` (which of its dims are made legible). - # Those two both DECODE meaning out of a latent the twin defined; this one CONSTITUTES the latent: its first - # `latent_dim - res_dim` dims are a mixture over a NAMED alphabet, and the remaining `res_dim` are a residual - # the alphabet cannot name (shaped by the twin/recon terms, so a self-supervised component survives). - # `state_field` names a per-row column of per-tick ACTIVE MEMBER index lists โ€” the same sparse shape as + # STATE + RESIDUAL emission. The named state is constituted from a fixed alphabet while the + # residual carries information outside that alphabet. `state_field` names a per-row column of per-tick ACTIVE MEMBER index lists โ€” the same sparse shape as # `concept_field` โ€” and `state_classes` is the alphabet size. Pair with `code_emit=True, n_codes=, # res_dim=` on the model. Off (None) = byte-identical. state_field: Optional[str] = None @@ -161,7 +123,7 @@ class TrainingArguments: # WHERE each named member's vector comes from. A name from strategies.CODE_SOURCES ("random" | "model" | # "orthogonal" | "twin") or your own (names, dim, model) -> [n_members, dim] callable. Resolved ONCE at # setup over the alphabet, then frozen โ€” a codebook that re-embeds during training is a learned codebook, - # with the drift/collapse problem the fixed FSQ grid exists to avoid. "random" is arbitrary but decodes + # which would let the geometry drift or collapse. "random" is arbitrary but decodes # losslessly (right for benchmarks: the readout cannot flatter the emission); "model"/"twin" put related # members near each other but give up exact recovery. `code_names` lists the member names in INDEX order, # so names[i] is the member whose index appears in the `state_field` column. @@ -170,14 +132,13 @@ class TrainingArguments: # MULTI-HOP training (scheduled sampling). For the first `train_hops` emitted positions (None = ALL), feed the # model's OWN predicted latent back with probability ss_prob instead of ground truth โ€” the exposure-bias fix that - # makes multi-hop rollout TRAINED rather than emergent. ss_sample samples the self-fed digits instead of argmax. + # makes multi-hop rollout TRAINED rather than emergent. # ss_prob SENTINEL: None = UNSET (resolved by the Trainer against the model). A multi_latent model rolls out # AUTOREGRESSIVELY at inference, so teacher-forced-only (ss_prob=0) is invalid โ€” the Trainer defaults an unset # ss_prob to 0.25 and REJECTS an explicit ss_prob=0 ("you cannot roll out what you did not train"). Single-latent # never rolls, so None -> 0.0 there. train_hops: Optional[int] = None ss_prob: Optional[float] = None - ss_sample: bool = False # KV-CACHE rollout (multi-latent): forward the prompt ONCE, then each latent token alone against the cached # prefix K/V, instead of re-running the whole growing sequence every tick. Numerically identical (verified to # ~1e-7 in test_kv_cache), but activation memory is ~1 prompt forward + n single tokens rather than n @@ -192,7 +153,7 @@ class TrainingArguments: # PARALLEL-QUERY emission (QueryBridgeEmission โ€” the non-AR continuous family, langset.bridge_emission). # n_queries = number of query slots = MAX emitted items per doc (the validity head gates the actual count โ‰ค this). - # Unused by the default FSQ/AR emission, so byte-identical there. + # Unused by autoregressive concept/codebook emission, so byte-identical there. n_queries: int = 16 bridge_lam_valid: float = 1.0 # weight on the validity/count BCE in the bridge's base loss bridge_pos_weight: float = ( @@ -215,10 +176,9 @@ class TrainingArguments: learn_field: Optional[str] = None learn_ratio: float = 0.0 - # multi-latent (variable-length FSQ latent-set emission). Active only when the model is `multi_latent=True`: + # multi-latent (variable-length vector emission). Active only when the model is `multi_latent=True`: # the Trainer reads a `target_texts` (list[str] per row) column, emits an EMA-twin target for each item, and - # trains the token-native FSQ emitter (per-dim digits + a learned STOP). Ignored on the single-latent path. - # `fsq_dim`/`fsq_levels` are MODEL internals (read from `model.head`), NOT args. + # trains the selected emission strategy. Ignored on the single-latent path. ema_m: float = 0.99 # EMA-twin momentum supplying the (stop-grad) target latents max_target_items: int = 12 # cap on target latents emitted/supervised per row max_steps: int = 16 # free-rollout cap used in multi-latent eval @@ -248,7 +208,7 @@ class TrainingArguments: # `seed_builder=cot_seed_texts` (see the injection block below) with a `cot_text` dataset column: the model is # co-trained to GENERATE each row's chain-of-thought (its own isolated backward, so the two autograd graphs never # coexist) and the latent forward is conditioned on seed+CoT. No injection / absent cot_text column = OFF - # (byte-identical FSQ path). Not a flag โ€” inject the strategies; this scalar just weights the CoT next-token CE. + # . Not a flag โ€” inject the strategies; this scalar just weights the CoT next-token CE. lam_cot: float = 1.0 # weight on the CoT next-token CE (used only with the CoT strategies) # SUPERPOSITION (multi-latent) โ€” to make one emitted latent hold a calibrated MIXTURE of several possible next @@ -264,11 +224,13 @@ class TrainingArguments: # ---- MULTI-LATENT STRATEGY INJECTION (dependency injection, not flags) -------------------------------------- # The multi-latent step is assembled from swappable pieces (see langset/strategies.py). Each field below holds # the STRATEGY ITSELF โ€” a class or callable โ€” and the trainer just calls it; there is NO `if use_x:` selection. - # The defaults reproduce the historical behavior byte-for-byte (guarded by test_trainer_multi_characterization). + # There is deliberately no implicit multi-vector emission default; choose the geometry explicitly. # To change a behavior, INJECT a different implementation, e.g. # TrainingArguments(target_source=SIGRegTarget) # EMA-free anti-collapse # rather than toggling a boolean the trainer then branches on. See strategies.py for the interface each must meet. - emission: Callable = FSQObjective # (model, args, dev, trainer) -> _EmissionObjective : seed->latents + base loss + emission: Optional[Callable] = ( + None # required for multi_latent: Concept/State/CodeSoftmax or QueryBridge + ) target_source: Callable = EMATwinTarget # (model, args, tok, dev) -> _TargetSource : the target latents + anti-collapse loss_terms: Callable = ( build_loss_terms # (args) -> list[_LossTerm] : the weighted aux separation/shaping terms diff --git a/tests/multi_state_helpers.py b/tests/multi_state_helpers.py new file mode 100644 index 0000000..45ac708 --- /dev/null +++ b/tests/multi_state_helpers.py @@ -0,0 +1,90 @@ +"""Shared deterministic fixtures for named-state multi-vector trainer tests.""" + +from __future__ import annotations + +import os +import tempfile + +import numpy as np +import torch + +from langset import LangSetModel, Trainer, TrainingArguments +from langset.strategies import CodeSoftmaxObjective + +TINY_MODEL = os.environ.get( + "LANGSET_TEST_MODEL", "hf-internal-testing/tiny-random-LlamaForCausalLM" +) + + +def _seed() -> None: + torch.manual_seed(0) + np.random.seed(0) + torch.use_deterministic_algorithms(True, warn_only=True) + + +def _build_model() -> LangSetModel: + _seed() + model = LangSetModel.from_pretrained( + TINY_MODEL, + bf16=False, + device="cpu", + multi_latent=True, + code_emit=True, + n_codes=8, + ) + model.head.set_code(torch.randn(8, model.latent_dim)) + return model + + +def _dataset() -> list[dict]: + topics = ["late payment", "theft", "property tax"] + stages = ["filing", "hearing", "ruling"] + return [ + { + "input_text": f"case {i}: a dispute concerning {topics[i % 3]}", + "target_texts": [ + f"{stage}: the {topics[i % 3]} matter at the {stage} stage" for stage in stages + ], + "stage": list(stages), + } + for i in range(8) + ] + + +def _args(out_dir: str, **over: object) -> TrainingArguments: + values: dict[str, object] = { + "epochs": 2, + "batch_size": 4, + "lr": 1e-3, + "max_len": 64, + "report_to": None, + "verbose": False, + "eval_every": 99, + "patience": 99, + "val_frac": 0.25, + "seed": 0, + "sup_field": "stage", + "lam_sup": 0.2, + "lam_phase": 0.1, + "emission": CodeSoftmaxObjective, + "output_dir": out_dir, + } + values.update(over) + return TrainingArguments(**values) # type: ignore[arg-type] + + +def _flat_trainable(model: LangSetModel) -> np.ndarray: + return np.concatenate( + [ + p.detach().float().cpu().numpy().ravel() + for _, p in sorted(model.named_parameters(), key=lambda item: item[0]) + if p.requires_grad + ] + ).astype(np.float64) + + +def _run(**arg_over: object) -> dict[str, np.ndarray]: + model = _build_model() + with tempfile.TemporaryDirectory() as td: + Trainer(model, _args(td, **arg_over), _dataset()).train() + return {"params_post": _flat_trainable(model)} diff --git a/tests/test_aux_heads.py b/tests/test_aux_heads.py index 8f13056..da33989 100644 --- a/tests/test_aux_heads.py +++ b/tests/test_aux_heads.py @@ -19,7 +19,7 @@ import torch from langset import Head, LangSetModel, Trainer, TrainingArguments -from langset.strategies import last_epoch_selector +from langset.strategies import CodeSoftmaxObjective, last_epoch_selector TINY_MODEL = os.environ.get( "LANGSET_TEST_MODEL", "hf-internal-testing/tiny-random-LlamaForCausalLM" @@ -32,11 +32,18 @@ def _seed() -> None: torch.use_deterministic_algorithms(True, warn_only=True) -def _build_model() -> LangSetModel: +def _build_model(*, bf16: bool = False) -> LangSetModel: _seed() - return LangSetModel.from_pretrained( - TINY_MODEL, bf16=False, device="cpu", multi_latent=True, fsq_dim=32, fsq_levels=8 + model = LangSetModel.from_pretrained( + TINY_MODEL, + bf16=bf16, + device="cpu", + multi_latent=True, + code_emit=True, + n_codes=8, ) + model.head.set_code(torch.randn(8, model.latent_dim)) + return model # 9 rows: seed -> a SET of stage descriptions, with per-item stage labels (phase / recon-ce), a per-ROW scalar @@ -80,6 +87,7 @@ def _args(out_dir: str, **over: object) -> TrainingArguments: val_frac=0.25, seed=0, sup_field="stage", + emission=CodeSoftmaxObjective, output_dir=out_dir, ) d.update(over) @@ -254,9 +262,7 @@ def test_bf16_hidden_and_recon_heads_train() -> None: Head(name="v", reads="hidden", target="value", loss="mse", dim=1, transient=False), Head(name="d", reads="recon", target="dense", loss="mse", dim=1, transient=False), ] - model = LangSetModel.from_pretrained( - TINY_MODEL, bf16=True, device="cpu", multi_latent=True, fsq_dim=32, fsq_levels=8 - ) + model = _build_model(bf16=True) with tempfile.TemporaryDirectory() as td: Trainer( model, _args(td, epochs=2, heads=heads), _dataset() diff --git a/tests/test_bridge_emission_smoke.py b/tests/test_bridge_emission_smoke.py index 626ed8f..f6c59f6 100644 --- a/tests/test_bridge_emission_smoke.py +++ b/tests/test_bridge_emission_smoke.py @@ -104,7 +104,6 @@ def test_query_bridge_emission_folds_hard_negatives_into_the_bank() -> None: emission=QueryBridgeEmission, target_source=FrozenEncoderTarget, lam_multi_nce=0.0, - lam_hard_neg=0.0, # bridge consumes hard_neg_field IN-BANK; separate term stays off hard_neg_field="hard_neg", # <-- pooled adjacent-confusables folded into the bridge's InfoNCE output_dir=out, ) diff --git a/tests/test_code_softmax.py b/tests/test_code_softmax.py index c9c7f9e..96ee55f 100644 --- a/tests/test_code_softmax.py +++ b/tests/test_code_softmax.py @@ -1,13 +1,12 @@ -"""Codebook emission (`code_emit`) โ€” one softmax over a fixed alphabet instead of the FSQ digit grid. +"""Codebook emission (`code_emit`) uses one softmax over a fixed alphabet. -The FSQ default emits `fsq_dim` independent digit softmaxes over a learned grid. `code_emit` replaces that with a -SINGLE softmax over a fixed codebook and commits to the resulting MIXTURE of codes, which is what an emission that +It commits to the resulting MIXTURE of codes, which is what an emission that is a SET (maze frontier: which of 256 cells) actually needs: mass has to be allocated among the members instead of each member being scored independently. These pin the three properties the maze arm depends on: - * the rollout is shape-clean at fsq_dim == 1 and the commit is the mixture, not a single winner; - * KV-cache stays numerically identical on this path too (it must, since the loops are shared with FSQ); + * the rollout is shape-clean and the commit is the mixture, not a single winner; + * KV-cache stays numerically identical on this path; * an orthonormal codebook round-trips: a membership target projects back to the uniform-over-members law, which is what CodeSoftmaxObjective trains toward, and the emitted mixture decodes to the right members. """ @@ -48,7 +47,7 @@ def _members(m, cells: list[list[int]]) -> torch.Tensor: def test_code_emit_shapes_and_mixture_commit() -> None: - """emit_logits collapses fsq_dim to 1 and widens to n_codes; commit returns the softmax mixture of codes.""" + """emit_logits returns one n_codes softmax; commit returns the softmax mixture of codes.""" torch.manual_seed(0) m = _code_model() m.eval() @@ -105,10 +104,10 @@ def test_code_emit_kv_cache_matches_recompute() -> None: 4, 5, dtype=torch.bool ) # self-feed: exercises commit() at each sequential hop with torch.no_grad(): - d0, s0, _, _ = m.rollout_train_codebook( + d0, s0, _, _ = m.rollout_train_state( ids, am, target, ss_prob=0.4, ss_mask=ss_mask, train_hops=3, kv_cache=False ) - d1, s1, _, _ = m.rollout_train_codebook( + d1, s1, _, _ = m.rollout_train_state( ids, am, target, ss_prob=0.4, ss_mask=ss_mask, train_hops=3, kv_cache=True ) assert d0.shape == d1.shape == (4, 6, 1, N_CODES), f"unexpected logit shape {tuple(d0.shape)}" @@ -167,7 +166,7 @@ def test_code_emit_recon_carries_emission_gradient() -> None: """`recon` must be the EMISSION with gradient, not the target. Every auxiliary loss term builds its query from `c.recon` (MultiNCETerm: "emitted, gradient flows here"; - HardNegTerm's continuous branch). The codebook path is lossless, so the rollout's own recon IS the + the auxiliary contrastive terms). The codebook path is lossless, so the rollout's own recon IS the stop-grad target -- handing that back makes in-batch InfoNCE compare each target against itself: near-zero loss, zero gradient, silently inert. Caught in a live run as a `loss - loss_code` gap pinned at 0.145 for five straight epochs while loss_code fell 5.17 -> 3.19.""" diff --git a/tests/test_cot_emit.py b/tests/test_cot_emit.py index 2b8a689..fee54c9 100644 --- a/tests/test_cot_emit.py +++ b/tests/test_cot_emit.py @@ -12,9 +12,9 @@ import torch sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) -import test_trainer_multi_characterization as M # noqa: E402 +import multi_state_helpers as M # noqa: E402 -from langset import Trainer, TrainingArguments # noqa: E402 +from langset import Trainer # noqa: E402 from langset.strategies import CoTGenTerm, MultiStepCtx # noqa: E402 @@ -46,10 +46,7 @@ def _cot_loss(model, padding_side: str) -> float: valid=None, target_lat=None, recon=None, - dim_lg=None, lmax=1, - fsq_levels=8, - lab_label=None, target_source=None, phase_head=None, phase_ids={}, @@ -75,9 +72,7 @@ def test_cot_loss_is_padding_side_invariant() -> None: def _cot_texts_for(rows: list[dict], column_mapping=None) -> list[str]: model = M._build_model() with tempfile.TemporaryDirectory() as td: # bare args: exercise only the __init__ cot_text load - args = TrainingArguments( - epochs=1, batch_size=4, output_dir=td, report_to=None, verbose=False - ) + args = M._args(td, epochs=1, batch_size=4, sup_field=None, lam_sup=0.0, lam_phase=0.0) tr = Trainer(model, args, rows, column_mapping=column_mapping) return tr.cot_texts diff --git a/tests/test_gradcache_multi.py b/tests/test_gradcache_multi.py index 879777b..b4fc924 100644 --- a/tests/test_gradcache_multi.py +++ b/tests/test_gradcache_multi.py @@ -16,14 +16,23 @@ import torch from langset import LangSetModel, Trainer, TrainingArguments +from langset.strategies import CodeSoftmaxObjective ARCH = os.environ.get("LANGSET_TEST_MODEL", "hf-internal-testing/tiny-random-LlamaForCausalLM") def _multi_model(): - return LangSetModel.from_pretrained( - ARCH, device="cpu", dropout=0.0, n_latents=1, multi_latent=True, fsq_dim=8, fsq_levels=4 + m = LangSetModel.from_pretrained( + ARCH, + device="cpu", + dropout=0.0, + n_latents=1, + multi_latent=True, + code_emit=True, + n_codes=4, ) + m.head.set_code(torch.eye(4, m.latent_dim)) + return m def test_shared_ss_mask_full_equals_chunk() -> None: @@ -40,13 +49,11 @@ def test_shared_ss_mask_full_equals_chunk() -> None: ss_mask = torch.rand(B, L) < 0.4 # the shared per-(row, hop) self-feed decisions with torch.no_grad(): - _, _, _, recon_full = m.rollout_train_codebook( - ids, am, target, ss_prob=0.4, ss_mask=ss_mask - ) + _, _, _, recon_full = m.rollout_train_state(ids, am, target, ss_prob=0.4, ss_mask=ss_mask) rows = [0, 1, 2] ri = torch.tensor(rows) with torch.no_grad(): - _, _, _, recon_chunk = m.rollout_train_codebook( + _, _, _, recon_chunk = m.rollout_train_state( ids[ri], am[ri], target[ri], ss_prob=0.4, ss_mask=ss_mask[ri] ) diff = (recon_full[ri] - recon_chunk).abs().max().item() @@ -74,6 +81,7 @@ def test_multi_gradcache_trains() -> None: lr=1e-3, max_len=32, ss_prob=0.25, + emission=CodeSoftmaxObjective, lam_recon=0.0, val_frac=0.125, eval_every=999, diff --git a/tests/test_kv_cache.py b/tests/test_kv_cache.py index 9724149..d8575a6 100644 --- a/tests/test_kv_cache.py +++ b/tests/test_kv_cache.py @@ -1,6 +1,6 @@ """KV-cache rollout exactness (multi-latent). -The KV-cache path (`rollout_train_codebook(..., kv_cache=True)`) forwards the prompt ONCE and then each latent +The KV-cache path (`rollout_train_state(..., kv_cache=True)`) forwards the prompt ONCE and then each latent token alone against the cached prefix K/V, instead of re-running the full growing sequence every tick. Its whole justification is that it is NUMERICALLY IDENTICAL to the recompute path โ€” same emit logits, same feeds โ€” while holding ~1 prompt forward + n single tokens of activations instead of n full-prefix forwards (killing the @@ -22,13 +22,21 @@ def _multi_model(): - return LangSetModel.from_pretrained( - ARCH, device="cpu", dropout=0.0, n_latents=1, multi_latent=True, fsq_dim=8, fsq_levels=4 + m = LangSetModel.from_pretrained( + ARCH, + device="cpu", + dropout=0.0, + n_latents=1, + multi_latent=True, + code_emit=True, + n_codes=4, ) + m.head.set_code(torch.eye(4, m.latent_dim)) + return m def _run(m, ids, am, target, ss_mask, kv_cache, train_hops=None): - return m.rollout_train_codebook( + return m.rollout_train_state( ids, am, target, ss_prob=0.4, ss_mask=ss_mask, kv_cache=kv_cache, train_hops=train_hops ) diff --git a/tests/test_random_init.py b/tests/test_random_init.py index 47e778d..9541c88 100644 --- a/tests/test_random_init.py +++ b/tests/test_random_init.py @@ -97,13 +97,20 @@ def test_from_scratch_trains_through_trainer() -> None: import tempfile sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) - import test_trainer_multi_characterization as M # noqa: E402 + import multi_state_helpers as M # noqa: E402 from langset import Trainer # noqa: E402 + M._seed() m = LangSetModel.from_scratch( - ARCH, device="cpu", multi_latent=True, arch_overrides={"num_hidden_layers": 2} + ARCH, + device="cpu", + multi_latent=True, + arch_overrides={"num_hidden_layers": 2}, + code_emit=True, + n_codes=8, ) + m.head.set_code(torch.randn(8, m.latent_dim)) before = torch.cat([p.detach().flatten() for p in m.parameters() if p.requires_grad]).clone() with tempfile.TemporaryDirectory() as td: Trainer(m, M._args(td), M._dataset()).train() diff --git a/tests/test_replay.py b/tests/test_replay.py index affb81a..7ac3f82 100644 --- a/tests/test_replay.py +++ b/tests/test_replay.py @@ -11,7 +11,7 @@ import pytest import torch -from test_trainer_multi_characterization import _args, _build_model +from multi_state_helpers import _args, _build_model from langset import Trainer from langset.trainer import _replay_ce, _tokenize_replay @@ -51,7 +51,7 @@ def test_replay_ce_is_padding_invariant(): def test_replay_training_runs_multi_latent(): # integration: emit rows + a couple varying-length learn rows, replay firing every step -> the rewired # multi-latent learn path must train end-to-end (exercises _tokenize_replay + _replay_ce on real batches). - from test_trainer_multi_characterization import _flat_trainable + from multi_state_helpers import _flat_trainable m = _build_model() rows = [ @@ -80,9 +80,6 @@ def test_replay_training_runs_multi_latent(): sup_field=None, lam_sup=0.0, hard_neg_field=None, - lam_hard_neg=0.0, - label_dims=None, - lam_label_dims=0.0, ) Trainer(m, args, rows).train() # must not raise; learn_loss fires each step assert ( @@ -111,9 +108,6 @@ def test_all_rows_learn_tagged_errors_clearly(): sup_field=None, lam_sup=0.0, hard_neg_field=None, - lam_hard_neg=0.0, - label_dims=None, - lam_label_dims=0.0, ) with pytest.raises(ValueError): # no emit rows left -> a clear, upfront error Trainer(model, args, _all_learn_rows()).train() diff --git a/tests/test_sigreg_smoke.py b/tests/test_sigreg_smoke.py index 8923ba2..1913dcd 100644 --- a/tests/test_sigreg_smoke.py +++ b/tests/test_sigreg_smoke.py @@ -1,6 +1,5 @@ """SMOKE test for the SIGReg (LeJEPA) anti-collapse target source. -The multi-latent characterization golden (test_trainer_multi_characterization.py) pins the DEFAULT EMA-twin path. SIGReg is a different training path (no EMA twin, live gradient targets, in-batch NCE gated off, an extra isotropic-Gaussian penalty), so it needs its own coverage. This is a fast CPU/fp32 smoke, not a golden: it just proves the injected `target_source=SIGRegTarget` path RUNS, stays finite, and actually changes training vs the @@ -15,9 +14,9 @@ import numpy as np import torch -# reuse the golden harness (tiny model, deterministic seed, 8-row dataset, aux-term-exercising args) +# Reuse the named-state harness (tiny model, deterministic seed, 8-row dataset). sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) -import test_trainer_multi_characterization as M # noqa: E402 +import multi_state_helpers as M # noqa: E402 from langset import Trainer # noqa: E402 from langset.sigreg import SIGReg # noqa: E402 diff --git a/tests/test_state_residual.py b/tests/test_state_residual.py index 44f99df..cf5731c 100644 --- a/tests/test_state_residual.py +++ b/tests/test_state_residual.py @@ -1,7 +1,6 @@ """STATE + RESIDUAL emission โ€” the named half and the unnamed half. -`fsq_dim` controls how finely the latent is quantized; `emb_slots` controls which of its dims are made legible. -Both DECODE meaning out of a latent the twin defined. This third mode CONSTITUTES the latent: its leading dims +This mode CONSTITUTES the latent: its leading dims are a mixture over a named alphabet, its trailing `res_dim` dims are a residual nothing names. The split is what makes adopting a named alphabet safe. An alphabet is a ceiling โ€” whatever it cannot name, the @@ -87,8 +86,8 @@ def test_rollout_returns_emit_hidden_only_when_asked() -> None: enc = m.tokenizer(["row %d text" % i for i in range(3)], padding=True, return_tensors="pt") target = F.normalize(torch.randn(3, 4, m.latent_dim), dim=-1) with torch.no_grad(): - four = m.rollout_train_codebook(enc["input_ids"], enc["attention_mask"], target) - five = m.rollout_train_codebook( + four = m.rollout_train_state(enc["input_ids"], enc["attention_mask"], target) + five = m.rollout_train_state( enc["input_ids"], enc["attention_mask"], target, return_emit_hidden=True ) assert len(four) == 4 and len(five) == 5 diff --git a/tests/test_superposition.py b/tests/test_superposition.py deleted file mode 100644 index 2557c67..0000000 --- a/tests/test_superposition.py +++ /dev/null @@ -1,121 +0,0 @@ -"""REPRODUCING test for the superposition selector bug surfaced in PR #7 review (Copilot comment 2). - -`last_epoch_selector` returns float(ep) so it should keep the FINAL epoch's weights. But checkpoint selection runs -only inside the `if ep % a.eval_every: continue` gate, so when eval_every > 1 and the last epoch isn't an eval -epoch, the selector never sees it and an earlier epoch is restored. Since evaluation is RNG-neutral (verified: with -no restore, eval_every doesn't change the trajectory), the result must be INVARIANT to eval_every under -last_epoch_selector. It currently is not -> RED until the trainer always evaluates the final epoch for such a -selector. Run: .venv/bin/python tests/test_superposition.py -""" - -from __future__ import annotations - -import sys -import tempfile - -import numpy as np - -sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) -import test_trainer_multi_characterization as M # noqa: E402 - -from langset import Trainer # noqa: E402 -from langset.strategies import last_epoch_selector # noqa: E402 - - -def _run(**over) -> np.ndarray: - model = M._build_model() - with tempfile.TemporaryDirectory() as td: - Trainer( - model, - M._args(td, epochs=5, verbose=False, selector=last_epoch_selector, **over), - M._dataset(), - ).train() - return M._flat_trainable(model) - - -def test_last_epoch_selector_invariant_to_eval_every() -> None: - """last_epoch_selector keeps the final epoch, so its result cannot depend on the eval cadence. epochs=5 with - eval_every=3 evaluates ep0,ep3 and would restore ep3 (not the true final ep4); eval_every=1 restores ep4.""" - every1 = _run(eval_every=1) # final epoch (ep4) is always an eval epoch -> correct - every3 = _run(eval_every=3) # ep4 skipped -> restores ep3 unless the trainer is fixed - max_delta = float(np.max(np.abs(every1 - every3))) - assert max_delta == 0.0, ( - f"last_epoch_selector result depends on eval_every (max param delta {max_delta:.2e}) โ€” the final epoch is " - f"skipped by the eval gate, so an earlier epoch is kept instead of the last" - ) - - -def test_direct_superposition_smoke() -> None: - """The DIRECT way to teach superposition (what the maze example does): each target text describes the SET of - next states, trained with last_epoch_selector; read back with return_soft, whose per-latent entropy is the - calibrated-uncertainty readout. Verify it trains finite and exposes that entropy.""" - model = M._build_model() - rows = [ - {"input_text": "a fair coin is flipped", "target_texts": ["it lands heads or tails"]}, - { - "input_text": "a six-sided die is rolled", - "target_texts": ["it shows one, two, three, four, five, or six"], - }, - { - "input_text": "water at sea level is heated past its boiling point", - "target_texts": ["it turns to steam"], - }, - ] * 3 - with tempfile.TemporaryDirectory() as td: - Trainer( - model, - M._args( - td, - epochs=2, - verbose=False, - selector=last_epoch_selector, - sup_field=None, - lam_sup=0.0, - hard_neg_field=None, - lam_hard_neg=0.0, - label_dims=None, - lam_label_dims=0.0, - ), - rows, - ).train() - p = M._flat_trainable(model) - assert np.isfinite(p).all(), "direct-superposition training went non-finite" - lat, lengths, soft, ent = model.rollout("a six-sided die is rolled", return_soft=True) - assert ent.shape[0] == lat.shape[0], ( - "return_soft must expose one entropy value per emitted latent" - ) - - -def test_snapshot_every_is_one_based() -> None: - """Copilot #3: snapshot_every=N snapshots AFTER epochs N, 2N, ... (1-based) โ€” not at ep0. epochs=4, N=2 -> only - {output_dir}_ep2 and _ep4 exist; _ep0 must not.""" - import os - - model = M._build_model() - with tempfile.TemporaryDirectory() as td: - out = f"{td}/run" - Trainer( - model, - M._args(td, epochs=4, verbose=False, output_dir=out, snapshot_every=2), - M._dataset(), - ).train() - assert not os.path.isdir(f"{out}_ep0"), "snapshot fired at ep0 (should be 1-based)" - assert os.path.isdir(f"{out}_ep2") and os.path.isdir(f"{out}_ep4"), ( - "expected snapshots after epochs 2 and 4" - ) - - -if __name__ == "__main__": - import torch - - torch.use_deterministic_algorithms(True, warn_only=True) - for name in ( - "test_last_epoch_selector_invariant_to_eval_every", - "test_direct_superposition_smoke", - "test_snapshot_every_is_one_based", - ): - try: - globals()[name]() - print(f"{name} PASS") - except AssertionError as e: - print(f"{name} FAIL -> {e}") diff --git a/tests/test_trainer_multi_characterization.py b/tests/test_trainer_multi_characterization.py deleted file mode 100644 index 2554698..0000000 --- a/tests/test_trainer_multi_characterization.py +++ /dev/null @@ -1,157 +0,0 @@ -"""GOLDEN characterization tests for the MULTI-LATENT (FSQ latent-set) trainer. - -Purpose: pin the EXACT numerics of `Trainer._train_multi` so the strategy-seams refactor (extracting the -emission objective / target source / aux loss terms into swappable strategies) is provably MECHANICALLY -IDENTICAL when wired with its DEFAULT strategies. The single-latent golden (test_trainer_characterization.py) -does NOT cover this path โ€” all the feature branching lives here. - -The golden exercises, in one deterministic run, every DEFAULT code path the refactor touches: - * FSQ emission objective (per-dim digit CE + folded STOP + cosine recon) - * EMA-twin target source (stop-grad targets + ema_update) - * the aux LossTerms that are on by default or lit by these args: multi_nce (+ its identical-text - negative mask), supcon, phase head, hard negatives -So an accidental change to ANY of those terms / their weighting / their order fails loudly. - -Everything runs on CPU/fp32 with a tiny model so it's deterministic and fast. - -Regenerate the golden ONLY when you intentionally change behavior: - .venv/bin/python tests/test_trainer_multi_characterization.py --update -""" - -from __future__ import annotations - -import os -import sys -import tempfile -from pathlib import Path - -import numpy as np -import torch - -from langset import LangSetModel, Trainer, TrainingArguments - -HERE = Path(__file__).parent -GOLDEN = HERE / "golden_trainer_multi.npz" -TINY_MODEL = os.environ.get( - "LANGSET_TEST_MODEL", "hf-internal-testing/tiny-random-LlamaForCausalLM" -) -RTOL, ATOL = ( - 1e-4, - 1e-5, -) # atol loose enough for cross-platform float reassociation (macOS golden vs Linux CI) - - -def _seed() -> None: - torch.manual_seed(0) - np.random.seed(0) - torch.use_deterministic_algorithms(True, warn_only=True) - - -def _build_model() -> LangSetModel: - _seed() # seed BEFORE building: LoRA-A / FSQ head init are random -> must be pinned - return LangSetModel.from_pretrained( - TINY_MODEL, bf16=False, device="cpu", multi_latent=True, fsq_dim=32, fsq_levels=8 - ) - - -def _dataset() -> list[dict]: - """8 rows, each a seed -> a SET of target descriptions, with per-item stage labels (sup/phase) and a - per-row hard negative โ€” so one run lights multi_nce, supcon, phase, and hard_neg together.""" - topics = ["late payment", "theft", "property tax"] - stages = ["filing", "hearing", "ruling"] - rows = [] - for i in range(8): - t = topics[i % 3] - rows.append( - { - "input_text": f"case {i}: a dispute concerning {t}", - "target_texts": [f"{s}: the {t} matter at the {s} stage" for s in stages], - "stage": list(stages), # per-item labels (sup_field / phase) - "hardneg": [f"an unrelated {topics[(i + 1) % 3]} document"], - } - ) - return rows - - -def _args(out_dir: str, **over: object) -> TrainingArguments: - d: dict[str, object] = dict( - epochs=2, - batch_size=4, - lr=1e-3, - max_len=64, - report_to=None, - verbose=False, - eval_every=99, - patience=99, - val_frac=0.25, - seed=0, # no eval/early-stop -> deterministic weights - sup_field="stage", - lam_sup=0.2, - lam_phase=0.1, # light: light up supcon + phase-head terms - hard_neg_field="hardneg", - lam_hard_neg=0.2, # light up the hard-negative term - label_dims={"stage": [1]}, - lam_label_dims=0.3, # light up the FSQ label-subspace term (reserved dim 1) - output_dir=out_dir, - ) - d.update(over) - return TrainingArguments(**d) # type: ignore[arg-type] - - -def _flat_trainable(model: LangSetModel) -> np.ndarray: - return np.concatenate( - [ - p.detach().float().cpu().numpy().ravel() - for name, p in sorted(model.named_parameters(), key=lambda kv: kv[0]) - if p.requires_grad - ] - ).astype(np.float64) - - -def _run(**arg_over: object) -> dict[str, np.ndarray]: - model = _build_model() - with tempfile.TemporaryDirectory() as td: - Trainer(model, _args(td, **arg_over), _dataset()).train() - return {"params_post": _flat_trainable(model)} - - -def generate_golden() -> None: - g = _run() - np.savez(GOLDEN, **g) - print(f"[golden] wrote {GOLDEN} (params={g['params_post'].size})") - - -def test_train_multi_identity() -> None: - assert GOLDEN.exists(), "golden missing -> run with --update" - g = np.load(GOLDEN) - out = _run() - np.testing.assert_allclose( - out["params_post"], - g["params_post"], - rtol=RTOL, - atol=ATOL, - err_msg="MULTI-LATENT train math changed (FSQ objective / EMA target / aux terms differ)", - ) - - -def test_golden_is_sensitive_to_nce() -> None: - """Prove the golden actually GUARDS behavior: dropping the multi_nce term (a real math change) must - diverge from the golden beyond tolerance, so the identity test would catch a silent term change.""" - assert GOLDEN.exists(), "golden missing -> run with --update" - g = np.load(GOLDEN) - out = _run(lam_multi_nce=0.0) - max_delta = float(np.max(np.abs(out["params_post"] - g["params_post"]))) - assert max_delta > 1e-4, ( - f"golden is NOT sensitive to multi_nce (max param delta {max_delta:.2e}) -> tighten the snapshot" - ) - - -if __name__ == "__main__": - if "--update" in sys.argv: - generate_golden() - else: - test_train_multi_identity() - print("train_multi_identity OK") - test_golden_is_sensitive_to_nce() - print("golden_is_sensitive OK") - print("ALL PASS")