bwdam-text8-word2gauss-50d

Text8 Word2Gauss controls for BW-DAM

This repository contains the Gaussian word-embedding checkpoints used to reproduce Claim 5 of Dense associative memory for Gaussian distributions, locked to arXiv:2509.23162v1. Each checkpoint represents the 10,000 most frequent Text8 words as 50-dimensional means with spherical covariance.

These are experiment artifacts, not a hosted inference model. They are intended for the Bures--Wasserstein dense associative memory evaluation in the accompanying reproduction bundle.

Main finding

The real-data claim is implementation-sensitive.

All evaluations use N=10000, d=50, 3,750 perturbed queries, ten retrieval updates, two query trials per source-control model, and the exact 20 numpy.logspace(-1,2,20) beta markers inferred from Figure 4.

Checkpoints

The source-control full checkpoints are also retained:

NPZ files contain mus, variances, words, frequencies, and metadata_json; the newer files also retain context means/variances. Read metadata_json before assuming identical schemas or optimizer semantics.

Training semantics

The five fast checkpoints implement the written Word2Gauss max-margin/KL objective, vocabulary construction, negative-sampling scheme, initialization scale, and parameter constraints in deterministic batched PyTorch. They are not optimizer-equivalent to the public Cython trainer: 275,952,136 hinge examples per epoch are represented by 1,038 averaged AdaGrad steps, so gradients may cancel before being squared.

The source controls use seomoz/word2gauss commit cdf5e7f5d0c8c7582fdf92e4599bd8efda5b98bc, one worker, immediate per-hinge AdaGrad/projection, the source's fixed seed, five epochs, window 5, and two negative samples. Text8 is split into 171 bounded documents to avoid an approximately 8 GB pair matrix; this omits 5,100 positive cross-boundary pairs per epoch, an estimated fraction 2.9990814e-5.

The literal source has a material spherical-KL covariance-gradient defect in embeddings.pyx: it reads sigma_ptr[j] instead of the context covariance sigmaj_ptr[0]. The controls therefore include both:

• literal source SHA-256: a9c00384ebdafdc7f9215f4b3fbfb67d8e432bb2b083972ada34785cbe0078ec;
• corrected source SHA-256: 56e247f63a4dbd0168d5dcc204fe6b8fe01da93bce3629cb8c8513f9e7e183f2;
• pinned words.py SHA-256: c04eb45b596a424c5f3e6b7113926283ba8b82b0dfb6085dc2641b79d6230826.

At beta 50, the five-model batched result reaches mean W2 about 3.1e-9 after one step. The literal and corrected source controls instead end at 0.06816 and 0.04244, respectively. This is why the repository exposes all variants rather than presenting one checkpoint as the authors' unreleased implementation.

Download and evaluate

hf download GwendalTsang/bwdam-text8-word2gauss-50d \
  --local-dir bwdam-checkpoints
hf download GwendalTsang/bwdam-icml-2026-reproduction-bundle \
  --repo-type dataset --local-dir bwdam-reproduction
cd bwdam-reproduction
python3 -m venv .venv
. .venv/bin/activate
python3 -m pip install -r requirements.txt
python3 experiments.py --device cuda --dtype float32 --backend model-card-rerun \
  claim5 --query-trials 2 \
  --embeddings \
    ../bwdam-checkpoints/text8_word2gauss_v1.npz \
    ../bwdam-checkpoints/text8_word2gauss_seed17.npz \
    ../bwdam-checkpoints/text8_word2gauss_seed29.npz \
    ../bwdam-checkpoints/text8_word2gauss_seed41.npz \
    ../bwdam-checkpoints/text8_word2gauss_seed53.npz

Evaluate the source controls separately so they are not averaged with the fast models:

python3 experiments.py --device cuda --dtype float32 --backend source-literal-rerun \
  claim5 --query-trials 2 \
  --embeddings ../bwdam-checkpoints/text8_word2gauss_source_literal.npz
python3 experiments.py --device cuda --dtype float32 --backend source-corrected-rerun \
  claim5 --query-trials 2 \
  --embeddings ../bwdam-checkpoints/text8_word2gauss_source_corrected.npz

To retrain a deterministic batched checkpoint from the canonical 31,344,016-byte Text8 archive:

python3 train_word2gauss.py --device cuda --seed 17 \
  --backend model-card-retrain --output data/text8_word2gauss_seed17.npz

The trainer downloads Text8 only when absent and prints its hash before training. The required archive SHA-256 is a6640522afe85d1963ad56c05b0ede0a0c000dddc9671758a6cc09b7a38e5232.

Independent cloud check

The completed main-account T4 Job evaluated the seed-5 staged checkpoint from a clean pytorch:2.5.1 image. It reproduced crossings 9.561 / 13.147 / 15.630 and Figure-4 RMSE 0.02480. The Job ran for 389 billed seconds at $0.0067/min, approximately $0.0434; its full manifest reports 382.77 seconds inside the program and estimated compute $0.04274.

A superseded 17-second preflight cost approximately $0.0019 and exited 2 after tests and Claim 3 because Python 3.11 parsed the exact subcommand option --d as an abbreviation of top-level --device/--dtype. The corrected parser sets allow_abbrev=False; the replacement command retained the explicit --d 50 and completed. The failure is retained as provenance, not counted as evidence.

Licensing and limitations

No single license is asserted for these checkpoint artifacts. The public Word2Gauss code is MIT-licensed, but that does not by itself determine the license of trained parameters derived from Text8. The Text8 archive is not redistributed in this model repository, and the upstream corpus page does not provide a model-weight license. Users must review the Text8 source terms and their intended use. This deliberate limitation is why the YAML metadata omits a license field.

The checkpoints encode corpus statistics and may inherit societal, factual, and representational biases from cleaned English Wikipedia. They have not been evaluated for downstream semantic quality, fairness, privacy, or safety. They should not be used for decisions about people.

Linked resources

• Canonical reproduction logbook
• Reproduction bundle dataset
• Persistent Job-output Bucket
• Trackio artifact Bucket
• Private agent-trace dataset
• Resource collection
• Exact public Word2Gauss commit

See PROVENANCE.md for the version lock, hashes, hardware, figure digitization, and Job manifest.