MobileNetV2's strided stem stage and its stem / head leaf ties #
Two pieces MobileNetV2's whole-net ties stand on (the per-example paper tie first, retired 2026-09-19):
convStridedBnRelu6PC_has_vjp_at— the STEM stage's certified VJP,relu6 ∘ bnPC ∘ flatConvStride2Xla. The repo had the strided-conv-with-relu peer (the per-example r34 stem, since retired) and the non-strided relu6 peer (convBnRelu6PC_has_vjp_at); this is the missing corner.convStridedBnRelu6PCBack_eq_vjp_backward/convBnRelu6PCBack_eq_vjp_backward— the stem and head leaf ties, both closing on one conv-leaf rewrite (flatConvStride2XlaBack_eq_vjp_backward/convFlatBack_eq_vjp_backward) and thenrfl: relu6's certified backward ISreluMaskBack (0 < · ∧ · < 6), and the pinned BN-back is definitionally the certified one.
⚠ Both are SMOOTH-POINT statements, as every HasVJPAt in this cone is: the stem's and head's
post-BN clamp windows (≠ 0 ∧ ≠ 6) are hypotheses.
noncomputable def
Proofs.convStridedBnRelu6PC_has_vjp_at
{ic oc h w kH kW : ℕ}
(W : Kernel4 oc ic kH kW)
(b : Vec oc)
(ε : ℝ)
(γ β : Vec oc)
(hε : 0 < ε)
(v : Vec (ic * (2 * h) * (2 * w)))
(h_smooth :
∀ (k : Fin (oc * h * w)),
bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 0 ∧ bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 6)
:
HasVJPAt (relu6 (oc * h * w) ∘ bnPerChannelTensor3 oc h w ε γ β ∘ flatConvStride2Xla W b) v
conv(stride-2) → per-channel-BN → relu6 VJP at a smooth point — MobileNetV2's STEM,
i.e. MobileNetV2FullVJP's convBnRelu6StridedPC_has_vjp_at.
Equations
- Proofs.convStridedBnRelu6PC_has_vjp_at W b ε γ β hε v h_smooth = Proofs.convBnRelu6StridedPC_has_vjp_at W b ε γ β hε v h_smooth
Instances For
theorem
Proofs.convStridedBnRelu6PC_differentiableAt
{ic oc h w kH kW : ℕ}
(W : Kernel4 oc ic kH kW)
(b : Vec oc)
(ε : ℝ)
(γ β : Vec oc)
(hε : 0 < ε)
(v : Vec (ic * (2 * h) * (2 * w)))
(h_smooth :
∀ (k : Fin (oc * h * w)),
bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 0 ∧ bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 6)
:
DifferentiableAt ℝ (relu6 (oc * h * w) ∘ bnPerChannelTensor3 oc h w ε γ β ∘ flatConvStride2Xla W b) v
theorem
Proofs.convStridedBnRelu6PCBack_eq_vjp_backward
{ic oc h w kH kW : ℕ}
(hkH : 2 * ((kH - 1) / 2) + 1 = kH)
(hkW : 2 * ((kW - 1) / 2) + 1 = kW)
(W : Kernel4 oc ic kH kW)
(b : Vec oc)
(ε : ℝ)
(γ β : Vec oc)
(hε : 0 < ε)
(v : Vec (ic * (2 * h) * (2 * w)))
(h_smooth :
∀ (k : Fin (oc * h * w)),
bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 0 ∧ bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) k ≠ 6)
:
(flatConvStride2XlaBack W ∘ (bnPerChannelTensor3_has_vjp oc h w ε hε γ β).backward (flatConvStride2Xla W b v) ∘ reluMaskBack fun (i : Fin (oc * h * w)) =>
0 < bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) i ∧ bnPerChannelTensor3 oc h w ε γ β (flatConvStride2Xla W b v) i < 6) = (convStridedBnRelu6PC_has_vjp_at W b ε γ β hε v h_smooth).backward
The STEM tie.
theorem
Proofs.convBnRelu6PCBack_eq_vjp_backward
{ic oc h w kH kW : ℕ}
(hkH : 2 * ((kH - 1) / 2) + 1 = kH)
(hkW : 2 * ((kW - 1) / 2) + 1 = kW)
(W : Kernel4 oc ic kH kW)
(b : Vec oc)
(ε : ℝ)
(γ β : Vec oc)
(hε : 0 < ε)
(v : Vec (ic * h * w))
(h_smooth :
∀ (k : Fin (oc * h * w)),
bnPerChannelTensor3 oc h w ε γ β (flatConv W b v) k ≠ 0 ∧ bnPerChannelTensor3 oc h w ε γ β (flatConv W b v) k ≠ 6)
:
(convFlatBack W ∘ (bnPerChannelTensor3_has_vjp oc h w ε hε γ β).backward (flatConv W b v) ∘ reluMaskBack fun (i : Fin (oc * h * w)) =>
0 < bnPerChannelTensor3 oc h w ε γ β (flatConv W b v) i ∧ bnPerChannelTensor3 oc h w ε γ β (flatConv W b v) i < 6) = (convBnRelu6PC_has_vjp_at W b ε γ β hε v h_smooth).backward
The HEAD tie.