🔬Modal Research Exchange

what does it mean for Kimi K3.1 to have full attn weights already trained, and sparse attn still in training? is there some training technique where you init from full attn and then distill to sparse?

[nan] glm 5 did the same. basically they will freeze full attn and train indexer, and then start training the sparse attn
[nan] there are many fancy stuff in rubin, you can have qk gemm to be the indexer so during training you dont need to think about top k indexer, and during inference you can just enable the adaptive compression
[nan] related on sparse attn https://arxiv.org/html/2502.11089v1
[rasool] yes, this is commonly training technique used where one family of model is more available than others. For example, if you have model which has causal attention (like autoregressive), then slowly move and adapt toward to non-causal attention. for causal --> non-causal https://arxiv.org/abs/2410.17891. Also this is much easier to make it work than training from scratch.

{
  "adjacent_experiment": "",
  "reproducibility_reason": "Needs a Spindle fork: add a trainable lightning-indexer with top-k sparse attention to a served Qwen3 model. The from-scratch sparse comparison needs pretraining, so the test is indexer warmup versus no warmup.",
  "reproducible_on_spindle": true,
  "sources": [
    {
      "brief": null,
      "error": "digest failed: ask_json failed after 2 attempts: no JSON object in response: '{\\n  \"title\": \"Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention\",\\n  \"claim\": \"NSA is a trainable sparse attention architecture (compressed tokens + selected blocks + sl'",
      "full_text": true,
      "kind": "paper",
      "title": "Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention",
      "url": "https://arxiv.org/html/2502.11089v1"
    },
    {
      "brief": null,
      "error": "ReadTimeout: The read operation timed out",
      "full_text": false,
      "kind": "link",
      "title": "https://arxiv.org/abs/2410.17891",
      "url": "https://arxiv.org/abs/2410.17891"
    }
  ],
  "spec": {
    "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"
  },
  "summary": "A sparse-attention indexer is trained first to copy a frozen model's full attention, then the model is LoRA-tuned under sparse attention; we compare long-context eval loss against a cold-started indexer.",
  "tags": [
    "attention",
    "architecture",
    "lora",
    "sft",
    "paper"
  ],
  "title": "Does distilling an indexer from full attention help sparse adaptation?",
  "value_estimate": 3,
  "value_rationale": "Dense-to-sparse conversion via indexer distillation is how DeepSeek V3.2 and GLM-5 got long-context efficiency, so a cheap check matters to anyone forking attention. It needs a fork and is mostly research-facing. Triage only saw the links, since both sources failed to fetch."
}