🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3-4B",
  "baseline": "Same fork, data, and 100 sparse LoRA steps with indexer KL during training, but no 50-step indexer warmup. The indexer starts from random init. Dense-attention LoRA is logged as an upper-bound reference.",
  "batch_size": 64,
  "budget": {
    "gpu_hours": 2.0,
    "max_wall_minutes": 120
  },
  "command": "python - <<'EOF'\n# custom: indexer warmup + sparse attention come from the forked Spindle; loop is stock tinker SDK\nimport os, tinker\nsc = tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\ntc = sc.create_lora_training_client(base_model='Qwen/Qwen3-4B', rank=32)\nbatches = load_pg19_chunks(seq_len=4096, batch_size=64)  # user helper, fixed seed\nWARMUP = int(os.environ.get('WARMUP', 50))  # 0 for baseline arm\nfor step, b in enumerate(batches[:WARMUP + 100]):\n    loss = 'indexer_kl' if step < WARMUP else 'cross_entropy'  # 'indexer_kl' added by fork\n    tc.forward_backward(b, loss).result()\n    tc.optim_step(tinker.AdamParams(learning_rate=1e-4)).result()\ntc.save_weights_for_sampler(name=f'sparse-idx-warmup{WARMUP}').result()\nEOF",
  "dataset": {
    "hf_path": "deepmind/pg19",
    "n_examples": 9600,
    "name": "PG19",
    "split": "train"
  },
  "eval": {
    "benchmark": "PG19 validation, 4096-token chunks, evaluated with sparse top-512 attention",
    "metric": "mean next-token cross-entropy (nats), plus top-k recall of dense attention mass",
    "n_examples": 200
  },
  "hypothesis": "Warming up a top-k indexer by KL distillation from frozen full attention, then LoRA-adapting under sparse attention, gives held-out PG19 loss at least 0.05 nats lower than the same sparse LoRA run with a randomly initialized indexer.",
  "lora_rank": 32,
  "max_seq_len": 4096,
  "method": "Needs a Spindle fork that adds a DSA-style lightning indexer to Qwen/Qwen3-4B. Qwen3-4B is used because it is pure full attention, while Qwen3.5 uses hybrid linear attention. The indexer is a few low-dim heads per layer that score past tokens. The fork exposes three things: (a) an indexer-KL loss against the aggregated full-attention distribution, with the main model frozen; (b) top-k (k=512) token-selection sparse attention in forward and backward; (c) the indexer as a trainable parameter alongside the LoRA weights. Treatment stage 1 runs 50 steps that train only the indexer with KL to dense attention, using dense forward and a frozen base. Treatment stage 2 runs 100 steps of LoRA rank 32 next-token SFT on 4096-token PG19 chunks under top-k sparse attention. During stage 2 the indexer keeps its KL loss on selected tokens, detached from the LM loss, as in DSA. Both arms use the same data order, seed and LR. A dense-attention LoRA run with the same 100 steps is logged as a reference ceiling, not as an arm. Neither linked source could be read, so the design follows the Slack description and the DeepSeek/GLM indexer recipe.",
  "optimizer": {
    "lr": 0.0001,
    "name": "adam"
  },
  "steps": 150,
  "success_criterion": "Validated if treatment eval loss is at least 0.05 nats below baseline and within 0.05 nats of the dense LoRA reference. Seed noise is to be checked to be under 0.02 nats via a second baseline seed. Failed if the gap to baseline is under 0.02 nats.",
  "title": "Indexer warmup from full attention versus cold-start indexer for sparse LoRA adaptation",
  "training_type": "sft"
}

treatment_warmup_eval_loss_nats

2.636

treatment_warmup_topk512_recall

0.8661

baseline_seed0_eval_loss_nats

2.635

baseline_seed0_topk512_recall

0.8578

gap_baseline_minus_treatment_nats

-0.00143

treatment_final10_train_loss

2.612

baseline_final10_train_loss

2.608

treatment_train_minutes

17.59

baseline_train_minutes

17.97

eval_examples

194

train_batch_size

2