You can deploy this model and start earning money today!

hyunseoki/
qwen3-0.6b-lambda-gates-nke

Name: qwen3-0.6b-lambda-gates-nke
Availability: InStock

Views

No views yet

Qwen3-0.6B Lambda Gates — NKE (Non-Entity KL) Variants

Extends the baseline by adding a non-entity KL retention loss to the forget batch. In the standard recipe, only entity tokens in the forget passages carry a forget loss; non-entity tokens are unsupervised and can drift. NKE adds

L_nke = unmasked_retain_weight · KL( p_base || p_gated ) # on non-entity tokens of the forget batch

keeping the gated model's behavior on surrounding (non-entity) text close to the base model.

Variants

This repo contains 3 training variants differing in unmasked_retain_weight and λ_f:

Folder	`λ_f`	`λ_r`	`unmasked_retain_weight`	Intended effect
`nke_w1p0/`	0.1	0.5	1.0	Heavy non-entity retention
`nke_optA_w0p05/`	0.1	0.5	0.05	Light touch — keep baseline forget strength
`nke_optB_lf1p0_w0p1/`	1.0	0.5	0.1	Stronger forget + moderate retention

All other hyperparameters are identical to baseline (β=4.0, distill T=2.0, forget_retain_ratio=1:2, lr=1e-2, cosine, 3 epochs, bf16).

Contents per variant

<variant>/
  lambda_logits.pt       # 86,016 per-neuron logits (28 layers × 3072)
  neuron_indices.json    # Knowledge neurons at threshold 0.5

Usage

python

import torch, json

# Load a specific variant
gate_state = torch.load("nke_optA_w0p05/lambda_logits.pt", map_location="cpu")
with open("nke_optA_w0p05/neuron_indices.json") as f:
    knowledge_neurons = json.load(f)

See the baseline README for complete usage instructions.

Why NKE?

The baseline recipe supervises only entity tokens with the forget loss:

✅ Clear signal on what to forget
❌ The model's response to surrounding non-entity text can silently shift, degrading fluency and in-context oracle accuracy.

NKE anchors the gated model to the base model on non-entity positions, reducing this collateral drift at the cost of slightly less aggressive forgetting.

hyunseoki/qwen3-0.6b-lambda-gates-nke