Masked Language Modeling for Hashed Tokens

CS 712 — Group 8

This project tackles a structured MLM task where the vocabulary consists of opaque hex-encoded hash tokens (e.g., 0c742e19). Because token identity carries no semantic string meaning, we exploit the bit-vector structure of the 32-bit hash values to initialize embeddings and train several model variants, culminating in a lex-invariant RoBERTa with a structure-aware lookup cascade.

See the paper for full methodology, ablations, and results: CS712_Masked_Language_Modeling_for_Hashed_Tokens.pdf

---

Environment

conda create -y -n g8 python=3.10
conda activate g8
pip install -r requirements.txt

---

Project Structure

../
├── train.py                   # Base BERT/RoBERTa/DeBERTa MLM trainer
├── train_lexinvariant_pro.py  # Lex-invariant RoBERTa (main model)
├── train_bit.py               # Bit-prediction binary classifier
├── train_bit_lex.py           # Bit model + lex-invariant consistency
├── train_router_bi.py         # Router model (bi-encoder variant)
├── utils/
│   ├── focal_loss.py          # Focal loss + CustomLossTrainer
│   ├── lookup_baseline.py     # Structure-aware cascade lookup
│   └── vocab_truncator.py
├── scripts/                   # Launch wrappers (nohup + pid tracking)
│   ├── run_train.sh
│   ├── run_train_lexinvariant.sh
│   ├── run_train_bit.sh
│   ├── run_train_bit_lex.sh
│   └── run_train_router.sh
├── z_configs/                 # YAML configs for each model variant
│   ├── config.yaml            # Base model (BERT/RoBERTa/DeBERTa)
│   ├── config_lexinvariant.yaml
│   ├── config_bit.yaml
│   ├── config_bit_lex.yaml
│   └── config_router.yaml
└── data/phase_1_data/         # Train/test/vocab splits

---

Key Ideas

Bit-vector embedding initialization

Every hex token is a 32-bit integer. Instead of random initialization, each token's embedding is seeded from its bit pattern — ±1 per bit, scaled to the hidden dimension. This gives the model a structured inductive bias from step one (train.py:449–458, token_to_bit_tensor).

Lex-invariant training (train_lexinvariant_pro.py)

Two views of each sequence are generated: a clean view (low UNK dropout) and an augmented view (higher noise + optional shuffle). A consistency loss (lambda_inv) penalizes divergence between the two views' masked-token distributions, encouraging representations that are robust to surface-level token identity changes.

total_loss = clean_CE + λ_aug·aug_CE + λ_inv·consistency + λ_rep·rep_loss

Structure-aware lookup cascade (utils/lookup_baseline.py)

Deterministic rules exploit known structural patterns in the hashed token sequence (e.g., certain tokens always follow others). The lookup overrides model predictions where it has high confidence (conf_threshold=1.0) and sufficient training support (min_support=3). Support thresholds matter much more than confidence alone:

Support	CV Accuracy
1	23.0%
2	58.8%
3–5	82.6%
6–10	94.5%
≥ 11	96.9%+

---

Inference (Final Submission)

Set infer_only: true and point infer_only_model_dir to your trained model in z_configs/config_lexinvariant.yaml:

infer_only: true
infer_only_model_dir: "runs/lexinvariant/<your_run>/model"

Then run:

bash scripts/run_train_lexinvariant.sh

This skips training, loads the model, and writes submission_epoch02.txt to the run directory.

Apply the lookup cascade on top

python utils/lookup_baseline.py \
  --train data/phase_1_data/train.txt \
  --test  data/phase_1_data/test.txt \
  --vocab data/phase_1_data/vocab_eligible.txt \
  --roberta_pred runs/lexinvariant/<your_run>/submission_epoch02.txt \
  --name final

Note: Use test.txt (not validation.txt) for --test to avoid misaligned context lookups (see Section 5.5 of the paper).

---

Training from Scratch

Step 1 — Base model (optional warm-start)

bash scripts/run_train.sh

Controlled by z_configs/config.yaml. Set model_type to bert, roberta, or deberta.

Step 2 — Lex-invariant fine-tuning

Set init_checkpoint in z_configs/config_lexinvariant.yaml to the model directory from Step 1, set infer_only: false, then:

bash scripts/run_train_lexinvariant.sh

For data expansion, use:

• data/phase_1_data/combined_train_val.txt for validation-stage training
• data/phase_1_data/combined_train_test.txt for final test-stage training

Bit-vector models (experimental)

bash scripts/run_train_bit.sh      # pure bit classifier
bash scripts/run_train_bit_lex.sh  # bit model + lex-invariance

---

Config Reference

Key fields shared across all configs:

Field	Description
model_type	bert / roberta / deberta (base trainer only)
hidden_size	Model hidden dimension (must be divisible by num_attention_heads)
unk_noise_prob	Probability of replacing a token with [UNK] during training
train_on_eligible_only	If true, only mask tokens in vocab_eligible.txt
bully_token	High-frequency token to optionally down-sample
bully_keep_prob	Fraction of bully-token examples to retain
val_split	Fraction held out for validation; 0.0 trains on full data
resume_from_checkpoint	Path to Trainer checkpoint dir, model dir, or .pt file
infer_only	Skip training; load model and generate submission directly

Lex-invariant specific:

Field	Description
aug_noise_prob	UNK dropout rate for the augmented view
aug_shuffle_prob	Context shuffle probability for the augmented view
lambda_inv	Weight on consistency loss between clean and augmented views
lambda_aug_ce	Weight on augmented-view cross-entropy
consistency_temperature	Softmax temperature for KL/JS consistency

---

Monitoring Runs

Each script launches in the background via nohup. Logs and PID files land in logs/ and pids/:

tail -f logs/train_lexinvariant_<timestamp>.log
ps -p $(cat pids/train_lexinvariant_<timestamp>.pid)
kill $(cat pids/train_lexinvariant_<timestamp>.pid)

Run outputs are saved under runs/<model_type>/run_<timestamp>/:

• model/ — saved model weights + vocab.json
• results/ — Trainer checkpoints
• logs/summary.log — training summary
• submission_*.txt / submission_*.zip — prediction files