stackformers 4.3.0__tar.gz → 4.3.1__tar.gz

{stackformers-4.3.0 → stackformers-4.3.1}/CHANGELOG.md

@@ -6,6 +6,27 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
 adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
 MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
 
+## [4.3.1] — 2026-07-17
+
+### Fixed
+
+- **`MLMWrapper`'s masked pass leaked gradient to whatever produced `input.x`, despite
+  the target being detached.** `corrupted_x` was built from the live (undetached)
+  `input.x` at unmasked positions; self-attention mixes those into the masked
+  positions' predictions, so `mlm_loss` still reached back through every unmasked
+  position — confirmed empirically at 432/512 nonzero gradient entries under a
+  realistic 15% mask ratio. The existing test only covered 100% masking, which zeroes
+  that path as a side effect of `torch.where`'s backward and never exercised the
+  realistic partial-masking case.
+
+  Fixed by detaching `input.x` once, up front
+  (`input = input._replace(x=input.x.detach())`), so both `corrupted_x` and `target`
+  are built from the same severed copy. `mlm_loss` now trains only `mask_token`,
+  `encoder`, and `head` — never whatever produced `input.x` — regardless of mask
+  ratio. Added a regression test at the realistic mask ratio
+  (`test_wrapper_severs_gradient_to_upstream_input_under_partial_masking`) alongside
+  the existing full-masking one.
+
 ## [4.3.0] — 2026-07-17
 
 ### Added

{stackformers-4.3.0 → stackformers-4.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 4.3.0
+Version: 4.3.1
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-4.3.0 → stackformers-4.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "4.3.0"
+version = "4.3.1"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-4.3.0 → stackformers-4.3.1}/stackformers/mlm/README.md

@@ -41,7 +41,7 @@ One call site covers both modes: in training, `mlm_wrapper` runs the encoder twi
 
 **The mask token is owned state, not an injected collaborator.** It's the wrapper's own learned parameter, analogous to a `nn.Linear`'s weight — not swappable behavior, so it isn't a `Protocol`.
 
-**The reconstruction target is always `input.x.detach()`, unconditionally.** Not a config flag: computing a detached target is the only behavior the masked pass ever has, so there is nothing to branch on. This severs the gradient path from the reconstruction loss back to whatever produced `input.x`, which removes the standard collapse shortcut for self-supervised regression targets (drive every token toward one constant vector, and reconstruction becomes trivial). See `tests/mlm/test_wrapper.py` for the gradient test that verifies this directly.
+**`input.x` is detached once, before it is used for anything — not only as the target.** The masked pass replaces `input` with `input._replace(x=input.x.detach())` up front, so both the unmasked context fed to the encoder and the reconstruction target come from that same detached copy. Detaching only the target is not sufficient: self-attention still mixes the live, undetached unmasked positions into the masked positions' predictions, so a gradient path back to whatever produced `input.x` would otherwise survive at every unmasked position — worse, it grows with the fraction of tokens left unmasked. Severing the whole input removes the standard collapse shortcut for self-supervised regression targets (drive every token toward one constant vector, and reconstruction becomes trivial) — not just the degenerate all-masked case, but the realistic partial-masking case a target-only detach would silently miss. See `tests/mlm/test_wrapper.py::test_wrapper_severs_gradient_to_upstream_input_under_partial_masking` for the regression test.
 
 **No layout dispatch in the wrapper.** `MaskingStrategy` returns a boolean tensor shaped like `input.x`'s leading dims — `(b, n)` for `PaddedInput`, `(nt,)` for `PackedInput`. `torch.where` and boolean advanced indexing both broadcast/gather correctly against either shape, so `MLMWrapper.forward` never matches on the sequence variant; only `RandomMasking` does, internally.
 

{stackformers-4.3.0 → stackformers-4.3.1}/stackformers/mlm/wrapper.py

@@ -38,9 +38,13 @@ class MLMWrapper(nn.Module):
     invoke this unconditionally in both modes without an if-training branch of their
     own — the same role self.training already plays in nn.Dropout or nn.BatchNorm.
 
-    The reconstruction target is always the encoder's own clean input embedding,
-    detached — so this loss trains the encoder and head, never whatever produced
-    input.x. That removes the representation-collapse shortcut a trainable tokenizer
+    The masked pass detaches input.x once, up front — both the unmasked context fed to
+    the encoder and the reconstruction target come from that same detached copy — so
+    mlm_loss trains the encoder, mask_token, and head, never whatever produced input.x.
+    Detaching only the target is not sufficient: self-attention still mixes the
+    unmasked (undetached) positions into the masked positions' predictions, so a live
+    path back to input.x would otherwise survive at every unmasked position. Severing
+    it entirely removes the representation-collapse shortcut a trainable tokenizer
     would otherwise have available (drive every token toward a constant vector to make
     reconstruction trivial).
     """
@@ -60,8 +64,13 @@ class MLMWrapper(nn.Module):
         self.head = head if head is not None else RegressionHead(config.dim)
 
     def _masked_loss(self, input: SequenceInput, encoder: EncoderLike) -> Tensor:
+        # Detached once, up front, and reused for both corrupted_x and target below —
+        # unmasked positions must not carry a live path back to whatever produced
+        # input.x either, or self-attention mixes them into the masked positions'
+        # predictions and mlm_loss reaches upstream anyway despite target being detached.
+        input = input._replace(x=input.x.detach())
+        target = input.x
         should_mask = self.masking_strategy(input)
-        target = input.x.detach()
         # should_mask's shape always matches input.x's leading dims, so this select and
         # the boolean indexing below both work for PaddedInput and PackedInput alike.
         corrupted_x = torch.where(should_mask.unsqueeze(-1), self.mask_token, input.x)

{stackformers-4.3.0 → stackformers-4.3.1}/tests/mlm/test_wrapper.py

@@ -129,16 +129,17 @@ def test_wrapper_does_not_own_encoder(device: torch.device) -> None:
     assert encoder_param_ids.isdisjoint(wrapper_param_ids)
 
 
-def test_wrapper_detaches_target_from_upstream_input(device: torch.device) -> None:
+def test_wrapper_severs_gradient_to_upstream_input_under_full_masking(
+    device: torch.device,
+) -> None:
     """Collapse guard (design doc §5): a trainable tokenizer upstream of x must get no
     gradient from this loss, or the optimizer can collapse every token to one constant
     vector to make reconstruction trivial.
 
-    Force every position masked, so torch.where's backward zeroes x's contribution
-    through the corrupted-x branch entirely (its condition is True everywhere, so the
-    "x" branch of the select is never taken). The only remaining path from loss back to
-    x would be through the target — and that's detached — so x's accumulated gradient
-    must be exactly zero.
+    With every position masked, x is detached before corrupted_x is built from it, so
+    corrupted_x is pure mask_token — independent of x entirely. x never appears in the
+    masked pass's graph at all, so x.grad must be exactly None after backward, not just
+    zero-valued.
     """
     config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
     encoder = _build_encoder(device, torch.float32)
@@ -147,8 +148,28 @@ def test_wrapper_detaches_target_from_upstream_input(device: torch.device) -> No
     mask = torch.ones(B, N, dtype=torch.bool, device=device)
     res = wrapper(make_padded_input(x, mask), encoder)
     res.mlm_loss.backward()
-    assert x.grad is not None
-    assert torch.all(x.grad == 0)
+    assert x.grad is None
+
+
+def test_wrapper_severs_gradient_to_upstream_input_under_partial_masking(
+    device: torch.device,
+) -> None:
+    """Regression test: detaching only the target is not enough to close the collapse
+    shortcut. With the default (partial) mask ratio, corrupted_x still holds the live x
+    at unmasked positions, and self-attention mixes those into the masked positions'
+    predictions — so mlm_loss could still reach x through that path unless x is
+    detached before it is used for anything, not only before it is used as the target.
+    This is the realistic training case; test_..._full_masking above is the degenerate
+    edge case that a target-only detach could still pass.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.15)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    x = torch.randn(B, N, D, device=device, requires_grad=True)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert x.grad is None
 
 
 def test_wrapper_accepts_custom_masking_strategy(device: torch.device) -> None:

{stackformers-4.3.0 → stackformers-4.3.1}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "4.3.0"
+version = "4.3.1"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },