Full-MLP, LoRA, and Sparse MoE Experts

What this part does

Experts are the independently owned transformations that specialize a converted MLP layer. The subpackage defines how an expert executes, which parameters it may train, what it serializes, and how it proves that the live dense anchor remains unchanged.

Full copied MLPs support the complete package-to-model lifecycle. Compact LoRA and block-sparse experts can be assembled through public interfaces, while fine-grained segmentation and product-key retrieval remain research building blocks with deliberately narrower runtime boundaries.

Where it fits

An expert sits between the immutable anchor and the routing runtime:

1. a backend moves the original MLP into a RoutedMLP wrapper;
2. an expert representation creates independently owned specialist state;
3. training updates only the state declared by that representation;
4. a router and selection policy decide whether the expert is useful;
5. inference executes or combines the selected expert paths;
6. registry support packages only representations with a validated wire format.

The experts subpackage does not score tokens, choose routes, manage capacity, install packages, or decide how several outputs are combined. Those boundaries remain separate so the same expert can be studied under different routing and training strategies.

Implemented approaches

Full copied MLP

FullMLPExpert deep-copies one compatible anchor MLP. It begins with exactly the same parameters and buffers, but every tensor has independent storage. Ordinary forced-path training can therefore update the complete specialist without allowing gradients or optimizer state to reach the anchor.

This is the conservative baseline and the only representation currently supported by package export, package materialization across converted layers, anchored top-1 synchronization, and composition reload.

PEFT/LoRA expert

PEFTLoRAExpert injects low-rank adapters into a private frozen copy of the anchor MLP. Only adapter tensors train or enter its safetensors payload. The expert still returns a complete MLP output, preserving the common Expert interface, and it can also expose its difference from an already-computed live anchor output. Its low-rank parameterization follows LoRA, with the private-copy execution difference documented explicitly.

That delta interface supports explicit always-on-anchor and parameter-merging experiments. Package export and package-driven LoRA materialization are not yet implemented, so LoRA is composable rather than integrated.

Block-sparse delta expert

BlockSparseDeltaExpert targets configurable linear projections with a fixed set of trainable blocks. Their coordinates are selected deterministically from the configuration and stored explicitly; the implementation multiplies only the compact blocks instead of constructing one dense delta matrix.

Like LoRA, it owns a private frozen base and returns a complete output while also satisfying the anchor-relative delta contract. Its ordinary PyTorch path makes no fused-kernel or speedup claim, and package binding remains future work.

Fine-grained micro-expert set

FineGrainedMicroExpertSet partitions the intermediate channels of a compatible bias-free gated MLP into deterministic contiguous segments. Each child maps back to the full hidden width, but one child is only a partial contribution. The complete set is the atomic Expert, training, and serialization boundary. The segmentation idea is motivated by DeepSeekMoE, while the anchored initialization and all-segment execution are project-specific.

Executing every segment in original channel order reconstructs the anchor exactly at initialization. Selecting an arbitrary subset, assigning mixture weights, balancing load, or attaching children to existing top-1/top-2 replacement dispatch is not implemented.

Product-key retrieval

ProductKeyMicroExpertRetriever indexes a square collection of micro-expert identifiers using two learned subkey tables. It accepts an already projected query, performs an exact bounded top-k search, and returns stable expert indices, key pairs, masks, and differentiable selected scores. The bounded search follows the product-key construction in Mixture of A Million Experts.

Retrieval deliberately stops before execution. Query projection, selected micro-expert dispatch, output aggregation, capacity, and package persistence remain outside the version-one primitive.

How it works

Preserve one execution contract

Complete experts accept the same hidden-state shape as the anchor MLP and return the same output shape, dtype, and device. They may store a full copy or only compact trainable updates, but callers do not need a new forward contract for each parameterization.

Construction is failure-atomic: target paths, tensor structure, configuration, and serialization metadata are validated before the supplied live anchor is frozen. Loading is strict across keys, shapes, dtypes, topology, and finite values.

Make trainability representation-specific

The public Expert.set_trainable() boundary delegates parameter ownership to the representation:

• a full expert toggles every copied parameter;
• a LoRA expert toggles adapter factors only;
• a block-sparse expert toggles compact block values only;
• a micro-expert set toggles its complete segmented parameter set atomically;
• a product-key retriever may toggle only its two key tables.

This explicit ownership lets trainers assert the optimizer contents rather than relying on naming conventions in notebook code.

Keep payloads independent from the anchor

Expert safetensors contain only expert-owned state and non-executable metadata. They do not duplicate the live anchor checkpoint. Compact experts reconstruct their private frozen base from a separately supplied compatible anchor during strict loading.

Mathematical intuition

For LoRA, a targeted frozen projection (W) receives a rank-(r) update (sBA). Only the two factors train; the private base weight does not:

Fine-grained experts divide one gated MLP’s intermediate channels into non-overlapping ranges (I_i). Every child is narrow, but the complete set retains all original channels:

Product-key retrieval splits a projected query and scores a Cartesian pair of subkeys:

Minimal usage

Create one complete expert from a converted layer’s anchor MLP:

from anchored_moe import FullMLPExpert

wrapper = model.anchor_model.model.layers[2].mlp
expert = FullMLPExpert.from_anchor(wrapper.anchor_mlp)

assert expert.parameter_count == sum(
    parameter.numel() for parameter in wrapper.anchor_mlp.parameters()
)

The copied values begin equal, but their storage and trainable state are independent. Attaching or packaging the expert is a separate explicit step.

Verified implementation details

anchor_output = anchor_mlp(hidden_states)
expert = FullMLPExpert.from_anchor(anchor_mlp)
initial_output = expert(hidden_states)

assert torch.equal(initial_output, anchor_output)
expert(hidden_states).square().mean().backward()
assert all(p.grad is None for p in anchor_mlp.parameters())
assert any(p.grad is not None for p in expert.parameters())

expert = PEFTLoRAExpert.from_anchor(
    anchor_mlp,
    config=LoRAExpertConfig(
        rank=4,
        alpha=8,
        target_modules=("gate_proj", "up_proj", "down_proj"),
    ),
)

assert torch.equal(expert(hidden_states), anchor_mlp(hidden_states))
assert expert.adapter_parameter_count < expert.base_parameter_count

expert = BlockSparseDeltaExpert.from_anchor(
    anchor_mlp,
    config=BlockSparseDeltaConfig(
        block_shape=(4, 4),
        blocks_per_module=2,
        seed=29,
    ),
)

assert torch.equal(expert(hidden_states), anchor_mlp(hidden_states))
assert expert.delta_parameter_count < expert.dense_delta_parameter_count

expert_set = FineGrainedMicroExpertSet.from_anchor(
    anchor_mlp,
    config=FineGrainedMicroExpertConfig(segmentation_factor=4),
)

assert expert_set.expert_count == 4
assert torch.equal(expert_set(hidden_states), anchor_mlp(hidden_states))
segmented_count = sum(
    parameter.numel()
    for child in expert_set.micro_experts
    for parameter in child.parameters()
)
assert segmented_count == expert_set.parameter_count

result = retriever(queries)
exhaustive = (
    (queries[..., :3] @ first_subkeys.T).unsqueeze(-1)
    + (queries[..., 3:] @ second_subkeys.T).unsqueeze(-2)
).flatten(start_dim=-2)

expected = torch.argsort(
    exhaustive, dim=-1, descending=True, stable=True
)[..., :3]
assert torch.equal(result.selected_indices, expected)

Current limitations

• Only full copied MLP experts have package export, multi-layer materialization, anchored top-1 synchronization, and composition-reload support.
• LoRA currently computes a complete private-base expert before subtracting the live anchor output; it does not provide a native delta-only kernel.
• LoRA and block-sparse payloads are standalone expert artifacts rather than package wire formats.
• The block-sparse implementation uses ordinary PyTorch slicing and small matrix operations. It does not claim fused sparse execution or a runtime speedup.
• Micro-expert children cannot yet be routed or combined as an arbitrary subset. The complete set is the only replacement-expert boundary.
• Product-key retrieval accepts preprojected queries and returns selection metadata only. It does not execute experts or aggregate their outputs.
• Mixed parameterizations, non-default LoRA initialization, learned or ragged sparse masks, biased micro-expert MLPs, non-square retrieval indices, accelerator kernels, and distributed validation remain future work.

References

Full copied experts and the fixed-mask block-sparse representation are project-specific framework infrastructure rather than implementations of a paper-derived MoE or pruning algorithm. Block-coordinate ordering uses SHA-256 as standardized by NIST only for deterministic ranking, not for security or statistical randomness.

The expert theory README contains the complete equations, algorithms, assumptions, tests, and documented differences. The architecture blueprint defines the shared immutable-anchor and expert-modularity invariants.

Related pages

• Read prompt-routed adapters for attention-inclusive LoRA sharing a frozen base, separate from these private-copy MLP experts.

• Read backends to see how the live anchor MLP is discovered and moved into a routed wrapper.

• Continue to routers to learn how expert utility scores are constructed and compiled into a frozen runtime.

• Continue to policies to distinguish selecting experts from combining their outputs.

• Read training for optimizer ownership.

• Read registry for package formats and lifecycle support.