The inverted-residual pieces every MBConv net's sync-BN step twin shares #
MobileNetV2 (MobileNetV2SyncStepTieB.lean) and EfficientNet-B0 (EfficientNetSyncStepTieG.lean)
both tie their sync-BN data-parallel step to the single-device step at R·N on
ResNet34SyncStepTieB.lean's four steps. Their blocks share the ops ResNet-34 does not have — the
depthwise conv and its symmetric-strided peer, the XLA-SAME strided stem conv, global average
pooling and the row-wise dense — and this file states each once, net-agnostically:
- §1 homogeneity — each input-VJP and weight-gradient node is linear in its cotangent;
- §2 sharding — each input-VJP is a per-example map, so it commutes with the batch cut;
- §3 the collectives — the replica mean of a weight-gradient node is
1/Rof the global node; - §4 the per-parameter DP ties — a replica family at
R ×the shards of a global cotangent gives the single-device node atR·N(*_of_scaled).
MobileNetV4's render uses the same three weight-gradient kinds. The XLA-SAME strided depthwise
is MobileNetV2's alone and stays in its twin.
theorem
Proofs.MBConvSyncTieB.dStridedInB_smul
(N : ℕ)
{c h w kH kW : ℕ}
(W : DepthwiseKernel c kH kW)
(b : Vec c)
(dy : Vec (N * (c * h * w)))
(s : ℝ)
:
theorem
Proofs.MBConvSyncTieB.rowDenseBackFlat_smul
(N a c : ℕ)
(W : Mat a c)
(dy : Vec (N * c))
(s : ℝ)
:
(StableHLO.rowDenseBackFlat N a c W fun (i : Fin (N * c)) => s * dy i) = fun (i : Fin (N * a)) =>
s * StableHLO.rowDenseBackFlat N a c W dy i
The row-wise input-VJP dX = W·dy (the classifier's, and the SE excite dense's) is linear in
dy.
theorem
Proofs.MBConvSyncTieB.hasVJP3_backward_smul
{c₁ h₁ w₁ c₂ h₂ w₂ : ℕ}
{f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂}
(hf : HasVJP3 f)
(x : Tensor3 c₁ h₁ w₁)
(a : ℝ)
(dy : Tensor3 c₂ h₂ w₂)
:
A HasVJP3 backward is linear in its cotangent — HasVJP.backward_smul's three-axis peer,
read off HasVJP3.correct. The stride-1 depthwise weight gradient is stated through one.
theorem
Proofs.MBConvSyncTieB.depthwiseWeightGradB_smul
{N c h w kH kW : ℕ}
(xN cotN : String)
(b : Vec c)
(x : Vec (N * (c * h * w)))
(W : DepthwiseKernel c kH kW)
(cot : Vec (N * (c * h * w)))
(s : ℝ)
(idx : Fin (c * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.depthwiseWeightGradB xN b x W
(StableHLO.SHlo.operand cotN fun (i : Fin (N * (c * h * w))) => s * cot i))
idx = s * StableHLO.den (StableHLO.SHlo.depthwiseWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx
theorem
Proofs.MBConvSyncTieB.depthwiseStridedWeightGradB_smul
{N c h w kH kW : ℕ}
(xN cotN : String)
(b : Vec c)
(x : Vec (N * (c * (2 * h) * (2 * w))))
(W : DepthwiseKernel c kH kW)
(cot : Vec (N * (c * h * w)))
(s : ℝ)
(idx : Fin (c * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.depthwiseStridedWeightGradB xN b x W
(StableHLO.SHlo.operand cotN fun (i : Fin (N * (c * h * w))) => s * cot i))
idx = s * StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx
theorem
Proofs.MBConvSyncTieB.convStridedXlaWeightGradB_smul
{N ic oc h w kH kW : ℕ}
(xN cotN : String)
(b : Vec oc)
(x : Vec (N * (ic * (2 * h) * (2 * w))))
(W : Kernel4 oc ic kH kW)
(cot : Vec (N * (oc * h * w)))
(s : ℝ)
(idx : Fin (oc * ic * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.convStridedXlaWeightGradB xN b x W
(StableHLO.SHlo.operand cotN fun (i : Fin (N * (oc * h * w))) => s * cot i))
idx = s * StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradB xN b x W (StableHLO.SHlo.operand cotN cot)) idx
theorem
Proofs.MBConvSyncTieB.dInB_shard
{R N c h w kH kW : ℕ}
(W : DepthwiseKernel c kH kW)
(b : Vec c)
(DY : Vec (R * N * (c * h * w)))
(r : Fin R)
:
EnetTiePoC.dInB N W b (batchShard R N (c * h * w) DY r) = batchShard R N (c * h * w) (EnetTiePoC.dInB (R * N) W b DY) r
theorem
Proofs.MBConvSyncTieB.dStridedInB_shard
{R N c h w kH kW : ℕ}
(W : DepthwiseKernel c kH kW)
(b : Vec c)
(DY : Vec (R * N * (c * h * w)))
(r : Fin R)
:
EnetTiePoC.dStridedInB N W b (batchShard R N (c * h * w) DY r) = batchShard R N (c * (2 * h) * (2 * w)) (EnetTiePoC.dStridedInB (R * N) W b DY) r
theorem
Proofs.MBConvSyncTieB.gapInB_shard
{R N : ℕ}
(c h w : ℕ)
(DY : Vec (R * N * c))
(r : Fin R)
:
EnetTiePoC.gapInB N c h w (batchShard R N c DY r) = batchShard R N (c * h * w) (EnetTiePoC.gapInB (R * N) c h w DY) r
theorem
Proofs.MBConvSyncTieB.rowDenseBackFlat_shard
{R N : ℕ}
(a c : ℕ)
(W : Mat a c)
(DY : Vec (R * N * c))
(r : Fin R)
:
StableHLO.rowDenseBackFlat N a c W (batchShard R N c DY r) = batchShard R N a (StableHLO.rowDenseBackFlat (R * N) a c W DY) r
The row-wise input-VJP is batchMap of W·, so it shards like every per-example lift.
theorem
Proofs.MBConvSyncTieB.den_allReduceMeanF_depthwiseWeightGradB_shard
{N c h w kH kW : ℕ}
(R : ℕ)
(hR : 0 < R)
(t xN cotN : String)
(ds : List ℕ)
(b : Vec c)
(W : DepthwiseKernel c kH kW)
(X DY : Vec (R * N * (c * h * w)))
(dy : Fin R → StableHLO.SHlo (N * (c * h * w)))
(hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (c * h * w) DY r)
(idx : Fin (c * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) =>
StableHLO.SHlo.depthwiseWeightGradB xN b (batchShard R N (c * h * w) X r) W (dy r))
idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.depthwiseWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx
P4 at the depthwise weight — each replica's Σ_n over its own examples, averaged, is 1/R
of the global batch's Σ_n.
theorem
Proofs.MBConvSyncTieB.den_allReduceMeanF_depthwiseStridedWeightGradB_shard
{N c h w kH kW : ℕ}
(R : ℕ)
(hR : 0 < R)
(t xN cotN : String)
(ds : List ℕ)
(b : Vec c)
(W : DepthwiseKernel c kH kW)
(X : Vec (R * N * (c * (2 * h) * (2 * w))))
(DY : Vec (R * N * (c * h * w)))
(dy : Fin R → StableHLO.SHlo (N * (c * h * w)))
(hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (c * h * w) DY r)
(idx : Fin (c * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) =>
StableHLO.SHlo.depthwiseStridedWeightGradB xN b (batchShard R N (c * (2 * h) * (2 * w)) X r) W (dy r))
idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx
P4 at the strided depthwise weight.
theorem
Proofs.MBConvSyncTieB.den_allReduceMeanF_convStridedXlaWeightGradB_shard
{N ic oc h w kH kW : ℕ}
(R : ℕ)
(hR : 0 < R)
(t xN cotN : String)
(ds : List ℕ)
(b : Vec oc)
(W : Kernel4 oc ic kH kW)
(X : Vec (R * N * (ic * (2 * h) * (2 * w))))
(DY : Vec (R * N * (oc * h * w)))
(dy : Fin R → StableHLO.SHlo (N * (oc * h * w)))
(hdy : ∀ (r : Fin R), StableHLO.den (dy r) = batchShard R N (oc * h * w) DY r)
(idx : Fin (oc * ic * kH * kW))
:
StableHLO.den
(StableHLO.SHlo.allReduceMeanF R hR t ds fun (r : Fin R) =>
StableHLO.SHlo.convStridedXlaWeightGradB xN b (batchShard R N (ic * (2 * h) * (2 * w)) X r) W (dy r))
idx = 1 / ↑R * StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradB xN b X W (StableHLO.SHlo.operand cotN DY)) idx
P4 at the XLA-SAME strided conv weight (the stem) — den_allReduceMeanF_convWeightGradB_shard's
peer. Only the certificate differs from the symmetric strided one; the batch split is the same.
def
Proofs.MBConvSyncTieB.DepthwiseWSync
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{c kH kW : ℕ}
(t xN cotN : String)
(b : Vec c)
(X : Vec (R * N * (c * h * w)))
(W : DepthwiseKernel c kH kW)
(cots : Fin R → Vec (N * (c * h * w)))
(COT : Vec (R * N * (c * h * w)))
:
A depthwise weight, DP-tied — the collective over the [c, 1, kH, kW] kernel the render
all-reduces, against the single-device node at the global batch.
Equations
- One or more equations did not get rendered due to their size.
Instances For
def
Proofs.MBConvSyncTieB.DepthwiseStridedWSync
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{c kH kW : ℕ}
(t xN cotN : String)
(b : Vec c)
(X : Vec (R * N * (c * (2 * h) * (2 * w))))
(W : DepthwiseKernel c kH kW)
(cots : Fin R → Vec (N * (c * h * w)))
(COT : Vec (R * N * (c * h * w)))
:
The strided depthwise weight, DP-tied.
Equations
- One or more equations did not get rendered due to their size.
Instances For
def
Proofs.MBConvSyncTieB.ConvStridedXlaWSync
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{ic oc kH kW : ℕ}
(t xN cotN : String)
(b : Vec oc)
(X : Vec (R * N * (ic * (2 * h) * (2 * w))))
(W : Kernel4 oc ic kH kW)
(cots : Fin R → Vec (N * (oc * h * w)))
(COT : Vec (R * N * (oc * h * w)))
:
The stem's XLA-SAME strided conv weight, DP-tied. Tags are the render's: the collective is
named for the parameter, over its [oc, ic, kH, kW] shape.
Equations
- One or more equations did not get rendered due to their size.
Instances For
theorem
Proofs.MBConvSyncTieB.depthwiseWSync_of_scaled
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{c kH kW : ℕ}
(t xN cotN : String)
(b : Vec c)
(X : Vec (R * N * (c * h * w)))
(W : DepthwiseKernel c kH kW)
(cots : Fin R → Vec (N * (c * h * w)))
(COT : Vec (R * N * (c * h * w)))
(hc : ∀ (r : Fin R), cots r = batchShard R N (c * h * w) (fun (i : Fin (R * N * (c * h * w))) => ↑R * COT i) r)
:
DepthwiseWSync R hR N h w t xN cotN b X W cots COT
theorem
Proofs.MBConvSyncTieB.depthwiseStridedWSync_of_scaled
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{c kH kW : ℕ}
(t xN cotN : String)
(b : Vec c)
(X : Vec (R * N * (c * (2 * h) * (2 * w))))
(W : DepthwiseKernel c kH kW)
(cots : Fin R → Vec (N * (c * h * w)))
(COT : Vec (R * N * (c * h * w)))
(hc : ∀ (r : Fin R), cots r = batchShard R N (c * h * w) (fun (i : Fin (R * N * (c * h * w))) => ↑R * COT i) r)
:
DepthwiseStridedWSync R hR N h w t xN cotN b X W cots COT
theorem
Proofs.MBConvSyncTieB.convStridedXlaWSync_of_scaled
(R : ℕ)
(hR : 0 < R)
(N h w : ℕ)
{ic oc kH kW : ℕ}
(t xN cotN : String)
(b : Vec oc)
(X : Vec (R * N * (ic * (2 * h) * (2 * w))))
(W : Kernel4 oc ic kH kW)
(cots : Fin R → Vec (N * (oc * h * w)))
(COT : Vec (R * N * (oc * h * w)))
(hc : ∀ (r : Fin R), cots r = batchShard R N (oc * h * w) (fun (i : Fin (R * N * (oc * h * w))) => ↑R * COT i) r)
:
ConvStridedXlaWSync R hR N h w t xN cotN b X W cots COT