Buckets:
| """cwm3/tests/test_chunked_no_target_leak.py — the chunked decoder F never sees a target (plan §3 F, no teacher forcing). | |
| (a) signature: ChunkedDecoder.forward(z, expect_fn) — no target argument exists; the joint model's forward takes only the input batch. | |
| (b) gradient: with every target tensor requiring grad, d(loss of chunk c) / d(targets of chunk c' != c) is exactly zero for all c, c' — | |
| the only target that reaches a chunk's loss is its own supervision; in particular no FUTURE target (chunk >= c) reaches chunk c's | |
| emissions through the latent advance. | |
| (c) causality across chunks: perturbing the emissions of chunk c (via s_{c+1}) cannot change chunks <= c — emissions of chunk c depend | |
| on z and s_c only; verified by comparing a forward whose GRUCell is zeroed after chunk c. | |
| (d) expected-emission feed uses model outputs: replacing the targets by random tensors leaves every emission bitwise unchanged. | |
| (e) parameter matching: chunked vs direct within ±5 % for the three branch dims (build_decoder asserts; re-checked here). | |
| Run: cd cwm3 && python -m pytest tests/test_chunked_no_target_leak.py -q | |
| """ | |
| from __future__ import annotations | |
| import inspect | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import heads as hd # noqa: E402 | |
| from models.decoders import EXPECT, ChunkedDecoder, DirectDecoder25, build_decoder, match_hidden, chunked_params, N_CHUNKS, STEPS # noqa: E402 | |
| from models.cwm3 import CWM3Joint # noqa: E402 | |
| from tests.test_zero_init_identity import _batch, _fresh_branches # noqa: E402 | |
| def test_forward_signature_has_no_targets(): | |
| sig = inspect.signature(ChunkedDecoder.forward) | |
| assert list(sig.parameters) == ["self", "z", "expect_fn"] | |
| sig = inspect.signature(CWM3Joint.forward) | |
| assert list(sig.parameters) == ["self", "b", "return_trunks"] | |
| src = inspect.getsource(ChunkedDecoder.forward) | |
| for tok in ("target", "tg[", "zV", "zU", "zD", "Ystd", "Kstd"): | |
| assert tok not in src, f"decoder source mentions {tok!r}" | |
| def test_no_target_gradient_across_chunks(): | |
| torch.manual_seed(0) | |
| dec = ChunkedDecoder(d_in=32, D=hd.D_E, n_expect=EXPECT["E"][1], hidden=48) | |
| z = torch.randn(6, 32) | |
| em = dec(z, EXPECT["E"][0]) # [6,25,12] | |
| # targets as leaf tensors requiring grad: chain z, z_prev, mask | |
| zt = (torch.rand(6, 25, 6) > 0.5).float().requires_grad_(True) | |
| zp = torch.cat([torch.zeros(6, 1, 6), zt[:, :-1]], 1).detach().requires_grad_(True) | |
| m = torch.ones(6, 25, 6, requires_grad=True) | |
| c = hd.split_E(em)["chains"] | |
| for ch in range(N_CHUNKS): | |
| sl = slice(ch * STEPS, (ch + 1) * STEPS) | |
| loss_c = hd.chain_bce(c[:, sl, :, 0], c[:, sl, :, 1], zp[:, sl], zt[:, sl], m[:, sl]).sum() | |
| g_zt, g_zp, g_m = torch.autograd.grad(loss_c, [zt, zp, m], retain_graph=True, allow_unused=True) | |
| for name, g in (("z", g_zt), ("z_prev", g_zp), ("mask", g_m)): | |
| if g is None: | |
| continue | |
| other = torch.ones(25, dtype=torch.bool) | |
| other[sl] = False | |
| assert g[:, other].abs().sum() == 0, f"chunk {ch}: loss gradient reached {name} targets of other chunks" | |
| def test_chunk_causality_and_target_independence(): | |
| torch.manual_seed(1) | |
| dec = ChunkedDecoder(d_in=20, D=hd.D_U, n_expect=EXPECT["U"][1], hidden=40) | |
| z = torch.randn(5, 20) | |
| dec.eval() | |
| with torch.no_grad(): | |
| full = dec(z, EXPECT["U"][0]) | |
| # zero the latent advance after chunk c: chunks <= c must be unchanged, chunks > c generally change | |
| for c in range(N_CHUNKS - 1): | |
| def hooked(z_, expect_fn, c=c): | |
| B = z_.shape[0] | |
| s = dec.s0(z_) | |
| outs = [] | |
| for k in range(dec.n_chunks): | |
| ce = dec.chunk_emb.weight[k][None].expand(B, -1) | |
| em = dec.head(torch.cat([z_, s, ce], 1)).view(B, dec.steps, dec.D) | |
| outs.append(em) | |
| if k + 1 < dec.n_chunks: | |
| e = expect_fn(em).reshape(B, -1) | |
| adv = dec.cell(torch.cat([s, torch.nn.functional.gelu(dec.e_proj(e))], 1), s) | |
| s = s + (adv if k < c else torch.zeros_like(adv)) | |
| return torch.cat(outs, 1) | |
| alt = hooked(z, EXPECT["U"][0]) | |
| assert torch.equal(alt[:, :(c + 1) * STEPS], full[:, :(c + 1) * STEPS]) | |
| assert not torch.equal(alt[:, (c + 1) * STEPS:], full[:, (c + 1) * STEPS:]) | |
| # (d) whole joint model: random targets do not change emissions (targets are not inputs) | |
| V, U, E = _fresh_branches() | |
| b, tg = _batch() | |
| m = CWM3Joint(V, U, E, decoder="chunked", k_head=True, seed=0, tok_source="region") | |
| m.eval() | |
| with torch.no_grad(): | |
| o1 = m(b) | |
| for br in tg: | |
| for k, v in tg[br].items(): | |
| if v.dtype.is_floating_point: | |
| tg[br][k] = torch.rand_like(v) | |
| o2 = m(b) | |
| for k in ("V", "U", "E", "F"): | |
| assert torch.equal(o1[k], o2[k]) | |
| def test_parameter_matching(): | |
| for br, d_in, hidden in (("V", 272, 256), ("U", 144, 128), ("E", 272, 256)): | |
| dec, info = build_decoder("chunked", br, d_in, hidden, 0.0) | |
| assert abs(info["ratio"] - 1) <= 0.05, info | |
| assert info["params"] == chunked_params(d_in, {"V": hd.D_V, "U": hd.D_U, "E": hd.D_E}[br], EXPECT[br][1], info["hidden"], info["d_s"], info["d_e"]) | |
| direct, dinfo = build_decoder("direct_fresh", br, d_in, hidden, 0.0) | |
| assert dinfo["params"] == info["direct_params"] | |
| z = torch.randn(3, d_in) | |
| assert dec(z, EXPECT[br][0]).shape == direct(z).shape == (3, 25, {"V": hd.D_V, "U": hd.D_U, "E": hd.D_E}[br]) | |
Xet Storage Details
- Size:
- 6.33 kB
- Xet hash:
- e08c93190fdde1b0c582245721f212e963d285661e2d33a1a2ff76e089a0c8cf
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.