feat: add DeepSeek equal-length history control
This commit is contained in:
wuyang
@@ -0,0 +1,552 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run a 2 x 3 message-history control on DeepSeek-V2-Lite.
|
||||
|
||||
The system factor is off/on. The history factor has three levels:
|
||||
|
||||
none: no completed turn before the target user message
|
||||
filler: a repeated-token user/assistant turn
|
||||
demo: the fixed meaningful user/assistant turn from the prior probe
|
||||
|
||||
The filler and demo histories add exactly the same number of official-template
|
||||
tokens and keep the same user role, assistant role, assistant EOS, target
|
||||
position, and generation prompt. This tests whether the prior attenuation of
|
||||
the system edge requires the specific demo text. It does not identify a pure
|
||||
distance effect because filler token identity and the added history remain
|
||||
coupled.
|
||||
|
||||
The layer forward path is intentionally reused from the preceding audited
|
||||
probe. This file owns the six-cell condition renderer, the shared-bootstrap
|
||||
2 x 3 statistics, the control metadata, and the finalized result contract.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
import v2_lite_routing_history_factorial_probe as base
|
||||
|
||||
|
||||
SYSTEM_MESSAGE = base.SYSTEM_MESSAGE
|
||||
DEMO_USER = base.DEMO_USER
|
||||
DEMO_ASSISTANT = base.DEMO_ASSISTANT
|
||||
FILLER_USER = "x x x x x x x x x"
|
||||
FILLER_ASSISTANT = "x"
|
||||
|
||||
HISTORY_LEVELS = ("none", "filler", "demo")
|
||||
CONDITIONS = (
|
||||
"s0_none",
|
||||
"s1_none",
|
||||
"s0_filler",
|
||||
"s1_filler",
|
||||
"s0_demo",
|
||||
"s1_demo",
|
||||
)
|
||||
FACTORS = {
|
||||
"s0_none": {"system": 0, "history": "none"},
|
||||
"s1_none": {"system": 1, "history": "none"},
|
||||
"s0_filler": {"system": 0, "history": "filler"},
|
||||
"s1_filler": {"system": 1, "history": "filler"},
|
||||
"s0_demo": {"system": 0, "history": "demo"},
|
||||
"s1_demo": {"system": 1, "history": "demo"},
|
||||
}
|
||||
SYSTEM_CELLS = {
|
||||
"none": ("s0_none", "s1_none"),
|
||||
"filler": ("s0_filler", "s1_filler"),
|
||||
"demo": ("s0_demo", "s1_demo"),
|
||||
}
|
||||
COMPARISONS = (
|
||||
("system_none", "s0_none", "s1_none"),
|
||||
("system_filler", "s0_filler", "s1_filler"),
|
||||
("system_demo", "s0_demo", "s1_demo"),
|
||||
("filler_at_s0", "s0_none", "s0_filler"),
|
||||
("filler_at_s1", "s1_none", "s1_filler"),
|
||||
("demo_at_s0", "s0_none", "s0_demo"),
|
||||
("demo_at_s1", "s1_none", "s1_demo"),
|
||||
("demo_vs_filler_s0", "s0_filler", "s0_demo"),
|
||||
("demo_vs_filler_s1", "s1_filler", "s1_demo"),
|
||||
)
|
||||
ALIGNMENT_COMPARISONS = (
|
||||
COMPARISONS[0],
|
||||
COMPARISONS[1],
|
||||
COMPARISONS[2],
|
||||
COMPARISONS[7],
|
||||
COMPARISONS[8],
|
||||
)
|
||||
SYSTEM_EDGE_CONTRASTS = {
|
||||
"filler_minus_none": ("none", "filler"),
|
||||
"demo_minus_none": ("none", "demo"),
|
||||
"demo_minus_filler": ("filler", "demo"),
|
||||
}
|
||||
|
||||
|
||||
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})
|
||||
if factors["history"] == "filler":
|
||||
messages.extend(
|
||||
[
|
||||
{"role": "user", "content": FILLER_USER},
|
||||
{"role": "assistant", "content": FILLER_ASSISTANT},
|
||||
]
|
||||
)
|
||||
elif factors["history"] == "demo":
|
||||
messages.extend(
|
||||
[
|
||||
{"role": "user", "content": DEMO_USER},
|
||||
{"role": "assistant", "content": DEMO_ASSISTANT},
|
||||
]
|
||||
)
|
||||
messages.append({"role": "user", "content": content})
|
||||
return messages
|
||||
|
||||
|
||||
def history_control_domain(
|
||||
loads: dict[str, np.ndarray],
|
||||
mode: str,
|
||||
replicates: int,
|
||||
seed: int,
|
||||
scope: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Compute all six cells with one shared source-bootstrap matrix."""
|
||||
shapes = {value.shape for value in loads.values()}
|
||||
if len(shapes) != 1:
|
||||
raise ValueError(f"history-control shape mismatch: {sorted(shapes)}")
|
||||
rows = next(iter(loads.values())).shape[0]
|
||||
rng = np.random.default_rng(base.scoped_seed(seed, scope))
|
||||
sampled = rng.integers(
|
||||
0,
|
||||
rows,
|
||||
size=(replicates, rows),
|
||||
endpoint=False,
|
||||
)
|
||||
point = {
|
||||
condition: base.distribution(value, mode)
|
||||
for condition, value in loads.items()
|
||||
}
|
||||
boot = {
|
||||
condition: base.bootstrap_distributions(value, mode, sampled)
|
||||
for condition, value in loads.items()
|
||||
}
|
||||
point_metrics = {
|
||||
condition: base.metric_vector(value)
|
||||
for condition, value in point.items()
|
||||
}
|
||||
boot_metrics = {
|
||||
condition: base.metric_vector(value)
|
||||
for condition, value in boot.items()
|
||||
}
|
||||
|
||||
def system_edges(
|
||||
values: dict[str, np.ndarray],
|
||||
) -> dict[str, np.ndarray]:
|
||||
return {
|
||||
history: values[after] - values[before]
|
||||
for history, (before, after) in SYSTEM_CELLS.items()
|
||||
}
|
||||
|
||||
def edge_contrasts(
|
||||
edges: dict[str, np.ndarray],
|
||||
) -> dict[str, np.ndarray]:
|
||||
return {
|
||||
name: edges[after] - edges[before]
|
||||
for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items()
|
||||
}
|
||||
|
||||
metric_system_edges: dict[str, Any] = {}
|
||||
metric_system_edge_contrasts: dict[str, Any] = {}
|
||||
for metric in point_metrics["s0_none"]:
|
||||
point_edges = system_edges(
|
||||
{
|
||||
condition: point_metrics[condition][metric]
|
||||
for condition in CONDITIONS
|
||||
}
|
||||
)
|
||||
boot_edges = system_edges(
|
||||
{
|
||||
condition: boot_metrics[condition][metric]
|
||||
for condition in CONDITIONS
|
||||
}
|
||||
)
|
||||
metric_system_edges[metric] = {
|
||||
history: {
|
||||
"point": float(point_edges[history][0]),
|
||||
"ci95": base.interval(boot_edges[history]),
|
||||
}
|
||||
for history in HISTORY_LEVELS
|
||||
}
|
||||
point_contrasts = edge_contrasts(point_edges)
|
||||
boot_contrasts = edge_contrasts(boot_edges)
|
||||
metric_system_edge_contrasts[metric] = {
|
||||
name: {
|
||||
"point": float(point_contrasts[name][0]),
|
||||
"ci95": base.interval(boot_contrasts[name]),
|
||||
}
|
||||
for name in SYSTEM_EDGE_CONTRASTS
|
||||
}
|
||||
|
||||
point_vectors = system_edges(point)
|
||||
boot_vectors = system_edges(boot)
|
||||
point_vector_contrasts = edge_contrasts(point_vectors)
|
||||
boot_vector_contrasts = edge_contrasts(boot_vectors)
|
||||
distribution_system_edge_contrasts = {}
|
||||
for name in SYSTEM_EDGE_CONTRASTS:
|
||||
point_magnitude = 0.5 * np.abs(
|
||||
point_vector_contrasts[name]
|
||||
).sum()
|
||||
boot_magnitude = 0.5 * np.abs(
|
||||
boot_vector_contrasts[name]
|
||||
).sum(axis=1)
|
||||
distribution_system_edge_contrasts[name] = {
|
||||
"half_l1_magnitude": {
|
||||
"point": float(point_magnitude),
|
||||
"ci95": base.interval(boot_magnitude),
|
||||
},
|
||||
"expert_share_difference_in_system_edges": (
|
||||
point_vector_contrasts[name].tolist()
|
||||
),
|
||||
"expert_share_difference_in_system_edges_ci95": (
|
||||
base.interval(boot_vector_contrasts[name])
|
||||
),
|
||||
}
|
||||
|
||||
system_edge_distances: dict[str, Any] = {}
|
||||
point_tv: dict[str, float] = {}
|
||||
boot_tv: dict[str, np.ndarray] = {}
|
||||
point_jsd: dict[str, float] = {}
|
||||
boot_jsd: dict[str, np.ndarray] = {}
|
||||
for history, (before, after) in SYSTEM_CELLS.items():
|
||||
point_delta = point[after] - point[before]
|
||||
point_tv[history] = float(0.5 * np.abs(point_delta).sum())
|
||||
boot_tv[history] = 0.5 * np.abs(
|
||||
boot[after] - boot[before]
|
||||
).sum(axis=1)
|
||||
point_jsd[history] = float(
|
||||
base.js_divergence(point[before], point[after])[0]
|
||||
)
|
||||
boot_jsd[history] = base.js_divergence(
|
||||
boot[before],
|
||||
boot[after],
|
||||
)
|
||||
system_edge_distances[history] = {
|
||||
"total_variation": {
|
||||
"point": point_tv[history],
|
||||
"ci95": base.interval(boot_tv[history]),
|
||||
},
|
||||
"js_divergence": {
|
||||
"point": point_jsd[history],
|
||||
"ci95": base.interval(boot_jsd[history]),
|
||||
"unit": "nats",
|
||||
},
|
||||
}
|
||||
|
||||
system_edge_distance_contrasts = {}
|
||||
for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items():
|
||||
system_edge_distance_contrasts[name] = {
|
||||
"total_variation_delta": {
|
||||
"point": point_tv[after] - point_tv[before],
|
||||
"ci95": base.interval(boot_tv[after] - boot_tv[before]),
|
||||
},
|
||||
"js_divergence_delta": {
|
||||
"point": point_jsd[after] - point_jsd[before],
|
||||
"ci95": base.interval(boot_jsd[after] - boot_jsd[before]),
|
||||
"unit": "nats",
|
||||
},
|
||||
}
|
||||
|
||||
lexical_replacement = {}
|
||||
for system in (0, 1):
|
||||
before = f"s{system}_filler"
|
||||
after = f"s{system}_demo"
|
||||
point_delta = point[after] - point[before]
|
||||
tv_boot = 0.5 * np.abs(
|
||||
boot[after] - boot[before]
|
||||
).sum(axis=1)
|
||||
jsd_boot = base.js_divergence(boot[before], boot[after])
|
||||
lexical_replacement[f"at_s{system}"] = {
|
||||
"total_variation": {
|
||||
"point": float(0.5 * np.abs(point_delta).sum()),
|
||||
"ci95": base.interval(tv_boot),
|
||||
},
|
||||
"js_divergence": {
|
||||
"point": float(
|
||||
base.js_divergence(point[before], point[after])[0]
|
||||
),
|
||||
"ci95": base.interval(jsd_boot),
|
||||
"unit": "nats",
|
||||
},
|
||||
}
|
||||
|
||||
return {
|
||||
"metric_system_edges": metric_system_edges,
|
||||
"metric_system_edge_contrasts": metric_system_edge_contrasts,
|
||||
"system_edge_distances": system_edge_distances,
|
||||
"system_edge_distance_contrasts": (
|
||||
system_edge_distance_contrasts
|
||||
),
|
||||
"distribution_system_edge_contrasts": (
|
||||
distribution_system_edge_contrasts
|
||||
),
|
||||
"lexical_replacement": lexical_replacement,
|
||||
}
|
||||
|
||||
|
||||
def layer_statistics(
|
||||
prompt_rows: list[dict[str, Any]],
|
||||
replicates: int,
|
||||
seed: int,
|
||||
layer_index: int,
|
||||
) -> dict[str, Any]:
|
||||
scopes = {}
|
||||
for load_scope, load_key in (
|
||||
("full_input", "full_load"),
|
||||
("target_content", "content_load"),
|
||||
):
|
||||
modes = {}
|
||||
for mode in ("token_weighted", "prompt_balanced"):
|
||||
conditions: dict[str, Any] = {}
|
||||
domain_loads: dict[str, dict[str, np.ndarray]] = {}
|
||||
for domain in base.DOMAIN_ORDER:
|
||||
rows = [
|
||||
row for row in prompt_rows
|
||||
if row["domain"] == domain
|
||||
]
|
||||
domain_loads[domain] = {
|
||||
condition: np.asarray(
|
||||
[
|
||||
row["conditions"][condition][load_key]
|
||||
for row in rows
|
||||
],
|
||||
dtype=np.int64,
|
||||
)
|
||||
for condition in CONDITIONS
|
||||
}
|
||||
for condition in CONDITIONS:
|
||||
conditions.setdefault(condition, {})[domain] = (
|
||||
base.bootstrap_domain(
|
||||
domain_loads[domain][condition],
|
||||
mode,
|
||||
replicates,
|
||||
seed,
|
||||
(
|
||||
f"layer={layer_index}|scope={load_scope}|"
|
||||
f"mode={mode}|condition={condition}|"
|
||||
f"domain={domain}"
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
comparisons: dict[str, Any] = {}
|
||||
for comparison, before, after in COMPARISONS:
|
||||
comparisons[comparison] = {}
|
||||
for domain in base.DOMAIN_ORDER:
|
||||
comparisons[comparison][domain] = base.paired_domain(
|
||||
domain_loads[domain][before],
|
||||
domain_loads[domain][after],
|
||||
mode,
|
||||
replicates,
|
||||
seed,
|
||||
(
|
||||
f"layer={layer_index}|scope={load_scope}|"
|
||||
f"mode={mode}|comparison={comparison}|"
|
||||
f"domain={domain}"
|
||||
),
|
||||
)
|
||||
|
||||
control = {
|
||||
domain: history_control_domain(
|
||||
domain_loads[domain],
|
||||
mode,
|
||||
replicates,
|
||||
seed,
|
||||
(
|
||||
f"layer={layer_index}|scope={load_scope}|"
|
||||
f"mode={mode}|history_control|domain={domain}"
|
||||
),
|
||||
)
|
||||
for domain in base.DOMAIN_ORDER
|
||||
}
|
||||
modes[mode] = {
|
||||
"conditions": conditions,
|
||||
"comparisons": comparisons,
|
||||
"history_control": control,
|
||||
}
|
||||
scopes[load_scope] = {"modes": modes}
|
||||
return scopes
|
||||
|
||||
|
||||
def install_control_contract() -> None:
|
||||
"""Install the six-cell renderer/statistics into the audited runner."""
|
||||
base.CONDITIONS = CONDITIONS
|
||||
base.FACTORS = FACTORS
|
||||
base.COMPARISONS = COMPARISONS
|
||||
base.ALIGNMENT_COMPARISONS = ALIGNMENT_COMPARISONS
|
||||
base.condition_messages = condition_messages
|
||||
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 finalize_result(path: Path) -> dict[str, Any]:
|
||||
result = json.loads(path.read_text(encoding="utf-8"))
|
||||
result["evidence_identity"] = (
|
||||
"X / official BF16 weights, official tokenizer chat template, "
|
||||
"paired 2x3 history-distance control on local truncated forward"
|
||||
)
|
||||
boundary = result["boundary"]
|
||||
boundary.pop("factorial_claim", None)
|
||||
boundary.update(
|
||||
{
|
||||
"history_control_claim": (
|
||||
"filler and demo histories have identical official-template "
|
||||
"token increments, roles, assistant EOS, and target position; "
|
||||
"their text identities differ"
|
||||
),
|
||||
"filler_is_semantics_free": False,
|
||||
"pure_distance_isolated": False,
|
||||
"lexical_replacement_control": True,
|
||||
"causal_boundary": (
|
||||
"none-to-filler still couples added history, repeated filler "
|
||||
"tokens, and distance; filler-to-demo isolates replacement of "
|
||||
"the fixed history text only within this protocol"
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
old_contract = result.pop("message_history_contract")
|
||||
selected = result["corpus_contract"]["selected"]
|
||||
increments = {}
|
||||
for system in (0, 1):
|
||||
base_condition = f"s{system}_none"
|
||||
for history in ("filler", "demo"):
|
||||
condition = f"s{system}_{history}"
|
||||
deltas = [
|
||||
row["conditions"][condition]["tokens"]
|
||||
- row["conditions"][base_condition]["tokens"]
|
||||
for row in selected
|
||||
]
|
||||
increments[f"{condition}_minus_{base_condition}"] = {
|
||||
"min": min(deltas),
|
||||
"max": max(deltas),
|
||||
"all_equal": len(set(deltas)) == 1,
|
||||
}
|
||||
result["history_control_contract"] = {
|
||||
"chat_template_revision": base.MODEL_REVISION,
|
||||
"chat_template": old_contract["chat_template"],
|
||||
"chat_template_sha256": old_contract["chat_template_sha256"],
|
||||
"system_message": SYSTEM_MESSAGE,
|
||||
"system_message_sha256": base.text_sha256(SYSTEM_MESSAGE),
|
||||
"demo_user": DEMO_USER,
|
||||
"demo_user_sha256": base.text_sha256(DEMO_USER),
|
||||
"demo_assistant": DEMO_ASSISTANT,
|
||||
"demo_assistant_sha256": base.text_sha256(DEMO_ASSISTANT),
|
||||
"filler_user": FILLER_USER,
|
||||
"filler_user_sha256": base.text_sha256(FILLER_USER),
|
||||
"filler_assistant": FILLER_ASSISTANT,
|
||||
"filler_assistant_sha256": base.text_sha256(FILLER_ASSISTANT),
|
||||
"target_role": "user",
|
||||
"add_generation_prompt": True,
|
||||
"conditions": FACTORS,
|
||||
"comparisons": [
|
||||
{"name": name, "before": before, "after": after}
|
||||
for name, before, after in COMPARISONS
|
||||
],
|
||||
"system_edge_contrasts": {
|
||||
name: {
|
||||
"before_history": before,
|
||||
"after_history": after,
|
||||
"definition": (
|
||||
f"system edge at {after} minus system edge at {before}"
|
||||
),
|
||||
}
|
||||
for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items()
|
||||
},
|
||||
"token_increment_validation": increments,
|
||||
"scope_split": {
|
||||
"full_input": (
|
||||
"all rendered BOS, system/history, target, newlines, EOS, "
|
||||
"and generation-prompt tokens"
|
||||
),
|
||||
"target_content": (
|
||||
"exact intersection of (relative character span, token ID) "
|
||||
"inside target user content across all six conditions"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
inference = result["inference_contract"]
|
||||
inference["batch_grouping"] = (
|
||||
"all six control 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 six "
|
||||
"cells for every history-control 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"
|
||||
)
|
||||
result["schema_version"] = 2
|
||||
path.write_text(
|
||||
json.dumps(result, indent=2, ensure_ascii=False) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def main() -> None:
|
||||
install_control_contract()
|
||||
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"
|
||||
],
|
||||
},
|
||||
indent=2,
|
||||
ensure_ascii=False,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user