Needle Shared-Prefix LLM Branches and LoRA Training

What this path does

A BranchedModel splits a Needle 3 model into a frozen prefix and independent last-layer suffixes. One suffix can classify a request, another can specialize in a domain, and the original frozen suffix remains the response fallback. Routing chooses one response branch for the whole request.

This path has its own owner, training strategies, and artifact format. It is separate from the Hugging Face MLP conversion and ExpertRegistry lifecycle. The backend preserves Needle’s attention, mHC residual lanes, QKV convolutions, and frozen engram computation. It does not replace Needle blocks with generic dense MLPs.

In JAX, the prefix runs once and its full state feeds both the router and selected response suffix. In the browser deployment, complete native models run separately while sharing the downloaded base bytes.

Install and construct branches

Install the optional reference-model and training dependencies:

python -m pip install -e ".[needle]"

The backend pins cactus-needle==3.0.4. Supply a local format-v2 master .safetensors checkpoint with configuration metadata; a quantized .cact file cannot serve as the training checkpoint. No weights are downloaded implicitly. Ordinary package import does not load JAX, Flax, Optax, or Needle.

This example assumes a checkpoint with more than 16 layers. The split is a layer count and must leave at least one layer on each side:

import numpy as np
from anchored_moe.branching import (
    BranchConfig,
    BranchedModel,
    BranchedModelConfig,
    FixedBranchSelection,
    LinearRequestScorer,
    NeedleBranchBackend,
)

backend = NeedleBranchBackend.from_checkpoint(
    "checkpoints/needle3.safetensors",
    shared_prefix_layers=16,
    model_id="Cactus-Compute/needle3",
    revision="<pinned-commit>",
    tokenizer_id="<pinned-tokenizer-identity>",
)
model = BranchedModel(
    backend,
    BranchedModelConfig(
        shared_prefix_layers=16,
        response_branches=(
            BranchConfig("base", "frozen"),
            BranchConfig("calendar", "lora", rank=8),
        ),
        router_branch=BranchConfig("router", "lora", rank=8),
    ),
    LinearRequestScorer(
        ("calendar",),
        np.zeros((backend.feature_size, 1), np.float32),
        [0.0],
    ),
    FixedBranchSelection("base"),
)

The Python selection policy above deliberately selects the base while branches are being prepared. The browser’s generative router is configured separately.

All choices keep the shared prefix, embeddings, final normalization, tied language-model head, and engram memory/projections frozen. A copied suffix trains its declared block tensors; it does not copy and train the entire model.

Choose a routing strategy

The JAX path separates scoring from selection. LinearRequestScorer reads pooled prefix features, or pooled features from the optional router suffix. Set router_branch=None for a small linear head directly on the prefix. ArgmaxSelection, ThresholdSelection, and FixedBranchSelection are interchangeable policies.

Utility scores estimate the improvement over base, whose evidence is zero; ties favor base. Probability-mode scores use a softmax over all responses. A utility value is not a probability. Use calibrate_utility_selection on held-out data when choosing activation thresholds. Updating a branch or scorer invalidates composition-bound calibration and requires recalibration.

For official-runtime deployment, the preferred router generates a single select_branch tool call through normal completion. It can read recent conversation context without a JavaScript JAX runtime. Train that suffix below, then use the browser guide to export and execute it. Training this suffix does not change the Python owner’s scorer or selection policy: route_request() and generate_tokens() still use those objects.

Train specialists independently

BranchTrainer.train_specialist(model, branch_id, examples, data_id=...) optimizes only the named response suffix. Supply ResponseTrainingExample objects with aligned input IDs and labels; mask prompt positions with -100. The objective is shifted causal cross-entropy, with optional precomputed anchor-logit replay KL and parameter drift penalties.

train_router() fits the affine scorer and, optionally, its private suffix. Use train_suffix=False for head-only fitting. Router examples must describe the request before its answer is generated. Training and calibration datasets must remain separate.

JaxAdamStrategy is the provided optimizer strategy. To replace it, implement BranchTrainingStrategy.optimize(parameters, loss_fn). A BranchObjective supplies validation and a differentiable scalar loss. Each training call owns only its supplied parameter tree and publishes a validated result atomically; a failed call leaves the model unchanged.

Train a context-aware generative router

Continue from the model above using the matching deployed tokenizer. The helper renders the same history envelope, system prompt, and tool schema used by browser inference, and masks prompt labels automatically:

from anchored_moe.branching import (
    BranchTrainer,
    GenerativeRouterConfig,
    JaxAdamStrategy,
    RoutingMessage,
    prepare_generative_router_example,
)
from needle.model.export import RefTokenizer

router = GenerativeRouterConfig(
    model.response_ids,
    fallback_branch=model.config.fallback_branch,
    descriptions=(
        ("base", "General requests"),
        ("calendar", "Calendar and scheduling requests"),
    ),
)
tokenizer = RefTokenizer.from_cact("checkpoints/needle3.cact")
example = prepare_generative_router_example(
    router,
    tokenizer,
    "Do the same tomorrow",
    "calendar",
    history=(
        RoutingMessage("user", "Schedule a meeting at noon"),
        RoutingMessage("assistant", "Meeting scheduled"),
    ),
    max_sequence_length=model.backend.max_sequence_length,
)
trainer = BranchTrainer(JaxAdamStrategy(steps=20))
trainer.train_generative_router(model, [example], data_id="route-training-v1")

This call trains only the router suffix. The prefix, response suffixes, and affine head remain unchanged. The single example demonstrates the API; train with representative labeled requests, context-dependent cases, and explicit base choices, then evaluate held-out routes after quantized export.

The renderer rejects token overflow, reserving generation capacity. History contains preceding messages only. See the context contract for truncation rules and runtime limits.

Generate with an incremental KV cache

For the Python owner, provide an unpadded integer token array shaped [1, sequence], rendered consistently with the affine router’s training data:

result = model.generate_tokens(input_ids, max_new_tokens=64, use_cache=True)
assert result.prefix_evaluations == 1

The route stays fixed for the response. Prefill runs once; subsequent tokens reuse global KV state, bounded local-attention rings, QKV convolution tails, and engram history. Each response has private immutable decode state. There is no automatic cache reuse between requests.

Caching is the default when the backend supports it. use_cache=False selects full-history recomputation for parity/debugging. benchmark_branch_generation compares both paths including routing, prefill, and decoding. This JAX greedy decoder has an optional logits processor but does not itself guarantee schema-valid tool calls.

Save or export a composition

Choose the artifact for the runtime you intend to use:

The reference package saves the learned router suffix, while GenerativeRouterConfig is supplied separately to browser export and stored in its manifest. Reference packages do not contain optimizer state.

from anchored_moe.branching import export_needle_branch, export_needle_bundle

artifact = export_needle_branch(
    model,
    "calendar",
    "deployment/calendar.cact",
    tokenizer="checkpoints/needle3.cact",
)
bundle = export_needle_bundle(
    model,
    "deployment/all-branches",
    tokenizer="checkpoints/needle3.cact",
)

Native export merges the selected suffix offline and writes CQ4 weights with 128-element groups. Exports refuse overwrite. The standard native bundle does not make the affine router executable in the official engine. Its response models each contain the full prefix and allocate separate native caches.

Quantized deployment can change both response quality and routing outcomes. Source-reference calibration is not native-quality validation. Original Needle probe heads are omitted because their behavior after these changes has not been validated.

Run the small-model checks

These examples use random local models and do not download weights. Run them after installing the needle extra:

import runpy
runpy.run_path("examples/verify_needle_branches.py", run_name="__main__")

import runpy
runpy.run_path("examples/verify_needle_deployment.py", run_name="__main__")

The checks establish small-model correctness. Real Needle weights, route quality, CQ4 quality, and target-device performance need separate evaluation. Current boundaries include single unpadded requests, isolated branch training, and no shared-prefix native ABI or learned engram updates.

Related pages

• Browser deployment covers byte patches, recent context, and official WASM execution.
• Routers explains the separate token-level router pipeline.
• Training covers MLP expert training and preservation objectives.
• Prompt-routed adapters covers the separate whole-model PEFT path.