tribunal-kit 9.2.3 → 9.2.5

package/.agent/skills/train-sentence-transformers/references/hf_jobs_execution.md

@@ -0,0 +1,173 @@
+# Hugging Face Jobs Execution
+
+Run training on Hugging Face's managed GPUs without provisioning any local infrastructure. The same training script runs locally and on Jobs. This reference covers only the Jobs-specific concerns.
+
+## Prerequisites
+
+- Hugging Face account with a **Pro, Team, or Enterprise** plan. Jobs are paid.
+- `HF_TOKEN` with **write** permission. Log in once locally with `hf auth login` (the modern command from the `hf` CLI, replacing the deprecated `huggingface-cli login`).
+- Access to the `hf_jobs()` MCP tool, or the `hf` CLI (`curl -LsSf https://hf.co/cli/install.sh | bash -s`).
+
+## The three submission paths
+
+### 1. Inline script via MCP (recommended in Claude Code)
+
+Pass the full training script as `script`. Dependencies come from the PEP 723 header.
+
+```python
+hf_jobs("uv", {
+    "script": """
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["sentence-transformers[train]>=5.0", "trackio"]
+# ///
+
+# <full training script content>
+""",
+    "flavor": "a10g-large",
+    "timeout": "3h",
+    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
+})
+```
+
+### 2. Script-from-URL via MCP
+
+Upload the script to the Hub (as a model or dataset repo file) or a Gist, then reference by URL:
+
+```python
+hf_jobs("uv", {
+    "script": "https://huggingface.co/USERNAME/scripts/resolve/main/train_bi_encoder.py",
+    "flavor": "a10g-large",
+    "timeout": "3h",
+    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
+})
+```
+
+Local file paths (`./train.py`, `/path/to/train.py`) **do not work**. Jobs run in isolated containers without access to your filesystem.
+
+### 3. CLI
+
+```bash
+hf jobs uv run \
+    --flavor a10g-large \
+    --timeout 3h \
+    --secrets HF_TOKEN \
+    "https://huggingface.co/USERNAME/scripts/resolve/main/train.py"
+```
+
+Syntax gotchas:
+- Command order is `hf jobs uv run`, **not** `hf jobs run uv`.
+- Flags (`--flavor`, `--timeout`, `--secrets`) go **before** the script URL.
+- `--secrets` (plural), not `--secret`.
+
+## Required script modifications for Jobs
+
+Add these to your `TrainingArguments`:
+
+```python
+args = SentenceTransformerTrainingArguments(
+    ...,
+    push_to_hub=True,
+    hub_model_id="your-username/my-model",
+    hub_strategy="every_save",        # push each checkpoint; timeout-safe
+    save_strategy="steps",
+    save_steps=0.1,                   # 10 saves/pushes per epoch; scales with dataset size
+)
+```
+
+Why each matters:
+
+| Argument | Why |
+|---|---|
+| `push_to_hub=True` | The Jobs container is destroyed after the job finishes. Without Hub push, all weights are lost. |
+| `hub_model_id` | Required to identify the destination repo. |
+| `hub_strategy="every_save"` | Default, but worth being deliberate about on Jobs: each checkpoint is pushed as it's written, so a timeout leaves all completed checkpoints on the Hub. `"end"` only pushes once `trainer.train()` returns, so a timeout loses everything. |
+| `save_strategy="steps"` + `save_steps=0.1` | Checkpoints must actually be saved for `hub_strategy="every_save"` to push them. Fractional `0.1` = save every 10% of training, auto-scales with dataset size. |
+
+## Secrets
+
+Secrets are environment variables injected into the Jobs container. They never appear in logs and are not part of the script.
+
+| Secret | Required when |
+|---|---|
+| `HF_TOKEN` | Always, for Hub push. Also covers Trackio auth. |
+| `WANDB_API_KEY` | Using `report_to="wandb"`. |
+| `MLFLOW_TRACKING_URI`, `MLFLOW_TRACKING_TOKEN` | Using MLflow with a remote server. |
+
+The `$HF_TOKEN` syntax in the job config references the value from your local environment at submission time. The literal string `$HF_TOKEN` is replaced with your token's value. Never hardcode tokens in the script itself.
+
+Trackio (the default tracker in this skill) uses `HF_TOKEN` for auth, so no extra secrets are needed. Only switch to the W&B / MLflow rows above if you're using those trackers.
+
+## Timeout
+
+Default is **30 minutes**, which is too short for almost any real training. Set explicitly:
+
+```python
+"timeout": "2h"       # 2 hours
+"timeout": "90m"      # 90 minutes
+"timeout": "1.5h"     # 90 minutes
+"timeout": 7200       # seconds, as integer
+```
+
+Rule: **estimated training time × 1.3**. The extra buffer covers model loading, dataset caching, checkpoint saving, and Hub push.
+
+On timeout, the container is killed immediately. Only data on the Hub (`hub_strategy="every_save"` saves you here) or in persistent volumes survives.
+
+## Dataset caching
+
+Hugging Face datasets are cached at `~/.cache/huggingface/datasets` by default. That's **inside the container**, which is destroyed after the job. Each Jobs run re-downloads the dataset.
+
+For large datasets (>5 GB), this matters. Options:
+
+- **Persistent `/data` volume** (Jobs feature, check current documentation): set `HF_DATASETS_CACHE=/data/datasets` so caches persist across jobs.
+- **Pre-cache locally, push to Hub**: if the dataset is on Hub already, nothing to do. If it's local-only, `dataset.push_to_hub(...)` once so subsequent jobs load from Hub.
+
+## Monitoring a running job
+
+```bash
+hf jobs ps [--all]                        # running (or all) jobs
+hf jobs inspect <job-id>                  # full config + status
+hf jobs logs <job-id> [--follow|--tail N] # tail or stream
+hf jobs cancel <job-id>
+hf jobs hardware                          # list flavors + hourly rates
+```
+
+`hf jobs logs <id> --follow` under `Bash run_in_background` pairs nicely with a `Monitor` watching for the `VERDICT:` line emitted by your training script's verdict block.
+
+MCP equivalents (signatures may vary by server version, so check the actual
+tool listing): `hf_jobs("ps")`, `hf_jobs("logs", {"job_id": ...})`,
+`hf_jobs("cancel", {"job_id": ...})`.
+
+For recurring runs, `hf jobs scheduled uv run "<cron>" <script> ...`
+schedules. `hf jobs scheduled ps/suspend/delete` manages.
+
+## Common failures
+
+### "Model not found on Hub" after a successful-looking run
+
+The run succeeded but `push_to_hub` was not enabled. The container is gone. The weights are gone.
+
+Fix: always set `push_to_hub=True` + `hub_model_id=...` + `secrets={"HF_TOKEN": "$HF_TOKEN"}`.
+
+### Tracker not connecting
+
+- **Trackio:** `HF_TOKEN` missing or lacks write permission. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN"}` and make sure the token has write access.
+- **W&B:** `WANDB_API_KEY` missing. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN", "WANDB_API_KEY": "$WANDB_API_KEY"}`.
+
+### OOM on first step
+
+Flavor too small. Move up one tier (see `hardware_guide.md`).
+
+### Training starts but eval hangs forever
+
+`eval_strategy="steps"` with no `eval_dataset`. Always provide an eval dataset, or set `eval_strategy="no"`.
+
+### Dataset download times out
+
+Large dataset or slow cold-cache. Increase `timeout` or pre-cache to a persistent volume.
+
+### `CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True` crash
+
+The cached losses are incompatible with gradient checkpointing. Disable `gradient_checkpointing`.
+
+After submission, the MCP returns a job ID. Monitor with `hf_jobs("logs", {"job_id": ...})` when you want an update. Don't poll in a tight loop. End-to-end submission templates live in `scripts/train_sentence_transformer_example.py` / `scripts/train_cross_encoder_example.py` / `scripts/train_sparse_encoder_example.py`. Wrap the script contents in the inline pattern from §1 above.

package/.agent/skills/train-sentence-transformers/references/losses_cross_encoder.md

@@ -0,0 +1,137 @@
+# Cross-Encoder Losses (reranker)
+
+All losses live in `sentence_transformers.cross_encoder.losses`.
+
+Cross-encoder losses fall into three families: **pointwise** (score each pair independently), **pairwise** (score pairs of pairs), and **listwise** (score a whole ranked list at once). Plus contrastive/distillation variants for specific setups.
+
+## Top-line decision table
+
+| You have | Use | Family |
+|---|---|---|
+| `(query, passage, label)` with `label ∈ {0, 1}` or `[0, 1]` | `BinaryCrossEntropyLoss` | Pointwise |
+| `(query, passage, class_id)` multi-class | `CrossEntropyLoss` | Pointwise |
+| `(query, passage)` with implicit positives, want contrastive | `CachedMultipleNegativesRankingLoss` | Contrastive |
+| `(query, passages, labels)` (one row per query, with parallel lists of candidate passages and their relevance scores) | `LambdaLoss` | Listwise |
+| Same listwise shape, want well-tested simpler loss | `ListNetLoss` or `ListMLELoss` | Listwise |
+| `(query, positive, negative)` pairwise | `RankNetLoss` | Pairwise |
+| `(query, passage, teacher_score)` distillation from stronger reranker | `MSELoss` or `MarginMSELoss` | Distillation |
+
+## Pointwise losses
+
+### `BinaryCrossEntropyLoss`
+
+**The default cross-encoder loss.** For `(query, passage, label)` where label is 0/1 or a continuous [0, 1] relevance score.
+
+```python
+loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(5.0))
+```
+
+- `pos_weight`: upweight positives when your dataset is imbalanced (e.g. 1 positive + N hard negatives). A good default is `pos_weight = num_hard_negatives`.
+- Handles both binary labels and continuous relevance scores. If you have graded relevance (0.0, 0.5, 1.0), BCE still works.
+
+### `CrossEntropyLoss`
+
+Multi-class classification over passages. Use with `CrossEncoder("...", num_labels=N)` where N is the number of classes (e.g. 3 for NLI: entailment / neutral / contradiction).
+
+## Contrastive losses (for reranker training)
+
+### `MultipleNegativesRankingLoss` (cross-encoder)
+
+Contrastive training mirroring bi-encoder MNRL. Applies in-batch negatives: for each `(query, positive)`, every other `positive` in the batch is used as a negative.
+
+- **Data**: `(query, positive)` or `(query, positive, negative_0, negative_1, ...)`.
+- For CrossEncoder training, **prefer `BinaryCrossEntropyLoss` with mined hard negatives** as the default. MNRL is the fallback when you only have `(query, positive)` pairs and want to avoid a separate mining pass.
+
+### `CachedMultipleNegativesRankingLoss`
+
+Cached variant (GradCache). Decouples per-device batch size from effective in-batch negatives, same as the bi-encoder cached version.
+
+```python
+loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)
+```
+
+- Incompatible with `gradient_checkpointing=True`.
+- Key choice for training rerankers with effective batch 256+ on a single GPU.
+
+## Distillation losses
+
+> **`activation_fn=nn.Identity()` is mandatory** for all distillation, listwise, and pairwise losses below. Only `BinaryCrossEntropyLoss` and `CrossEntropyLoss` tolerate the default `Sigmoid`. The loss sees raw logits during training, but the model's `activation_fn` is applied at evaluation via `predict()`. The default `Sigmoid` (with `num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently collapsing eval ranking (training loss looks healthy while nDCG crashes from ~0.59 to ~0.14). Construct as `CrossEncoder("...", num_labels=1, activation_fn=torch.nn.Identity())`. See `troubleshooting.md` ("CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training") for the failure-mode walkthrough.
+
+### `MSELoss` (cross-encoder)
+
+Regress the student's output score to the teacher's score. Construct the model with `activation_fn=nn.Identity()` (see callout above).
+
+- **Data**: `(query, passage, teacher_score)`.
+- The teacher is typically a larger/stronger cross-encoder. Precompute its scores once, store as the label.
+
+### `MarginMSELoss` (cross-encoder)
+
+Regress the **difference** between positive and negative scores against the teacher's margin. Typically gives better distillation than plain MSE.
+
+- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
+- Popular recipe for MS MARCO-style distillation.
+
+## Listwise losses
+
+All listwise losses expect the dataset to have per-query **lists** of candidate passages and their scores, typically via a collator that groups rows by query.
+
+> Listwise and pairwise losses below also require `activation_fn=nn.Identity()` (see the Distillation-section callout above).
+
+### `LambdaLoss`
+
+The state-of-the-art listwise ranking loss. Optimizes a surrogate of nDCG via weighted pairwise comparisons.
+
+**Data shape**: `(query, [doc1...docN], [score1...scoreN])` per row. One query, a list of candidate documents, and a parallel list of relevance scores. Use `mine_hard_negatives(..., output_format="labeled-list", ...)` to build this from `(query, positive)` pairs. For ID-only datasets with separate text splits (e.g. `lightonai/ms-marco-en-bge`), `sentence_transformers.util.resolve_ids(..., output_format="lists")` produces this nested shape directly.
+
+```python
+import torch.nn as nn
+from sentence_transformers.cross_encoder.losses import LambdaLoss, NDCGLoss2PPScheme
+
+model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
+loss = LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())
+```
+
+- Strong default when you have multiple candidates per query and graded relevance.
+- `weighting_scheme`: `NDCGLoss2PPScheme` (default), `NDCGLoss2Scheme`, `LambdaRankScheme`. The `NDCGLoss2PPScheme` default was shown to reach the strongest performance in the original LambdaLoss paper.
+
+#### LambdaLoss-specific operational notes
+
+- **OOM recovery order**: drop `mini_batch_size` first (chunking inside the loss preserves the K-list semantic), then `per_device_train_batch_size` paired with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list length) only as a last resort. Lowering K changes the experiment. The others don't.
+- **Tiny loss at large K is expected.** With `NDCGLoss2PPScheme`, the loss normalizes by the discount-weighted pair count. At K=128 the loss can scale to ~`1e-4` numerically. That's not "training broken". Read `eval_NanoBEIR_R100_mean_ndcg@10` (or your equivalent) instead of the training loss to judge progress.
+- For very large K (>=128), the O(K²) weight buffer per query is materialized outside of the forward chunking, so even small `mini_batch_size` may not be enough. Consider scoring strategy: top-K hard negatives often beat random K.
+
+### `ListNetLoss`
+
+Treats the ranking as a probability distribution (via softmax) and minimizes cross-entropy against the teacher distribution.
+
+### `ListMLELoss`
+
+Maximum-likelihood estimation over permutations. Simpler than LambdaLoss, decent default.
+
+### `PListMLELoss`
+
+Position-aware ListMLE. Weights higher-ranked items more heavily. Often outperforms plain ListMLE on top-k metrics.
+
+### `RankNetLoss`
+
+Pairwise classification: for each pair of candidates, predict which is ranked higher via cross-entropy.
+
+- Simpler and faster than LambdaLoss.
+- Scales quadratically with list length. Not great for long candidate lists (>20).
+
+### `ADRMSELoss`
+
+Alternative listwise formulation (Approx Discounted Rank MSE) from the Rank-DistiLLM paper. Same data shape as `LambdaLoss`. In practice LambdaLoss is the stronger default. `RankNetLoss` was reported as marginally more effective (~0.002 nDCG@10) than ADRMSELoss on the paper's LLM-distillation setup, and LambdaLoss generally beats both.
+
+Hard-negative mining is essential for any contrastive reranker (random negatives teach nothing). See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.
+
+## Gotchas
+
+- **`BinaryCrossEntropyLoss` without `pos_weight`** when you have 5+ hard negatives per positive: the loss under-weights the positive signal. Set `pos_weight=num_hard_negatives`.
+- **`CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
+- **Listwise losses with wildly different list lengths per query**: some losses don't handle ragged lists well. Pad or truncate to a fixed length.
+- **`MarginMSELoss` without precomputed teacher score diffs**: this loss needs the `score_diff` label column populated from a teacher pass. It does not compute the teacher internally.
+- **Training a cross-encoder with `num_labels=1` then using `CrossEntropyLoss`**: mismatch. BCE is for num_labels=1, CE is for num_labels>=2.
+- **Default `Sigmoid` activation under any non-BCE loss**: silently destroys eval ranking. Pass `activation_fn=torch.nn.Identity()` to `CrossEncoder(...)` for distillation, listwise, and pairwise losses (everything except BCE/CE). See the callouts in the Distillation and Listwise sections.
+- **Custom CE head writing to the wrong feature key**: a custom scoring head must populate `features["scores"]` (not `features["sentence_embedding"]`). Otherwise `CrossEncoder.predict()` raises `KeyError: 'scores'` at inference time, even though training succeeds.
+- **Loading a CE with a custom-class head fails from a different script**: if you defined a `class ClassifierHead(nn.Module)` inline in `train.py` and saved the model, `modules.json` records `__main__.ClassifierHead`. Loading via `CrossEncoder("path")` from any other entry point raises `ImportError: Module '__main__' does not define a 'ClassifierHead' attribute`. Either move the class into an importable file (`my_pkg/heads.py`), build the same shape from stock ST modules (`Dense + LayerNorm + Dense`), or document that the model is only loadable from the same script.

package/.agent/skills/train-sentence-transformers/references/losses_multi_vector_encoder.md

@@ -0,0 +1,100 @@
+# Multi-Vector-Encoder Losses (ColBERT / late-interaction)
+
+All losses live in `sentence_transformers.multi_vector_encoder.losses`.
+
+Multi-vector models emit **one embedding per token** and score `(query, document)` via MaxSim (for each query token, take the max similarity to any document token, then sum across query tokens). This is fundamentally different from single-vector cosine, which changes what "temperature" means: MaxSim is an unbounded sum over query-token similarities, so for it a scale near 1.0 makes sense. **The default `scale=1.0` (temperature=1.0) is correct for unnormalized MaxSim: do not copy `scale=20.0` from bi-encoder MNRL.** With MeanMaxSim scoring (`similarity_fct=mean_colbert_scores`, and set `model.similarity_fn_name = "meanmaxsim"` so evaluation matches), each score is divided by its query's token count, so a scale of roughly the average length is a reasonable start, although some works use larger scales.
+
+## Top-line decision table
+
+| You have | Use |
+|---|---|
+| `(anchor, positive)` or `(anchor, positive, negative)` triplets | `MultiVectorMultipleNegativesRankingLoss` |
+| Same, want effective batch size of 128+ | `CachedMultiVectorMultipleNegativesRankingLoss` |
+| Cross-encoder teacher scores, `(query, positive, negative, score_diff)` | `MultiVectorMarginMSELoss` |
+| Listwise distillation `(query, [doc_1..doc_N], scores)` | `MultiVectorDistillKLDivLoss` |
+
+Hard-negative mining is essential for competitive results, because random in-batch negatives leave a lot on the table for late-interaction models. See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.
+
+## Contrastive losses
+
+### `MultiVectorMultipleNegativesRankingLoss`
+
+The default late-interaction contrastive loss. In-batch positives plus explicit hard negatives, scored with MaxSim (or XTR).
+
+```python
+from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss
+
+loss = MultiVectorMultipleNegativesRankingLoss(model=model)  # scale=1.0 (temperature=1.0) is the correct default
+```
+
+- **Data**: `(anchor, positive)` or `(anchor, positive, negative_1, ..., negative_n)`. The collator stamps each column's task (column 0 becomes query, others become document). Pass `task=...` on a column to override.
+- Set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args (same reason as bi-encoder MNRL).
+- `similarity_fct=colbert_scores` by default. Pass `XTRScores(top_k=...)` for the sparser XTR-style scoring (from `sentence_transformers.multi_vector_encoder.scoring`). XTR applies to training only: evaluation always scores with MaxSim (see Gotchas).
+- `scale=1.0` (temperature=1.0) matches PyLate and is correct for unnormalized MaxSim. **Do NOT copy `scale=20.0` from bi-encoder MNRL.** With MeanMaxSim (`similarity_fct=mean_colbert_scores`), raise the scale to roughly the average query length instead, which is the same objective at the same strength.
+
+### `CachedMultiVectorMultipleNegativesRankingLoss`
+
+GradCache variant: chunked embedding forward, cached gradients, and a second re-embedding pass. Decouples per-device batch size from effective in-batch negatives.
+
+```python
+loss = CachedMultiVectorMultipleNegativesRankingLoss(
+    model=model,
+    mini_batch_size=8,           # or use mini_batch_num_tokens=... for token-budgeted packing
+    score_mini_batch_size=4,     # optional, smaller trims the transient (Q, Q*N, q_tok, d_tok) buffer
+)
+```
+
+- **Incompatible with `gradient_checkpointing=True`** (same as every `Cached*` loss).
+- `mini_batch_num_tokens=N` packs sequences until N real tokens per mini-batch instead of a fixed `mini_batch_size`. Big win on variable-length data with flash-attention / input flattening.
+- `score_mini_batch_size` chunks the SCORING phase (which builds `(Q, Q*N, q_tokens, d_tokens)` intermediates) independently from the embedding phase. Drop it first when hitting OOM in the loss stage.
+- `chunk_elements` is the second lever on that same phase, and it cuts the other axis: `score_mini_batch_size` chunks queries, `chunk_elements` chunks documents inside each MaxSim call. Pass it through the scorer, e.g. `similarity_fct=partial(colbert_scores, chunk_elements=1_000_000)`. Scores and gradients are unchanged (it is a pure memory knob), and it composes with `score_mini_batch_size`.
+- `gather_across_devices=True` gathers document embeddings across DDP ranks. Enables cross-rank in-batch negatives.
+
+## Distillation losses
+
+Distillation is where multi-vector models learn most efficiently: cross-encoder teachers (e.g. `gte-modernbert-base`) score `(query, doc)` pairs offline, and the student MaxSim model regresses to that signal.
+
+For the standard MS MARCO KD format (LightOn's `ms-marco-en-bge` and similar: `(query_id, document_ids, scores)` with separate `queries` / `documents` splits), you can use `sentence_transformers.util.resolve_ids` to join the IDs against the text splits at load time. It's a factory that returns a batched transform: pass a mapping from each input ID column to its lookup dataset, e.g. `dataset.set_transform(resolve_ids({"query_id": queries, "document_ids": documents}, max_list_length=32))` (lazy, no caching) or via `dataset.map(..., batched=True, remove_columns=["query_id", "document_ids"])` (eager, cached: `map` keeps input columns unless told to drop them, so pass the ID columns to `remove_columns`), yielding `(query, document_1, ..., document_32, scores)` rows ready for `MultiVectorDistillKLDivLoss`. The same mapping shape covers triplets with IDs (`{"query_id": queries, "positive_id": documents, "negative_id": documents}`) and ID-only contrastive datasets without a `scores` column.
+
+### `MultiVectorMarginMSELoss`
+
+Regress the **margin** between positive and negative MaxSim scores against the teacher's margin.
+
+```python
+loss = MultiVectorMarginMSELoss(model=model)
+```
+
+- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`. Raw teacher scores `[score(q, pos), score(q, neg_1), ...]` also work, converted to margins internally.
+- Popular recipe from PyLate / colpali-engine.
+- Teacher scores are precomputed once, stored as the label column. The loss does not run the teacher inline.
+- Same `similarity_fct` convention as the other losses, but pairwise-shaped: defaults to `colbert_scores_pairwise`, or pass `xtr_scores_pairwise` for XTR.
+
+### `MultiVectorDistillKLDivLoss`
+
+Listwise KL-div: student's softmax distribution over N candidates should match the teacher's.
+
+```python
+loss = MultiVectorDistillKLDivLoss(model=model)
+```
+
+- **Data**: `(query, document_1, ..., document_N, scores)` where `scores` is a list of N teacher scores per row. One column per candidate document (the standard positional multi-column convention), and the label column must use a recognized name (`label`, `labels`, `score`, `scores`). `resolve_ids` expands a stored ID list into the numbered columns.
+- Stronger training signal than `MarginMSE` when you have full `N`-way teacher scores (not just positive/negative margins).
+- `similarity_fct` defaults to `colbert_kd_scores`, the listwise KD variant, not the `colbert_scores` used by the MNRL family. `XTRKDScores` is the XTR counterpart.
+- `temperature` softens both distributions before the softmax, and can be overridden per side with `student_temperature` / `teacher_temperature`. Set `teacher_temperature` from the spread of your teacher's scores rather than from a value copied out of a paper: a float32 softmax underflows to exact zeros once a row's spread divided by the temperature exceeds about 100, and every candidate that underflows drops out of the KL, which is the ranking information you are distilling. Sharpening `student_temperature` far below the student's own score spread collapses that side too, and once both are one-hot the loss and its gradient are exactly zero. **Start at 1.0 / 1.0 and only move a side you have a measured reason to move.**
+- The loss is scaled by the student temperature squared, which keeps gradient magnitudes comparable only for a shared temperature at or above 1.0 (the regime Hinton et al. cover). Below that the factor shrinks the loss instead, by 1e-6 at `student_temperature=0.001`: when combining KD with a contrastive loss, scale the KD weight back up accordingly.
+- **OOM**: drop `per_device_train_batch_size` first and raise `gradient_accumulation_steps` to hold the effective batch. Only reduce `N` (candidate-list length) as a last resort, since lowering N changes the experiment.
+
+## Data-shape gotchas
+
+- **Column 0 is always the query.** Losses call `self.model(sf, task="query" if idx==0 else "document")`. Use the collator's `router_mapping` to override if you need a non-standard column layout.
+- **Cross-column varlen `T`** is handled: the batch's positive and negative columns can have different token counts (Qwen2-VL family etc.). `stack_padded_token_embeddings` pads to the per-batch max before MaxSim.
+- **Mask flow**: MVE reads the SCORING mask from the model OUTPUT dict (`MultiVectorMask` rewrites the input mask). Custom loss code that reads `sentence_features[i]["attention_mask"]` directly instead of `outputs[i]["attention_mask"]` will silently score against skiplisted tokens.
+
+## Gotchas
+
+- **`scale=20.0` copied from bi-encoder MNRL**: misscaled for the unnormalized MaxSim sum. Keep `scale=1.0` there. Only length-normalized MeanMaxSim wants a larger scale, roughly the average query length.
+- **Missing `Normalize` at the token level in the pipeline**: `colbert_scores` assumes L2-normalized token embeddings. If your custom pipeline drops the token-level `Normalize`, either add one or pass `normalize_embeddings=True` semantics via a wrapper.
+- **`CachedMultiVectorMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
+- **`MultiVectorMarginMSELoss` without precomputed `score_diff`**: label column must be populated from a teacher pass ahead of training. The loss does not compute the teacher inline.
+- **Expecting XTR scoring at eval**: `XTRScores` is a train-only `similarity_fct`. XTR takes its top-k across the whole candidate set, so a `(query, document)` pair has no standalone score, which means it cannot be the model's `similarity_fn_name` and the evaluators reject it outright. Evaluation and inference score with MaxSim, also for XTR-trained models. That is by design, not a mismatch to fix. To score a fixed candidate set ad hoc, call `xtr_scores` directly.
+- **Distillation with a weak teacher**: multi-vector students easily match a small-model teacher and then plateau. Use a strong cross-encoder (e.g. `gte-modernbert-base`, `mxbai-rerank-large-v2`) for the teacher pass.

package/.agent/skills/train-sentence-transformers/references/losses_sentence_transformer.md

@@ -0,0 +1,246 @@
+# Bi-Encoder Losses (SentenceTransformer)
+
+All losses live in `sentence_transformers.sentence_transformer.losses`.
+
+Losses are grouped by data shape. The #1 rule: **pick a loss that matches your data**, not the other way around.
+
+## Top-line decision table
+
+| You have | Use |
+|---|---|
+| `(anchor, positive)` pairs | `MultipleNegativesRankingLoss` (or Cached variant for large batches) |
+| `(anchor, positive, negative)` triplets | `MultipleNegativesRankingLoss` (it handles triplets natively) |
+| `(text1, text2, score)` with `score ∈ [-1, 1]` or `[0, 1]` | `CoSENTLoss` (strongly recommended) |
+| `(text1, text2, label)` with `label ∈ {0, 1}` | `OnlineContrastiveLoss` |
+| `(text, class_id)` single-column with integer class | `BatchAllTripletLoss` |
+| `(query, positive, negative, score_diff)` | `MarginMSELoss` (distillation) |
+| `(text, teacher_embedding)` | `MSELoss` (embedding distillation) |
+| Want multiple output dims from one training | Wrap any of the above in `MatryoshkaLoss` |
+| No labels at all, just sentences | `DenoisingAutoEncoderLoss` or `ContrastiveTensionLossInBatchNegatives` |
+
+## Contrastive losses (pairs + triplets, no labels)
+
+### `MultipleNegativesRankingLoss` (MNRL)
+
+The default bi-encoder loss. Uses **in-batch negatives**: every other `positive` in the batch acts as a negative for the current `anchor`.
+
+```python
+loss = MultipleNegativesRankingLoss(model, scale=20.0)  # similarity_fct defaults to cos_sim
+```
+
+- **Data**: `(anchor, positive)` or `(anchor, positive, negative)`. More columns = more explicit hard negatives per row.
+- **Scale**: temperature. Default `scale=20.0` multiplies similarities by 20 (equivalent to softmax temperature 0.05). Tune only if cosine similarities end up saturated.
+- **Critical**: set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args. Otherwise duplicate anchors create false negatives.
+- **Tip**: batch size matters a lot. More in-batch negatives = better gradients.
+
+### `CachedMultipleNegativesRankingLoss`
+
+Same loss, but with gradient caching (GradCache): forwards in mini-batches but computes the contrastive loss over the full batch. Use this when you want effective batch size of 256+ but your GPU can only fit 32 forwards.
+
+```python
+loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)
+```
+
+- **Incompatible with `gradient_checkpointing=True`**.
+- Set `mini_batch_size` to whatever `per_device_train_batch_size` would be if you couldn't use this. Then crank the actual `per_device_train_batch_size` to what you want the effective batch to be (256+, 1024+).
+
+### `MultipleNegativesSymmetricRankingLoss`
+
+MNRL computed bidirectionally. Scores positives from both (anchor -> positive) and (positive -> anchor) directions. Slightly better on retrieval tasks where the "anchor" and "positive" distinctions are soft (paraphrase, deduplication).
+
+### `CachedMultipleNegativesSymmetricRankingLoss`
+
+Cached variant of the above.
+
+### `GISTEmbedLoss`
+
+Like MNRL, but uses a **guide model** (a separate pretrained Sentence Transformer) to **filter out false negatives** before computing the contrastive loss. The guide model scores each potential negative. If it looks too similar to the positive, it's excluded.
+
+```python
+guide = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+loss = GISTEmbedLoss(model, guide=guide)
+```
+
+- Expensive: guide model does a forward pass per batch.
+- Strong when your in-batch negatives are noisy (e.g. small corpus with many near-duplicates).
+
+### `CachedGISTEmbedLoss`
+
+Cached + GIST.
+
+### `MegaBatchMarginLoss`
+
+In-batch margin-based triplet: for each anchor, find the hardest negative in the batch and apply a margin loss. Older pattern, usually outperformed by MNRL.
+
+### `TripletLoss`
+
+Classic triplet margin loss on explicit `(anchor, positive, negative)`. Uses a fixed margin and the hardest in-batch is not considered, only the provided triplet.
+
+```python
+loss = TripletLoss(model, distance_metric=TripletDistanceMetric.EUCLIDEAN, triplet_margin=5)
+```
+
+- Simpler than MNRL, less powerful when the batch has useful negatives.
+- Good for cases where you have *trusted* pre-mined triplets and want to avoid in-batch noise.
+
+## Batch-triplet losses (single column + integer label)
+
+These mine triplets *within* the batch from samples sharing a label.
+
+### `BatchAllTripletLoss`
+
+For each anchor, form triplets with all positive/negative combinations in the batch. Max signal per batch.
+
+```python
+loss = BatchAllTripletLoss(model, margin=5)
+```
+
+- **Data**: single text column + integer `label`. Needs multiple samples per label in each batch (set `batch_sampler=BatchSamplers.GROUP_BY_LABEL`).
+
+### `BatchHardTripletLoss`
+
+Same, but only the single hardest positive + hardest negative per anchor.
+
+### `BatchSemiHardTripletLoss`
+
+Semi-hard mining: negatives harder than the positive but easier than the margin. Often more stable than fully-hard.
+
+### `BatchHardSoftMarginTripletLoss`
+
+Variant with a soft margin (log-sum-exp) instead of a fixed margin hinge.
+
+**When to use batch-triplet losses**: classification-style datasets (labels are class IDs, not pair annotations). E.g. "train an embedder where samples from the same class are close."
+
+## Scored regression losses (labeled pairs, float score)
+
+### `CoSENTLoss`
+
+**The recommended regression loss** for `(text1, text2, score)`. Trains on pairwise ranking: for any two pairs `(a, b)` and `(c, d)` with `score(a,b) > score(c,d)`, the model should score `(a, b)` higher. Much better than squared error.
+
+```python
+loss = CoSENTLoss(model, scale=20.0)
+```
+
+- **Data**: `(text1, text2, float_score)`. Labels can be `[0, 1]` or `[-1, 1]`.
+- Works well with small datasets (STS-B, 5k pairs).
+
+### `AnglELoss`
+
+Similar to CoSENT but uses angle-based optimization in complex space. Sometimes outperforms CoSENT on tasks with fine-grained similarity gradations. Strong alternative.
+
+### `CosineSimilarityLoss`
+
+Squared-error loss on cosine similarity: `mse(cos(text1, text2), label)`. Simpler than CoSENT, usually worse. Keep for legacy / reproducibility.
+
+## Contrastive labeled losses (labeled pairs, binary label)
+
+### `ContrastiveLoss`
+
+For `(text1, text2, label)` where `label ∈ {0, 1}`. Minimizes distance for positives. Pushes negatives past a margin.
+
+```python
+loss = ContrastiveLoss(model, margin=0.5, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE)
+```
+
+### `OnlineContrastiveLoss`
+
+Same setup but **ignores "easy" pairs** (positives already close, negatives already far) and only optimizes hard ones. Much more robust to label noise.
+
+```python
+loss = OnlineContrastiveLoss(model, margin=0.5)
+```
+
+**Preferred over `ContrastiveLoss`** for most practical labeled pair datasets.
+
+### `SoftmaxLoss`
+
+A classifier head on concatenated `(u, v, |u-v|)` embeddings, trained with cross-entropy. Useful when you have NLI-style multi-class labels (entailment / neutral / contradiction) and want a categorical loss. Historically important (it trained the first popular sentence embedding models) but **generally outperformed by MNRL**.
+
+## Distillation losses
+
+### `MSELoss`
+
+Regress the student's embedding to match a teacher's embedding.
+
+- **Data**: `(text, teacher_embedding)`. The teacher embedding is a fixed vector per row.
+- Use when distilling a smaller bi-encoder from a larger one.
+
+### `MarginMSELoss`
+
+For `(query, positive, negative, score_diff)`: minimize `mse(student_score_diff, teacher_score_diff)`. The teacher is typically a cross-encoder that produced the score differences.
+
+- **Data**: 3 text columns + 1 float column.
+- Workhorse of dense retriever training from cross-encoder teachers (ms-marco distillation).
+
+### `DistillKLDivLoss`
+
+KL-divergence distillation: student's softmax distribution over candidates should match teacher's.
+
+- **Data**: `(query, passages[], teacher_scores[])`.
+- Good for list-wise distillation when you have multiple candidates per query.
+
+See `../scripts/train_sentence_transformer_distillation_example.py` for the end-to-end pattern (its docstring covers Embedding MSE / Margin MSE / Listwise KL with full recipes).
+
+## Regularizer / wrapper losses
+
+These don't have their own data shape. They wrap another loss and add a regularization objective.
+
+### `MatryoshkaLoss`
+
+Train **once**, deploy at any of several dimensions. Wraps any loss and computes it at multiple truncated dimensions, adding them weighted.
+
+```python
+base_loss = MultipleNegativesRankingLoss(model)
+loss = MatryoshkaLoss(
+    model,
+    base_loss,
+    matryoshka_dims=[768, 512, 256, 128, 64],
+    matryoshka_weights=[1, 1, 1, 1, 1],   # relative weighting per dim
+)
+```
+
+- At inference, `SentenceTransformer(..., truncate_dim=128)` gives 128-dim output with ~95% of full quality.
+- Default matryoshka_dims: pick the dims you want to deploy at. Smaller dims improve compression at the cost of slight quality drop.
+
+### `Matryoshka2dLoss`
+
+2D-Matryoshka: reduce dimension **and** number of transformer layers in a single wrapper. Internally composes `MatryoshkaLoss` + `AdaptiveLayerLoss`, so you only need this one (don't wrap it in AdaptiveLayerLoss yourself). Deploy at any (dim, layer) pair at inference.
+
+### `AdaptiveLayerLoss`
+
+Wrap any loss. Adds a term that trains each of the transformer's layers to be a valid exit point. Deploy with fewer layers at inference for faster encoding.
+
+```python
+loss = AdaptiveLayerLoss(
+    model,
+    base_loss,
+    n_layers_per_step=1,
+    last_layer_weight=1.0,
+    prior_layers_weight=1.0,
+)
+```
+
+### `GlobalOrthogonalRegularizationLoss` (GOR)
+
+Stand-alone regularizer (not a wrapper, despite living in this section). Penalizes embedding pairs whose dot product deviates from orthogonality, encouraging the model to spread embeddings across the full vector space. Use it alongside a primary contrastive loss by summing the two outputs in your own training step. Can help with downstream retrieval diversity.
+
+## Unsupervised losses
+
+### `DenoisingAutoEncoderLoss` (TSDAE)
+
+Sentence-level denoising autoencoder: corrupt a sentence (drop tokens), force the model to reconstruct it. Pretraining-style, useful for domain adaptation when you have unlabeled in-domain sentences.
+
+### `ContrastiveTensionLoss`
+
+Unsupervised contrastive: two copies of the model encode the same sentence. They should agree. Pure self-supervised.
+
+### `ContrastiveTensionLossInBatchNegatives`
+
+CT with in-batch negatives. Stronger than vanilla CT.
+
+## Gotchas
+
+- **`MultipleNegativesRankingLoss` without `BatchSamplers.NO_DUPLICATES`** will include duplicate anchors in the same batch, destroying training signal. Always set the sampler.
+- **Any `Cached*` loss + `gradient_checkpointing=True` = crash.** Pick one.
+- **`TripletLoss` with bad negatives** (too easy) = loss hits zero fast and model stops learning. Mine hard negatives first.
+- **Matryoshka wrapping `CachedMultipleNegativesRankingLoss`**: supported, but the cached-loss's mini-batch semantics apply to the base loss only. Think twice before combining.