When Small Losses and Gradients Disappear
Permanent archive: Zenodo · 10.5281/zenodo.22696183.
Zero can look reassuring while hiding a lost, representable quantity. Three compact cases in MLX and nntrainer expose two ways that floating-point algebra erases useful information.
This audit executes actual MLX and nntrainer implementations. It checks mathematical values and derivatives, rather than a chatbot’s answers. Three implementation reports belong to two failure families: binary cross-entropy tail cancellation, and an asymmetric derivative of logaddexp.
| Implementation, FP32 | Input | Original | Mathematical target |
|---|---|---|---|
| MLX binary cross entropy | logit 20, label 1 | loss 0 | ≈ 2.0611536203 × 10⁻⁹ |
| MLX logaddexp VJP and JVP | logaddexp(20,b), derivative in b at b=0 | derivative 0 | ≈ 2.0611536182 × 10⁻⁹ |
| nntrainer sigmoid cross entropy | logit 20, label 1; loss_scale=1 | loss 0; gradient 0 | loss ≈ 2.061 × 10⁻⁹; gradient ≈ −2.061 × 10⁻⁹ |
These quantities are inside the normal FP32 range. Rounding to zero is avoidable here. This experiment does not establish a measurable effect on model quality or a failure of every product from either organization.
BCE: a small positive loss disappears
The MLX logits formula is logaddexp(0,x) − x·y. For x=20 and y=1, the small correction has already been lost when it is added to 20; subtracting 20 cannot recover it. The reflected input, x=−20 and y=0, retains the correction. A label/sign symmetry check reveals the mismatch.
linear = x*(1-y) if x >= 0 else -x*y
loss = linear + log1p(exp(-abs(x)))
The local patch evaluates this identity using a sign branch that also gives the correct derivative at zero. It preserves weights, reductions, soft labels and the existing probability-input branch. Small-value tests use relative tolerance with zero absolute tolerance; references are evaluated from the actual rounded inputs with mpmath at 100 decimal digits.
Logaddexp: exchanging arguments changes the gradient
Logaddexp is symmetric in its arguments. Its derivatives should exchange with them. MLX computed the second derivative factor as 1 − sigmoid(a−b). At a−b=20, the sigmoid rounds to one, so the subtraction erases a representable derivative. Putting the differentiated value in the first argument restores it.
d/da logaddexp(a,b) = sigmoid(a-b)
d/db logaddexp(a,b) = sigmoid(b-a)
The patch computes both factors directly in the native VJP and JVP implementations. A fresh C++ build of the audited main commit reproduces the failures on CPU; all 438 checks pass after the change. The extra execution cost has not been benchmarked.
nntrainer: preserve the tail in forward and backward
The sigmoid cross-entropy layer uses log(1 + exp(-abs(x))), so adding a tiny exponential to one can lose it before the logarithm. Its subtraction of large matching linear terms is a second source of cancellation. Backward subtracts the label from an already rounded sigmoid.
The FP32 CPU patch uses log1p, a sign-dependent linear term and double intermediates. Backward computes the small sigmoid tail directly. Tests execute the actual layer through RunLayerContext, verify analytic gradients and finite differences, and preserve labels and global-mean behavior across batch shapes. The previously reported loss_scale issue is separate; these cases use a scale of one.
What passed, and what was actually built
| Suite | Original | Patched |
|---|---|---|
| MLX BCE, CPU | 23 failed / 8 passed | 31 passed |
| MLX BCE, Metal | Minimal original failures reproduced | 31 passed |
| MLX native C++ logaddexp, CPU | 72 failures / 438 checks | 0 failures / 438 checks |
| nntrainer losses, FP32 CPU | 4 failed / 73 passed | 77 passed |
- MLX: main at
ce916dbb. Python BCE source runs with official MLX/core 0.32.2 on CPU and Metal. The logaddexp C++ core was built from that source on CPU. - Native Metal limitation: the installed Command Line Tools lack
xcrun metal. The original logaddexp defect is reproduced on official Metal 0.32.2; its native patch was tested on CPU only. - nntrainer: main at
a7ea056e, real C++ FP32 CPU execution. The 77 tests include five new sigmoid tests, six prior softmax tests and 66 existing loss semantics tests. A clean reproduction of the five new tests needs only this round’s patches. - Machine and coverage: Apple M4, macOS 15.5. MLX tests cover FP32, FP16 and BF16 with suitable representability limits. Full repository CI, CUDA, mobile backends, performance and model-training outcomes were not assessed.
Duplicate review and evidence boundary
The package preserves public searches covering open and closed issues and pull requests, along with the nearest matches and reasons for excluding them. The existing MLX probability-input BCE fix concerns a different mode. nntrainer loss scaling is a known, separate issue. No exact match was identified for these cases within the recorded search; this is not a guarantee of absolute novelty.
QuantizeLinear and earlier activation reports are excluded from the new-case count. Large-input normalization and cosine candidates remain preliminary. No new Anthropic/Bloom finding is claimed in this release. Experiments, local patches and publication materials were prepared with AI assistance; the numerical conclusions come from the recorded execution and mathematical references.
Watch and discuss
Watch the 50-second YouTube Short → · Read the LinkedIn post →
The video uses synthetic English narration and original diagrams, with captions and links to the same evidence.
Download and reproduce
Download tests, logs, patches and checksums →
Open the reproducible GitHub package →
Archive SHA-256: 95e949524548b3599922354cf997b38e9fed9a841cbda5fbbbc98f0d0fc5211e
