feat: isolate DeepSeek role marker head token

This commit is contained in:
wuyang
2026-07-29 20:44:43 +08:00
parent 8ef2d33437
commit ea49ed0511
15 changed files with 3676450 additions and 33 deletions
+54
View File
@@ -371,3 +371,57 @@ See `research/DEEPSEEK_ROUTING_HISTORY_BOUNDARY_TOKEN_AUDIT.md` for all paired
intervals, per-token route alignment, scope split, BF16 batch-shape audit,
primary literature, and the boundary between an input-ID intervention and a
chat-turn semantic claim.
## Role-marker-head single-token control
`v2_lite_routing_role_marker_head_control.py` keeps the official filler
history, EOS, colon, target span, generation prompt, mask, length, and 32-row
batch shape, while changing one ordinary token ID:
```text
system off/on × official / pre-target User→Assistant /
pre-target User→x / post-target Assistant→User
```
The pinned tokenizer maps `User`, `Assistant`, `:`, and `x` to IDs `5726`,
`77398`, `25`, and `87`. The pre-target controls identify the effect of the
first role-marker token only; the colon remains. The post-target replacement
is a causal suffix negative control. None of the three counterfactuals is a
valid official chat serialization.
```bash
PYTHONPATH=/path/to/transformers-4.41.2-deps:/usr/lib/python3/dist-packages \
python -B experiments/deepseek/v2_lite_routing_role_marker_head_control.py \
--artifact-dir /path/to/deepseek-v2-lite \
--human-eval /path/to/HumanEval.jsonl.gz \
--gsm8k /path/to/gsm8k/test.jsonl \
--tnews /path/to/tnews/test.json \
--tnews-archive /path/to/tnews_public.zip \
--wikitext /path/to/wikitext-validation.parquet \
--output src/data/deepseek-v2-lite-routing-role-marker-head-control.json \
--per-domain 32 \
--content-tokens 23 \
--batch-prompts 4 \
--layers 7 \
--bootstrap 2000 \
--seed 20260729 \
--captured-at 2026-07-29T12:30:00+00:00
```
The eight cells add 2,044,224 real top-6 route selections. All 256
source×system groups preserve length and target position; every one of the
768 counterfactuals differs from its official sequence at exactly one ID. The
committed run and independent rerun are byte-exact:
```text
9dc0e37fbce6581269428dcfcb84c7a17b466c5239c8171e6741f66d5eeb8caf
```
For exact target content, direct replacement TV is about `.02–.03`, but the
system-edge contrast has mixed direction: `12↑12↓` for User→Assistant and
`13↑11↓` for User→x. The post-target control is exact for all 34,488 aligned
target-token ordered top-6 routes, with zero target TV, JSD, and ΔCV. See
`research/DEEPSEEK_ROUTING_ROLE_MARKER_HEAD_AUDIT.md` for all paired intervals,
depth maps, token-level alignment, cross-experiment BF16 batch-content audit,
primary sources, and the boundary against full role semantics or Chat-model
behavior.
@@ -43,6 +43,7 @@ FILLER_USER = prior.FILLER_USER
FILLER_ASSISTANT = prior.FILLER_ASSISTANT
BOUNDARY_LEVELS = ("eos", "x", "period", "newline")
REFERENCE_LEVEL = "eos"
BOUNDARY_TEXT = {
"x": "x",
"period": ".",
@@ -297,7 +298,7 @@ def boundary_control_domain(
metric_system_edges: dict[str, Any] = {}
metric_system_edge_contrasts: dict[str, Any] = {}
for metric in point_metrics["s0_eos"]:
for metric in point_metrics[f"s0_{REFERENCE_LEVEL}"]:
point_edges = system_edges(
{
condition: point_metrics[condition][metric]
@@ -398,7 +399,7 @@ def boundary_control_domain(
direct_substitutions = {}
for boundary in BOUNDARY_LEVELS:
if boundary == "eos":
if boundary == REFERENCE_LEVEL:
continue
direct_substitutions[boundary] = {}
direct_point_tv: dict[int, float] = {}
@@ -406,7 +407,7 @@ def boundary_control_domain(
direct_point_jsd: dict[int, float] = {}
direct_boot_jsd: dict[int, np.ndarray] = {}
for system in (0, 1):
before = f"s{system}_eos"
before = f"s{system}_{REFERENCE_LEVEL}"
after = f"s{system}_{boundary}"
point_delta = point[after] - point[before]
tv_boot = 0.5 * np.abs(
@@ -0,0 +1,573 @@
#!/usr/bin/env python3
"""Run a 2 x 4 role-marker-head control on DeepSeek-V2-Lite.
Every condition keeps the same repeated-token user/assistant history, the
official assistant EOS boundary, and the same target user content. The system
factor is off/on. The second factor is one of four role-head conditions:
official: target ``User`` and suffix ``Assistant`` stay official
target_assistant: target ``User`` -> ``Assistant`` (one pre-target ID)
target_x: target ``User`` -> ``x`` (one pre-target ID)
suffix_user: suffix ``Assistant`` -> ``User`` (one post-target ID)
The colon token, sequence length, target position, attention mask, official
EOS, and padded batch shape are held fixed. The post-target suffix edit is a
causal-mask negative control: it must not alter earlier target-content routes.
All three edited sequences are explicit token-ID counterfactuals rather than
valid official chat serializations.
The official BF16 forward path and shared source-bootstrap statistics are
reused from the audited boundary-token runner. This file owns the renderer,
role-head contract, result schema, and causal-suffix validation.
"""
from __future__ import annotations
import contextlib
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any
import v2_lite_routing_history_boundary_token_control as prior
base = prior.base
SYSTEM_MESSAGE = prior.SYSTEM_MESSAGE
FILLER_USER = prior.FILLER_USER
FILLER_ASSISTANT = prior.FILLER_ASSISTANT
ROLE_LEVELS = (
"official",
"target_assistant",
"target_x",
"suffix_user",
)
REFERENCE_LEVEL = "official"
CONDITIONS = tuple(
f"s{system}_{level}"
for level in ROLE_LEVELS
for system in (0, 1)
)
FACTORS = {
condition: {
"system": int(condition[1]),
"history": "filler",
"assistant_boundary": "official_eos",
"role_head": condition.split("_", 1)[1],
}
for condition in CONDITIONS
}
SYSTEM_CELLS = {
level: (f"s0_{level}", f"s1_{level}")
for level in ROLE_LEVELS
}
SYSTEM_EDGE_CONTRASTS = {
f"{level}_minus_official": (REFERENCE_LEVEL, level)
for level in ROLE_LEVELS
if level != REFERENCE_LEVEL
}
COMPARISONS = tuple(
[
(f"system_{level}", f"s0_{level}", f"s1_{level}")
for level in ROLE_LEVELS
]
+ [
(
f"{level}_at_s{system}",
f"s{system}_{REFERENCE_LEVEL}",
f"s{system}_{level}",
)
for level in ROLE_LEVELS
if level != REFERENCE_LEVEL
for system in (0, 1)
]
)
ALIGNMENT_COMPARISONS = COMPARISONS
RENDER_AUDIT: list[dict[str, Any]] = []
ROLE_TOKEN_IDS: dict[str, int] = {}
def condition_messages(
content: str,
condition: str,
) -> list[dict[str, str]]:
factors = FACTORS[condition]
messages: list[dict[str, str]] = []
if factors["system"]:
messages.append({"role": "system", "content": SYSTEM_MESSAGE})
messages.extend(
[
{"role": "user", "content": FILLER_USER},
{"role": "assistant", "content": FILLER_ASSISTANT},
{"role": "user", "content": content},
]
)
return messages
def role_token_ids(tokenizer: Any) -> dict[str, int]:
texts = {
"user": "User",
"assistant": "Assistant",
"x": "x",
"colon": ":",
"eos": tokenizer.eos_token,
}
ids: dict[str, int] = {}
for name, text in texts.items():
encoded = list(
tokenizer(text, add_special_tokens=False).input_ids
)
if len(encoded) != 1:
raise RuntimeError(
f"{name} role control is not one token: {encoded}"
)
ids[name] = int(encoded[0])
if ids["eos"] != tokenizer.eos_token_id:
raise RuntimeError("encoded EOS differs from tokenizer.eos_token_id")
if any(
ids[name] in tokenizer.all_special_ids
for name in ("user", "assistant", "x", "colon")
):
raise RuntimeError("ordinary role-head control is unexpectedly special")
if len({ids["user"], ids["assistant"], ids["x"]}) != 3:
raise RuntimeError(f"role-head token IDs are not distinct: {ids}")
return ids
def render_role_variant(
tokenizer: Any,
content: str,
condition: str,
) -> dict[str, Any]:
messages = condition_messages(content, condition)
rendered = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
official_ids = list(
tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
)
token_ids, offsets = base.tokenize_with_offsets(
tokenizer,
rendered,
add_special_tokens=False,
)
if token_ids != official_ids:
raise RuntimeError(
f"{condition} rendered IDs differ from apply_chat_template"
)
ids = role_token_ids(tokenizer)
ROLE_TOKEN_IDS.update(ids)
content_start = rendered.rfind(content)
if content_start < 0:
raise RuntimeError(f"{condition} target content is absent")
content_end = content_start + len(content)
positions, records, crossing = base.content_positions(
token_ids,
offsets,
content_start,
content_end,
)
if not positions:
raise RuntimeError(f"{condition} has no target-content tokens")
target_first = min(positions)
target_last = max(positions)
target_user_candidates = [
index
for index, token_id in enumerate(official_ids)
if token_id == ids["user"] and index < target_first
]
if len(target_user_candidates) < 2:
raise RuntimeError(
f"{condition} cannot locate filler and target User heads"
)
target_user_position = max(target_user_candidates)
suffix_assistant_candidates = [
index
for index, token_id in enumerate(official_ids)
if token_id == ids["assistant"] and index > target_last
]
if len(suffix_assistant_candidates) != 1:
raise RuntimeError(
f"{condition} expected one suffix Assistant head; "
f"got {suffix_assistant_candidates}"
)
suffix_assistant_position = suffix_assistant_candidates[0]
if official_ids[target_user_position + 1] != ids["colon"]:
raise RuntimeError("target User head is not followed by colon")
if official_ids[suffix_assistant_position + 1] != ids["colon"]:
raise RuntimeError("suffix Assistant head is not followed by colon")
if ids["eos"] not in official_ids[:target_user_position]:
raise RuntimeError("official history EOS is absent before target role")
level = FACTORS[condition]["role_head"]
edit_position: int | None = None
before_id: int | None = None
after_id: int | None = None
edit_side = "none"
if level == "target_assistant":
edit_position = target_user_position
before_id = ids["user"]
after_id = ids["assistant"]
edit_side = "pre_target"
elif level == "target_x":
edit_position = target_user_position
before_id = ids["user"]
after_id = ids["x"]
edit_side = "pre_target"
elif level == "suffix_user":
edit_position = suffix_assistant_position
before_id = ids["assistant"]
after_id = ids["user"]
edit_side = "post_target"
elif level != REFERENCE_LEVEL:
raise RuntimeError(f"unknown role-head level: {level}")
token_ids = list(token_ids)
if edit_position is not None:
if token_ids[edit_position] != before_id:
raise RuntimeError(
f"{condition} edit source ID differs at {edit_position}"
)
token_ids[edit_position] = int(after_id)
differing_ids = sum(
left != right
for left, right in zip(official_ids, token_ids, strict=True)
)
expected_differences = 0 if level == REFERENCE_LEVEL else 1
if differing_ids != expected_differences:
raise RuntimeError(
f"{condition} changed {differing_ids} IDs; "
f"expected {expected_differences}"
)
if edit_side == "pre_target" and edit_position >= target_first:
raise RuntimeError("pre-target edit is not before target content")
if edit_side == "post_target" and edit_position <= target_last:
raise RuntimeError("post-target edit is not after target content")
decoded = tokenizer.decode(
token_ids,
skip_special_tokens=False,
clean_up_tokenization_spaces=False,
)
RENDER_AUDIT.append(
{
"content_sha256": base.text_sha256(content),
"condition": condition,
"role_head": level,
"edit_side": edit_side,
"edit_position": edit_position,
"before_token_id": before_id,
"after_token_id": after_id,
"target_user_position": target_user_position,
"suffix_assistant_position": suffix_assistant_position,
"official_eos_token_id": ids["eos"],
"input_tokens": len(token_ids),
"target_first_position": target_first,
"target_last_position": target_last,
"changed_token_ids_vs_official": differing_ids,
"counterfactual_decoded_sha256": base.text_sha256(decoded),
}
)
return {
"condition": condition,
"messages_sha256": base.canonical_hash(messages),
"rendered_sha256": base.text_sha256(rendered),
"token_ids": token_ids,
"tokens": len(token_ids),
"token_ids_sha256": base.canonical_hash(token_ids),
"content_positions": positions,
"content_records": records,
"content_tokens": len(positions),
"wrapper_tokens": len(token_ids) - len(positions),
"boundary_crossing_tokens": crossing,
}
def layer_statistics(
prompt_rows: list[dict[str, Any]],
replicates: int,
seed: int,
layer_index: int,
) -> dict[str, Any]:
scopes = prior.layer_statistics(
prompt_rows,
replicates,
seed,
layer_index,
)
for scope in scopes.values():
for mode in scope["modes"].values():
mode["role_marker_control"] = mode.pop("boundary_control")
return scopes
def install_control_contract() -> None:
"""Install the eight-cell role renderer and shared statistics."""
prior.BOUNDARY_LEVELS = ROLE_LEVELS
prior.REFERENCE_LEVEL = REFERENCE_LEVEL
prior.CONDITIONS = CONDITIONS
prior.FACTORS = FACTORS
prior.SYSTEM_CELLS = SYSTEM_CELLS
prior.SYSTEM_EDGE_CONTRASTS = SYSTEM_EDGE_CONTRASTS
prior.COMPARISONS = COMPARISONS
prior.ALIGNMENT_COMPARISONS = ALIGNMENT_COMPARISONS
base.CONDITIONS = CONDITIONS
base.FACTORS = FACTORS
base.COMPARISONS = COMPARISONS
base.ALIGNMENT_COMPARISONS = ALIGNMENT_COMPARISONS
base.condition_messages = condition_messages
base.render_variant = render_role_variant
base.layer_statistics = layer_statistics
def output_path_from_argv() -> Path:
try:
return Path(sys.argv[sys.argv.index("--output") + 1])
except (ValueError, IndexError) as error:
raise ValueError("--output is required") from error
def render_contract_summary() -> dict[str, Any]:
if not RENDER_AUDIT:
raise RuntimeError("render audit is empty")
by_content: dict[str, list[dict[str, Any]]] = {}
for row in RENDER_AUDIT:
by_content.setdefault(row["content_sha256"], []).append(row)
if len(by_content) != 128 and "--per-domain" not in sys.argv:
raise RuntimeError(
f"expected 128 selected contents; observed {len(by_content)}"
)
equal_lengths = 0
equal_target_positions = 0
official_exact = 0
one_id_replacements = 0
pre_target_edits = 0
post_target_edits = 0
for rows in by_content.values():
for system in (0, 1):
cells = [
row
for row in rows
if FACTORS[row["condition"]]["system"] == system
]
if len(cells) != len(ROLE_LEVELS):
raise RuntimeError(
"render audit lacks one or more role-head cells"
)
if len({row["input_tokens"] for row in cells}) == 1:
equal_lengths += 1
target_spans = {
(
row["target_first_position"],
row["target_last_position"],
)
for row in cells
}
if len(target_spans) == 1:
equal_target_positions += 1
official_exact += sum(
row["role_head"] == REFERENCE_LEVEL
and row["changed_token_ids_vs_official"] == 0
for row in cells
)
one_id_replacements += sum(
row["role_head"] != REFERENCE_LEVEL
and row["changed_token_ids_vs_official"] == 1
for row in cells
)
pre_target_edits += sum(
row["edit_side"] == "pre_target"
for row in cells
)
post_target_edits += sum(
row["edit_side"] == "post_target"
for row in cells
)
groups = 2 * len(by_content)
replacements = 3 * groups
return {
"selected_contents": len(by_content),
"system_groups": groups,
"equal_input_length_groups": equal_lengths,
"equal_target_position_groups": equal_target_positions,
"official_cells_exact": official_exact,
"one_id_counterfactual_cells_exact": one_id_replacements,
"expected_one_id_counterfactual_cells": replacements,
"pre_target_one_id_edits": pre_target_edits,
"post_target_one_id_edits": post_target_edits,
"all_group_lengths_equal": equal_lengths == groups,
"all_target_positions_equal": equal_target_positions == groups,
"all_counterfactuals_change_exactly_one_id": (
one_id_replacements == replacements
),
}
def finalize_result(path: Path) -> dict[str, Any]:
result = json.loads(path.read_text(encoding="utf-8"))
result["schema_version"] = 3
result["evidence_identity"] = (
"X / official BF16 weights and tokenizer; official role heads plus "
"paired one-token pre-target and post-target role-head "
"counterfactuals on local truncated forward"
)
boundary = result["boundary"]
boundary.pop("factorial_claim", None)
boundary.update(
{
"role_marker_head_control": True,
"official_serialization_by_role_head": {
"official": True,
"target_assistant": False,
"target_x": False,
"suffix_user": False,
},
"single_input_id_intervention": True,
"official_assistant_eos_held_fixed": True,
"colon_token_held_fixed": True,
"target_position_held_fixed": True,
"causal_suffix_negative_control": True,
"complete_role_semantics_identified": False,
"task_performance": False,
"causal_boundary": (
"official-to-control comparisons intervene on exactly one "
"role-head input ID in this fixed forward contract; the "
"colon and all other protocol tokens remain, the edited "
"sequences are not official chats, and no answer quality "
"is measured"
),
}
)
old_contract = result.pop("message_history_contract")
render_validation = render_contract_summary()
result["role_marker_head_contract"] = {
"chat_template_revision": base.MODEL_REVISION,
"chat_template": old_contract["chat_template"],
"chat_template_sha256": old_contract["chat_template_sha256"],
"official_sequence": (
"Assistant: {filler} + eos_token + User: {target} + "
"generation-prompt Assistant:"
),
"system_message": SYSTEM_MESSAGE,
"system_message_sha256": base.text_sha256(SYSTEM_MESSAGE),
"filler_user": FILLER_USER,
"filler_user_sha256": base.text_sha256(FILLER_USER),
"filler_assistant": FILLER_ASSISTANT,
"filler_assistant_sha256": base.text_sha256(FILLER_ASSISTANT),
"role_levels": list(ROLE_LEVELS),
"role_token_ids": ROLE_TOKEN_IDS,
"conditions": FACTORS,
"comparisons": [
{"name": name, "before": before, "after": after}
for name, before, after in COMPARISONS
],
"system_edge_contrasts": {
name: {
"before_role_head": before,
"after_role_head": after,
"definition": (
f"system edge at {after} minus system edge at {before}"
),
}
for name, (before, after) in (
SYSTEM_EDGE_CONTRASTS.items()
)
},
"render_validation": render_validation,
"target_role": "user",
"add_generation_prompt": True,
"scope_split": {
"full_input": (
"all official or counterfactually edited BOS, system/history, "
"target, newline, and generation-prompt token IDs"
),
"target_content": (
"exact intersection of (relative character span, token ID) "
"inside target user content across all eight conditions"
),
},
}
inference = result["inference_contract"]
inference["batch_grouping"] = (
"all eight role-head variants of one source prompt execute in the "
"same right-padded batch"
)
statistical = result["statistical_contract"]
statistical["paired_indices"] = (
"one sampled source-prompt index matrix is reused across all eight "
"cells for every role-head contrast within each "
"domain/layer/scope/mode"
)
statistical.pop("interaction_distribution_magnitude", None)
statistical["system_edge_contrast_distribution_magnitude"] = (
"0.5 * L1 norm of the signed difference between two system-edge "
"expert-share vectors; this is not labeled standard TV"
)
path.write_text(
json.dumps(result, indent=2, ensure_ascii=False) + "\n",
encoding="utf-8",
)
return result
def main() -> None:
install_control_contract()
if {"-h", "--help"} & set(sys.argv[1:]):
base.main()
return
output = output_path_from_argv()
with open(os.devnull, "w", encoding="utf-8") as sink:
with contextlib.redirect_stdout(sink):
base.main()
result = finalize_result(output)
payload = output.read_bytes()
print(
json.dumps(
{
"output": str(output),
"sha256": hashlib.sha256(payload).hexdigest(),
"bytes": len(payload),
"source_prompts": result["inference_contract"][
"total_source_prompts"
],
"prompt_variants": result["inference_contract"][
"total_prompt_variants"
],
"input_tokens_by_condition": result[
"inference_contract"
]["input_tokens_by_condition"],
"total_routes": result["inference_contract"][
"total_routes_all_conditions_all_moe_layers"
],
"render_validation": result[
"role_marker_head_contract"
]["render_validation"],
},
indent=2,
ensure_ascii=False,
)
)
if __name__ == "__main__":
main()