Lfm2Bi: LFM2.5 bidirectional encoder for long-context span detection

[adaamko]

Problem

The encoder detectors run bidirectional token classification (ModernBERT / mmBERT / EuroBERT), which pays full-attention cost over the whole context. LFM2 is a hybrid architecture — mostly O(n) short-convolution layers with only a handful of full-attention layers — so a bidirectional LFM2 encoder should give long-context bidirectional encoding at a fraction of a standard encoder's cost, with a span-detection head on top.

Liquid AI has already proven the causal-to-bidirectional conversion: their LFM2.5 retrievers replace the causal attention mask with full bidirectional attention and make the short convolutions non-causal, releasing a 350M LFM2.5-Embedding (CLS pooling) and LFM2.5-ColBERT (per-token MaxSim) covering 11 languages. The per-token ColBERT variant shows the backbone already produces useful token-level representations. But the release targets short-context retrieval/embeddings — not long-context token classification. That gap is this issue.

Current behavior

Work has started on main; a contributor picks this up mid-flight rather than from zero:

• scripts/train_lfm_span_detector.py exists and trains a token-level span tagger on the bidirectional backbone. LfmForSpanTagging (scripts/train_lfm_span_detector.py:39) wraps LiquidAI/LFM2.5-ColBERT-350M (loaded with trust_remote_code, exposing Lfm2BidirectionalModel) with LayerNorm + dropout + a linear head; the LayerNorm is required because the retriever's raw hidden states are not scaled for classification (large magnitudes saturate the logits — see the comment at scripts/train_lfm_span_detector.py:50-52). Row loading, label alignment, tokenization, and metrics are imported verbatim from scripts/train_span_detector.py (load_rows/tokenize_split/compute_metrics, lines 43/97/134 there), so it stays a thin backbone swap. Defaults: --model-name LiquidAI/LFM2.5-ColBERT-350M, --max-length 8192 (train_lfm_span_detector.py:88-90).
• A related decoder-side baseline is already operational: the generative LFM detector (LFM 8B SFT) is served with scripts/experiments/serve_lfm.sh (vLLM, 32k context) and evaluated per-source / per-language with scripts/experiments/run_lfm_eval.sh via scripts/evaluate_generative_model.py. That is a separate track from the bidirectional encoder, but it shares the metric stack, so encoder results are directly comparable.
• The comparison target and shared scoring already exist: jhu-clsp/mmBERT-base trained by scripts/train_span_detector.py on KRLabsOrg/lettucedetect-code-hallucination + KRLabsOrg/lettucedetect-prose-hallucination, scored with scripts/evaluate_span_model.py (char-span and example-level P/R/F1, --by dataset / --by language) on top of scripts/span_eval_metrics.py.

What to do

What remains, in order:

1. Baseline the existing script: run scripts/train_lfm_span_detector.py on the two published datasets and score with scripts/evaluate_span_model.py --split test — establish where the un-adapted backbone lands vs mmBERT-base before touching anything.
2. Long-context adaptation: the backbone was trained for short retrieval inputs. Design a length curriculum (512 → 2k → 8k) on multilingual prose + code, and verify the non-causal convolution + positional handling actually behaves at 8k positions (the training script already accepts --max-length 8192; the open question is representation quality there, not plumbing).
3. Kill-gates, evaluated head-to-head against mmBERT-base on the same test split: (a) token-representation quality at 6-8k positions comparable to short-context level; (b) binary hallucination span char-F1 competitive with mmBERT-base. If either gate fails after honest effort, write the result up — a documented negative result closes this bet and is a valid contribution.
4. Measure the actual motivation: throughput (samples/sec) and peak memory vs mmBERT-base at 4k and 8k input lengths, same hardware, same batch size.
5. Optional: reproduce Liquid's bidirectional patches from a stock LFM2 checkpoint (bidirectional attention mask + non-causal symmetric convs) instead of starting from the ColBERT release, if the released backbone proves limiting.

Acceptance

This is a research issue: experiments, design notes, and negative results are all valid contributions. Concretely:

• results are produced by scripts/evaluate_span_model.py on the public test split, so they are reproducible and comparable to the published encoder numbers;
• the report includes the throughput/memory comparison (item 4) — speed at long context is the point of the bet;
• training changes stay in scripts/train_lfm_span_detector.py (or a sibling script with a descriptive name), runnable end-to-end with --limit for smoke tests;
• any code that moves into the lettucedetect/ package gets unit tests under tests/ following the test_*_pytest.py pattern (tests/pytest.ini), run with python -m pytest tests/ -v; python tests/run_pytest.py stays green throughout.

Non-goals

• No production release commitment for the resulting model.
• No embedding / retrieval work — Liquid's release covers that; this is token classification only.
• No changes to the generative LFM SFT track (scripts/experiments/serve_lfm.sh / run_lfm_eval.sh) — it is a comparison point, not the subject.
• No new metric code — scripts/span_eval_metrics.py is the shared scorer.

Start here

git clone https://github.com/KRLabsOrg/LettuceDetect.git
cd LettuceDetect
pip install -e ".[dev]"
python tests/run_pytest.py   # should be green before you change anything

# The code to read first:
sed -n '1,60p' scripts/train_lfm_span_detector.py   # existing backbone wrapper + rationale
grep -n "def load_rows\|def tokenize_split\|def compute_metrics" scripts/train_span_detector.py
sed -n '1,14p' scripts/evaluate_span_model.py       # how results get scored

[avirup112]

Thanks for writing this up — the LFM2.5 retriever release is good evidence the bidirectional patches actually work in practice, not just in theory, so I think this is worth pursuing seriously.

I'd like to take this on. I have open-ended access to GPU compute, so I'm planning to work through all four steps rather than just spiking the architecture piece. Splitting it into sub-tasks below mainly so progress is visible and so the riskiest part gets validated before I sink time into the rest.

Plan

1. Bidirectional LFM2 — get it loading and training at all
transformers>=4.55 already ships Lfm2Model/Lfm2Config natively (no trust_remote_code needed), so this is a real subclassing target rather than a from-scratch reimplementation.

• First check whether Liquid's released bidirectional weights (LFM2.5-Embedding-350M / LFM2.5-ColBERT-350M) are reusable as a starting point — strip the pooling/ColBERT head and see what's left underneath.
• If not, hand-patch a causal LFM2-350M/LFM2-700M checkpoint: zero out the causal attention mask in the GQA layers, and make the short-conv layers symmetric (centered kernel instead of causal/left-padded).
• Exit criterion: MLM-style forward+backward pass on a toy corpus, loss going down, no NaNs.

2. Long-context curriculum (512 → 2k → 8k)
MLM (or similar denoising) adaptation on multilingual prose + code at increasing sequence lengths, since Liquid only validated short-context retrieval. I'll checkpoint at each stage so representation quality at each length is checkable independently rather than only at the end.

3. Wire into the existing harness
scripts/train.py currently does AutoModelForTokenClassification.from_pretrained(model_name), with context/question tokens masked to -100 and answer tokens labeled 0/1. Whether this is a one-line --model-name swap or needs a custom Lfm2ForTokenClassification head depends on whether step 1's config resolves through the auto-class mapping — I'll confirm this early. TinyLettuce already set precedent for swapping in a non-default encoder (Ettin), so the harness is probably flexible enough either way.

4. Kill-gates vs mmBERT-base

• (a) MLM/representation quality at 6–8k ≈ short-context level
• (b) span-F1 head-to-head on the existing test split via scripts/evaluate.py's char-level scorer

I'll report back at the end of step 1 with what I find (whether Liquid's checkpoints are reusable as-is, and whether AutoModelForTokenClassification resolves cleanly) before going further into the long-context curriculum, in case anything there changes the plan. Will keep this thread updated as I go — happy to take feedback on the approach before I'm too deep into it.

[adaamko]

Thanks for the thorough write-up — this is a strong, well-scoped plan, and the LFM2.5 retriever release is exactly the evidence that the bidirectional conversion holds up in practice. Great direction to explore; a few pointers that might save you time:

Reusing Liquid's weights (step 1). Worth trying as a warm-start, with two caveats: (i) the released LFM2.5-Embedding/ColBERT-350M are tuned for short-context retrieval, so even if the backbone strips cleanly you'll still be doing most of the long-context MLM work in step 2; (ii) check the LFM2.5 license terms before redistributing any derivative weights. If the backbone doesn't reuse cleanly, the hand-patch path (zero the causal mask in the GQA layers + symmetric short-conv kernels) is the right fallback — essentially what VAGOsolutions' code_modification/lfm2_bi.py does.

Harness (step 3). Confirm early whether transformers exposes an Lfm2ForTokenClassification — if it only ships Lfm2Model/Lfm2ForCausalLM, you'll need a thin token-classification head subclass (linear layer over last_hidden_state, reusing the -100 masking scripts/train.py already does). TinyLettuce/Ettin is the precedent for a non-default encoder.

Kill-gate (b) — concrete target. Use scripts/evaluate.py's char-level scorer on the unified test split so it's comparable to lettucedect-v2-mmbert-base: ALL span-F1 ≈ 0.64 / example-F1 ≈ 0.87, code-agent span-F1 ≈ 0.51, PsiloQA ≈ 0.71. Data is the public KRLabsOrg/lettucedetect-code-hallucination + prose sources.

Sharing your step-1 findings before the curriculum is the right checkpoint. A draft PR whenever you have something running is welcome.