Contextual word-drift pipeline (BERT, single DGX Spark, on-the-fly aggregation)

[scheuclu]

Summary

Upgrade the drift analysis from Word2Vec/TWEC (300-d, distributional) to contextual drift: run the historical FineWeb corpus (2014–2025) through a frozen transformer encoder and measure how each word's contextual sense shifts over time (e.g. zoom lens→video-call, covid/mask/remote post-2020, delve becoming LLM-ese). This captures polysemy/sense change that Word2Vec cannot.

This issue supersedes the Antigravity-drafted ticket DATA-402 ("DGX-Spark cluster + PySpark"), which assumed infrastructure that does not exist. We run on a single DGX Spark, so the design is a single-process PyTorch pipeline with on-the-fly aggregation.

Verified environment (we run ON the target box)

• aarch64, NVIDIA GB10, 128 GB unified CPU+GPU memory, ~273 GB/s mem bandwidth, ~1 PFLOP FP4.
• torch 2.12.0+cu130 (CUDA OK), transformers 5.10.2.
• 519 GB free of a 916 GB volume.
• Shared vocab data/vocab/vocab.json = 119,466 words (<UNK>=0).
• No Spark / HDFS / cluster. Everything in-process.

Why we discard the original ticket's design

The ticket frets about 36 TB of token vectors and 2 TB of shuffle space and specs a 3-node / 24×A100 cluster with PySpark mapPartitions/reduceByKey. All of that is an artifact of a distributed design we are not using. With on-the-fly aggregation (never store per-token vectors) on one device:

• Total disk footprint is ~20 GB (a non-issue at 519 GB free).
• The real budget is GPU-hours, set by tokens/year.

Decisions (confirmed with user)

• Encoder: bert-base-uncased (768-d). Fast WordPiece tokenizer with clean word_ids() alignment, already lowercase (matches the existing vocab), SDPA attention. Frozen — AutoModel, take last_hidden_state.
• Corpus: full 1 B tokens/year × 12 = 12 B tokens (max coverage; multi-day run → throughput probe up front + solid checkpointing).
• Scope this round: data pipeline + local eval only. Web packing/frontend deferred to a follow-up.

Correctness — where the ticket's snippet was wrong

Ticket approach	What we do instead
decode each token id, strip "##"	word_ids() from the fast tokenizer maps subwords→words (works for WordPiece/BPE/SP; ## does not). First-subword pooling for v1 ("mean" deferred knob).
python dict[str, np.ndarray] per token	Vectorized GPU index_add_ into preallocated sum[V,768], sumsq[V,768], count[V], keyed by the shared vocab id. Accumulate in fp32 even though the forward runs bf16.
raw contextual cosine	Per-year mean-centering before cosine — BERT vectors are strongly anisotropic (random words sit at cos ≈ 0.3–0.6, which would swamp the drift signal). Optional all-but-top-k PCA removal as a knob.
Procrustes / TWEC compass per year	None. A frozen shared encoder puts every year in one coordinate system already → drift(w,y) = 1 − cos(c̄_{w,y}, c̄_{w,2018}) on centered centroids.
store/segment true sentences	Non-overlapping ~128-word windows (v1, no extra dependency).

Reuse the existing shared vocab to index the accumulator → result is row-comparable to the Word2Vec word set (same words, different basis). Nothing under models/embeddings*, models/aligned/*, or existing drift artifacts is touched.

Architecture / data flow

re-stream FineWeb RAW text (per snapshot)        ── reuses snapshot_registry + the
   filter language_score >= 0.65                    streaming/checkpoint contract of
        |                                           pipeline/data_pipeline.py
        v  lowercase + whitespace split, slice into 128-word windows
GPU forward (bf16, SDPA, inference_mode) -> last_hidden_state [B,L,768]
        |  word_ids() -> first-subword positions -> target vocab-id tensor [B,L] (-1 = ignore)
        v
ContextualAccumulator  sum[V,768] += h ; sumsq[V,768] += h*h ; count[V] += 1   (index_add_)
        |  checkpoint per completed snapshot (resumable)
        v  per year
finalize: centroid = sum/count ; dispersion = sumsq/count - centroid^2
          per-year mean-center (anisotropy fix) -> models/contextual/{year}.npy (+ _centered, _count, _dispersion)
        v
drift vs 2018 (reuse analysis/drift.py) -> models/contextual/drift_vs_2018.parquet + summary
        v
local eval: noise-floor / SNR vs the Word2Vec drift; sanity words

Key reuse: mirror the streaming + checkpoint contract of pipeline/data_pipeline.py (but read raw row["text"] — the existing pipeline lowercases/strips and discards raw text, so we must re-stream); pipeline/snapshot_registry.py::get_snapshots; analysis/drift.py (compute_drift_from_base, compute_drift_summary); the --smoke/--device entry-point style of scripts/twec_full.py.

New files

• pipeline/contextual_model.py — load frozen model + fast tokenizer (attn_implementation="sdpa", .eval().requires_grad_(False)); encode_windows(...) with is_split_into_words=True so word_ids() is clean; bf16 inference_mode forward. Width from model.config.hidden_size (never hardcode 768).
• pipeline/contextual_aggregate.py — ContextualAccumulator: GPU sum/sumsq/count; add_batch(...) builds the target-id tensor (ignore special/continuation/OOV/stopword/len<MIN_WORD_LENGTH/digit) + three index_add_s; save_state/load_state (npz), to_numpy.
• pipeline/contextual_stream.py — iter_windows(year, tokenizer): re-stream FineWeb per snapshot (same load_dataset(..., streaming=True) + language_score>=0.65 filter), raw text, lowercase + 128-word windows, stop per snapshot at target_tokens // len(snapshots). run_year(...) batches → forward → accumulator.add_batch, checkpointing each completed snapshot.
• pipeline/contextual_finalize.py — centroids, dispersion (on raw vectors — centering doesn't change within-word variance), per-year mean-centering (+optional PCA), write models/contextual/{year}*.npy.
• analysis/contextual_drift.py — load centered centroids into dict[int,np.ndarray], call analysis/drift.py::compute_drift_from_base(base_year=ANCHOR_YEAR) + compute_drift_summary, gate on a per-(word,year) min-count valid_mask; emit parquet + summary with a dispersion column.
• scripts/contextual_benchmark.py — run FIRST. Pull ~2k windows from one 2018 snapshot (no disk write), warm up, time batches, sweep batch ∈ {128,256,512}; print windows/sec, tokens/sec, GPU mem, and an hours-per-year + total-wall-clock table to confirm the 12 B-token ETA and pick batch size.
• scripts/contextual_drift.py — entry point mirroring scripts/twec_full.py (sys.path prelude, from config import ...): --smoke, --device, --model, --years, --stage {stream,finalize,drift,all}. Resumable (skip finalized {year}.npy, resume partials). --stage stream runs overnight in background; finalize/drift after.

config.py additions (append only — import, never hardcode)

# --- Contextual drift pipeline ---
CONTEXTUAL_MODEL = "bert-base-uncased"
CONTEXTUAL_LAYER = -1                       # last_hidden_state
CONTEXTUAL_POOLING = "first"                # "first" | "mean" (mean deferred)
CONTEXTUAL_MAX_LEN = 128                    # non-overlapping window length
CONTEXTUAL_BATCH_SIZE = 256                 # tune from benchmark
CONTEXTUAL_TARGET_TOKENS_PER_YEAR = TARGET_TOKENS_PER_YEAR   # 1B (full)
CONTEXTUAL_MIN_COUNT = 50                   # min obs for a trusted (word,year) centroid
CONTEXTUAL_CENTERING = True                 # per-year mean-centering (anisotropy fix)
CONTEXTUAL_PCA_REMOVE_K = 0                 # all-but-top-k; 0 = off
CONTEXTUAL_USE_STOPWORDS = True
CONTEXTUAL_STOPWORDS = frozenset({...})     # small English list
CONTEXTUAL_DIR = MODELS_DIR / "contextual"
CONTEXTUAL_STATE_DIR = DATA_DIR / "contextual_state"

Checkpoint format (resumable multi-day run)

Per year in data/contextual_state/:

• {year}.partial.npz — sum(V×768 f32), sumsq(V×768 f32), count(V f32) ≈ 0.73 GB/yr. Atomic write (.tmp + os.replace).
• {year}_checkpoint.json — {completed_snapshots, total_tokens, total_windows, model, hidden_dim, pooling, max_len, config_hash, updated_at}. Resume loads the npz and skips completed snapshots only if model/config_hash match (mismatch → refuse). On year completion, finalize then delete partial + checkpoint (mirrors data_pipeline.py).

Budget

• Disk: corpus streamed (writes nothing); 12 partials ≈ 8.8 GB + final artifacts ≈ 8.8 GB → < 20 GB of 519 GB free. Not a constraint.
• Compute: the binding cost. Benchmark sets the real number; full 12 B tokens through bert-base is expected to be a multi-day, checkpointed run. 128 GB unified memory is ample (model <1 GB, accumulators ~0.73 GB/yr, activations a few GB).

Verification

1. Smoke (--smoke, 1 snapshot for {2014, 2018}, ~2 M tokens → models/contextual_smoke/): counts>0 for common words; {year}.npy shape (119466, 768); checkpoint deletes on completion; kill mid-run + relaunch → completed snapshot skipped, counts continue.
2. Throughput probe → confirm ETA, set batch size (reconsider 1 B target if untenable).
3. Anisotropy: random word-pair mean cosine high on raw, ≈ 0 after mean-centering (else enable PCA_REMOVE_K=1..3).
4. Stable words: cos(centered("house",2014), centered("house",2025)) high / low drift; same for table, water.
5. Drifters: covid/zoom/mask/remote spike 2020+, delve 2023+; rank by total drift, confirm near top while function words sit low.
6. Dispersion: polysemous (bank, apple, python) > monosemous; invariant to centering.
7. Quality vs Word2Vec: noise-floor / SNR vs shipped TWEC drift; confirm contextual signal adds info on polysemy-driven cases.

Build order

1. config additions
2. pipeline/contextual_model.py
3. scripts/contextual_benchmark.py (run, set batch/ETA)
4. pipeline/contextual_aggregate.py
5. pipeline/contextual_stream.py
6. scripts/contextual_drift.py --smoke (verify resume)
7. pipeline/contextual_finalize.py
8. analysis/contextual_drift.py + artifacts
9. full run (background, checkpointed)
10. eval (Verification section)

Out of scope (follow-up phase)

Web packing (scripts/pack_web_data.py 14400-byte stride / vecs.bin), the web/lib/* 300-d hardcodes, and a /contextual-drift view. Contextual artifacts are 768-d; the follow-up parametrizes the dim and publishes a new data/vN alongside the existing TWEC data rather than replacing it.

---

Origin: Antigravity drafts linear_ticket_dgx_spark_drift_pipeline.md / dgx_spark_drift_pipeline.md. Planned with Claude Code; full plan also at ~/.claude/plans/cozy-greeting-music.md.