Documentation

LeanMlir.Proofs.Foundation.BceLossCot

BCE-with-logits: the loss, and the cotangent the RSB renders emit #

SmoothedLossCot.lean is this file's twin for the label-smoothed softmax cross-entropy every other net trains on. ResNet-50's RSB-A2/A3 recipes use binary cross-entropy with logits instead, and ResNet50RenderB's bce := true path swaps the loss cotangent for a three-op chain:

%sm = sigmoidB(logits)      %d0 = %sm − %onehot      %dy = %d0 / (B·K)

so dy = (σ(z) − t)/(B·K), against softmax-CE's five ops and /B. This file proves that chain is the gradient of a loss, which ResNet-50's BCE step tie needs.

What is proved #

What is NOT claimed #

Note: The divisor is B·K, not B, and the theorem carries it as a binder rather than asserting it. bceLossCotGraph_row_committed pins it to N·K, which is what ResNet50RenderB bakes. timm's BinaryCrossEntropy is reduction='mean' over B×C, not the mean of the per-example sum over classes; at K = 1000 the two differ by 1000× on the effective step, and RSB-A2's lr 5e-3 is tuned to this form. bceLogits is the per-example SUM over classes, so the /K half of the divisor is what turns it into the mean.

Note: No label smoothing on this path, and that is the recipe. timm's a3 arg string is ls0.0; the soft targets reach %onehot from mixup/cutmix on the host, not from a smoothing constant. So smoothTarget does not appear here and the render emits three ops where CE emits five.

Note: %loss itself is report-only. The renderer's reportBceLoss block is hand-written text, not pretty of an AST node, and nothing here is about those lines. What is proved is about the COTANGENT chain, which is on the gradient path and is pretty(provenGraph).

Note: One replica, as everywhere: under *dp* each gradient node feeds the allReduceMeanF node (DataParallel.Node), which composes this with the replica mean.

source
noncomputable def Proofs.softplus (z : ℝ) :
ℝ

Softplus, log(1 + eᶻ). The renderer computes it as max(z,0) + log(1 + e^−|z|), which is this function written so that no positive number is ever exponentiated; that rearrangement is a floating-point concern and not a different map.

Equations
Instances For
    source
    theorem Proofs.softplus_hasDerivAt (z : ℝ) :
    HasDerivAt softplus (sigmoidScalar z) z

    d/dz softplus(z) = σ(z) — the one derivative BCE-with-logits needs, and the reason the loss is written in this form at all.

    source
    theorem Proofs.softplus_neg (z : ℝ) :
    softplus (-z) = softplus z - z

    softplus(−z) = softplus(z) − z — the identity that collapses the reference's two softplus calls into one, and the bridge to log σ.

    source
    theorem Proofs.log_sigmoidScalar (z : ℝ) :
    Real.log (sigmoidScalar z) = -softplus (-z)

    log σ(z) = −softplus(−z).

    source
    theorem Proofs.one_sub_sigmoidScalar (z : ℝ) :
    1 - sigmoidScalar z = sigmoidScalar (-z)

    1 − σ(z) = σ(−z), hence log(1 − σ(z)) = −softplus(z).

    source
    noncomputable def Proofs.bceLogits (K : ℕ) (t z : Vec K) :
    ℝ

    BCE-with-logits at one example, Σ_k (softplus(z_k) − t_k·z_k) — the SUM over the K classes, in the stable form the renderer's own %loss block computes. The emitted cotangent divides by B·K, so the /K half of that divisor is what makes the shipped objective the MEAN over B×K rather than the mean of these sums (see the module header's note on the divisor).

    Equations
    Instances For
      source
      theorem Proofs.bceLogits_eq_logSigmoid (K : ℕ) (t z : Vec K) :
      bceLogits K t z = ∑ k : Fin K, -(t k * Real.log (sigmoidScalar (z k)) + (1 - t k) * Real.log (1 - sigmoidScalar (z k)))

      It IS binary cross-entropy: class by class, softplus(z) − t·z is −[t·log σ(z) + (1−t)·log(1 − σ(z))].

      This is what earns bceLogits its name. Without it, bceLogits_grad would be circular — a function defined to have the derivative the render emits, then proved to have it. The proof is softplus_neg and the two log σ identities; nothing analytic.

      source
      theorem Proofs.bceLogits_grad (K : ℕ) (t z : Vec K) (j : Fin K) :
      pdiv (fun (z' : Vec K) (x : Fin 1) => bceLogits K t z') z j 0 = sigmoidScalar (z j) - t j

      The emitted cotangent's numerator is this loss's gradient: ∂/∂z_j Σ_k (softplus(z_k) − t_k·z_k) = σ(z_j) − t_j.

      Note: NO hypothesis on t. Softmax-CE's gradient needed Σ_k t_k = 1 to collapse (Σ t)·softmax − t; BCE is per-class and separable, so the mixup targets — which are convex combinations of one-hots, and the a3 recipe's only source of soft labels — need no clause.

      source
      def Proofs.bceCotTail {N n : ℕ} (bk : ℝ) (bStr : String) (sg oh : StableHLO.SHlo (N * n)) :
      StableHLO.SHlo (N * n)

      The two ops after the sigmoid: (sg − oh) / bk. bceLossCotGraph is this applied to its sigmoid; it is computable, so the render prints its cotangent as pretty of this node over the sigmoid's name (RenderKit.bceCotB).

      Equations
      Instances For
        source
        noncomputable def Proofs.bceLossCotGraph (N K : ℕ) (bk : ℝ) (bStr logN ohN : String) (logits t : Vec (N * (1 * K))) :
        StableHLO.SHlo (N * (1 * K))

        The three-op BCE cotangent chain ResNet50RenderB emits under bce := true, at one row per example (m = 1, n = K) and batch N. logits is the head's output, t the graph input %onehot, and bk the baked divisor.

        Note: The render emits the sigmoid at (N := B, n := K) and re-enters the subtraction through an .operand at N*(1*K); the two indices are equal and NOT definitionally so at a variable K, which the render's own comment records. Stated here at the one index throughout, which is legitimate because sigmoid is elementwise and carries no row structure — smoothedLossCotGraph nests through the same .operand seam for the same reason.

        Equations
        Instances For
          source
          theorem Proofs.bceLossCotGraph_den (N K : ℕ) (bk : ℝ) (bStr logN ohN : String) (logits t : Vec (N * (1 * K))) (i : Fin (N * (1 * K))) :
          StableHLO.den (bceLossCotGraph N K bk bStr logN ohN logits t) i = (sigmoid (N * (1 * K)) logits i - t i) / bk

          What the chain denotes, coordinatewise: (σ(logits) − t) / bk.

          source
          theorem Proofs.bceLossCotGraph_row (N K : ℕ) (bk : ℝ) (bStr logN ohN : String) (logits t : Vec (N * (1 * K))) (n : Fin N) (j : Fin K) :
          StableHLO.den (bceLossCotGraph N K bk bStr logN ohN logits t) (finProdFinEquiv (n, finProdFinEquiv (0, j))) = pdiv (fun (z' : Vec K) (x : Fin 1) => bceLogits K (Mat.unflatten (StableHLO.batchSlice N (1 * K) t n) 0) z') (Mat.unflatten (StableHLO.batchSlice N (1 * K) logits n) 0) j 0 / bk

          Each row of the emitted cotangent is BCE-with-logits' gradient at that example's logits, divided by the baked constant. SmoothedLossCot's smoothedLossCotGraph_row at this loss, and with no hypothesis at all where that one needs the target's mass.

          source
          theorem Proofs.bceLossCotGraph_row_committed (N K : ℕ) (bStr logN ohN : String) (logits t : Vec (N * (1 * K))) (n : Fin N) (j : Fin K) :
          StableHLO.den (bceLossCotGraph N K (↑N * ↑K) bStr logN ohN logits t) (finProdFinEquiv (n, finProdFinEquiv (0, j))) = pdiv (fun (z' : Vec K) (x : Fin 1) => bceLogits K (Mat.unflatten (StableHLO.batchSlice N (1 * K) t n) 0) z') (Mat.unflatten (StableHLO.batchSlice N (1 * K) logits n) 0) j 0 / (↑N * ↑K)

          Note: The committed divisor is N·K, the mean over B×K. ResNet50RenderB bakes {B * nClasses}.0; softmax-CE's peer bakes {B}.0. timm's BinaryCrossEntropy is reduction='mean' over B×C, not the mean of the per-example sum over classes, and at K = 1000 the two differ by 1000× on the effective step — RSB-A2's lr 5e-3 is tuned to this form. Stated separately from bceLossCotGraph_row so that the divisor is a checked fact about the artifact rather than a binder nobody instantiated.