| """ST-TENP: timestep-conditioned TENP + structural sensitivity + redundancy. |
| |
| This is a separate pruning method from protocol TENP. Physical slicing is the |
| same; selection is not. Adaptive expert budget is not included. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
|
|
| from .apply import prune_moe_layer |
| from .scoring import DEFAULT_MODALITY_FUSE |
|
|
| logger = logging.getLogger("sttenp") |
|
|
| METHOD = "st-tenp" |
| VARIANTS = ("ta", "st", "st-redundancy") |
| DEFAULT_VARIANT = "st" |
|
|
|
|
| def variant_flags(variant: str) -> tuple[str, bool]: |
| """Map a named ST-TENP variant to (temporal, redundancy).""" |
| if variant not in VARIANTS: |
| raise ValueError(f"unknown ST-TENP variant {variant!r}; expected one of {VARIANTS}") |
| if variant == "ta": |
| return "ta", False |
| if variant == "st": |
| return "st", False |
| return "st", True |
|
|
|
|
| def prune_sttenp_layer( |
| mlp, |
| collector, |
| p_mid: float, |
| variant: str = DEFAULT_VARIANT, |
| modality_fuse: str = DEFAULT_MODALITY_FUSE, |
| n_min: int = 32, |
| temporal_beta: float = 0.3, |
| temporal_topk: int = 2, |
| ) -> dict: |
| """Prune one SparseMoEMLP with ST-TENP. Does not change the TENP default path.""" |
| temporal, redundancy = variant_flags(variant) |
| meta = prune_moe_layer( |
| mlp, |
| collector, |
| p_mid=p_mid, |
| modality_fuse=modality_fuse, |
| score="projection", |
| apply_audio_discount=False, |
| compensate="none", |
| temporal=temporal, |
| redundancy=redundancy, |
| n_min=n_min, |
| temporal_beta=temporal_beta, |
| temporal_topk=temporal_topk, |
| ) |
| for spec in meta.values(): |
| spec["method"] = METHOD |
| spec["variant"] = variant |
| spec["score"] = METHOD |
| logger.info("ST-TENP variant=%s temporal=%s redundancy=%s p_mid=%s", variant, temporal, redundancy, p_mid) |
| return meta |
|
|
