Atila Vahedian / Cotangent
Machine learning systems · Apple Metal

Cotangent

Faster gradient handling, with the arithmetic intact.

A study of the work between backward and the next parameter update. Cotangent batches gradient norms and packs fused AdamW updates while preserving the tested native clipping behavior. The full model gradients remain in the computation.

15.03%Less time to the quality target
100 paired seeds · eager replication
23.55–34.07%Additional training-time reduction
Same compiled model in each pair
4 configurationsTwo transformer sizes, gated GQA,
and causal convolution · one M5 Pro
01 / Results

Does it improve the training step?

The original 200-run replication passed its fixed time and quality requirements. The stronger 128-run extension compares Cotangent with compiled training and isolates the two changes that produce the saving.

Beyond the compiled baseline

Twelve paired seeds per configuration, 1,200 steps per arm. Both arms use the same compiled forward/backward. The baseline uses native clipping and fused AdamW; the candidate adds Cotangent gradient handling. Intervals below are descriptive, two-sided 99% paired Student-t intervals.

ConfigurationParametersBaseline / CotangentTime reduction99% interval
Small transformer3.35M15.51 / 11.86 s23.55%22.97–24.13%
Medium transformer19.28M56.32 / 37.16 s34.07%32.87–35.26%
Gated GQA9.74M35.03 / 23.99 s31.55%30.73–32.37%
Causal convolution3.75M17.24 / 12.54 s27.30%26.95–27.64%
Four paired compiled-training comparisons and full-test quality differences with descriptive 99 percent confidence intervals
All 48 paired observations are retained, including 0 slower candidate pairs and 31 pairs with worse full-test quality. Positive time reduction is faster. Positive quality difference is worse. The dashed quality line shows the original 0.01-BPB tolerance for reference.
Full-test quality and completion costs

The descriptive quality interval extends above 0.01 BPB for Small transformer, Medium transformer, Gated GQA. The narrow tolerance is therefore not established for every configuration. The extension is a descriptive breadth study; the independently powered quality decision belongs to the original replication.

ConfigurationBaseline / Cotangent BPBDifference99% intervalTotal-time reduction
Small transformer3.0583 / 3.0620+0.00365-0.00940 to +0.0167017.89%
Medium transformer2.8992 / 2.8854-0.01382-0.07474 to +0.0471128.73%
Gated GQA3.0413 / 3.0685+0.02723-0.00819 to +0.0626425.28%
Causal convolution2.3100 / 2.3100-0.00000-0.00005 to +0.0000519.61%

BPB is bits per byte; lower is better. Total completion time includes setup, progress diagnostics, full validation/test evaluation and checkpoint writing. Compiled measurements use a shared cache already exercised by pilots. All weight files remain local.

Sampled driver allocation includes allocator/cache effects. Small transformer peaks: 219 MiB baseline / 1075 MiB Cotangent. Causal convolution peaks: 180 MiB baseline / 1068 MiB Cotangent. The complete memory and completion-time table is in the paper.

The independent eager replication

Across 100 fresh paired seeds, time to the fixed validation target fell by 15.03% (99.5% CI, 13.50–16.56%). Fixed-budget training time fell by 16.97% (16.26–17.69%).

Mean full-test change: +0.00138 BPB (99.5% CI, −0.00451 to +0.00728). Its upper bound is below the predeclared 0.01-BPB tolerance. Ninety-nine pairs reached the target faster; one was slower. Fifty pairs had worse candidate test BPB. Every observation remains in the analysis.

Explore the complete replication report →
02 / Ablation

Which change accounts for the saving?

Eight four-treatment seed blocks cross scalar or bucketed norms with native or packed AdamW. Every treatment trains the same small transformer for 800 steps. Execution order is balanced across the four cells.

Norm batching, parameter packing and combined training-time reductions from eight factorial blocks
Norm batching alone saves 15.92%, parameter packing 2.16%, and their combination 17.86%. These contrasts use the native scalar treatment as the reference. The complete main-effect and interaction estimates are in the paper and public records.
03 / Method

Batch dispatch. Preserve the reduction.

Clipping first computes one norm per parameter, then combines those scalar norms. A flat reduction has the same real-arithmetic value, but floating-point additions and kernel selection can change its numerical result.

01Complete backward
02Equal-length norm buckets
03Original joint clipping norm
04Packed native fused AdamW

A layout chosen for the backend

Each equal-length bucket is viewed as B × (K/2) × 2. Reducing both trailing axes keeps the pinned MPS generic norm path. Scalar norms return to original parameter order before the final reduction. Short and odd lengths use native scalar fallbacks.

Consistent optimizer coordinates

Parameters, gradients and moment states share the same coordinate ordering. AdamW’s coordinatewise update commutes with that reindexing under the stated assumptions. All stacking, copying, concatenation and temporary buffers count in the measurements.

Mathematical argument and implementation →

128Full-model numerical checks across four configurations and FP32/BF16 forward precision
288Norm cases across 32 lengths, three bucket sizes and three gradient dtypes
16Compiled composition checks with matched deterministic embedding derivatives

These controlled checks isolate numerical update behavior. Official training retains native embedding backward, whose accumulation can vary across independently executed runs. The implementation requires dense, contiguous, complete gradients and one optimizer hyperparameter group. Equivalence on other devices, versions, strides or execution paths needs separate verification.

04 / Evidence

Every compiled comparison, including adverse pairs.

Inspect the 48 breadth pairs directly. Timing is measured at fixed training budget; quality uses the full official test split. Positive quality difference means the candidate was worse.

ConfigurationSeedBaseline (s)Cotangent (s)Time reductionTest Δ BPB
Protocols, raw measurements and reproducibility

Frozen extension protocol · Complete extension analysis · All paired records · Reproduction instructions

The extension contains 96 breadth runs and 32 factorial runs. Tags protocol-v5 and protocol-v6 precede their official outcomes. Earlier failures stay in the experimental record. All 128 extension snapshots were restored and checked locally.

python3 analysis/v6/public_audit.py

This standard-library audit verifies public source/record hashes and independently recomputes the paired statistics. It does not execute Metal training or verify the local weight files.

05 / Paper

The research, in full.

Cotangent: Preserving Reduction Arithmetic in Batched Gradient Handling

The paper develops the coordinate argument, explains the reduction dispatch, reports the independent replication and stronger compiled comparisons, and examines the method’s numerical and hardware boundaries.

Measured scope

One M5 Pro, PyTorch 2.14.1, byte language models from 3.35M to 19.28M parameters, and a reused WikiText-2 corpus. The saving concerns gradient handling and optimizer application. Full activations and derivative products remain, and temporary buffers add memory. Compiled comparisons use warm caches and per-step synchronization. Other Apple devices, CUDA, distributed training and larger model scales remain untested.

The original replication supports its predeclared quality tolerance. The extension supplies per-configuration descriptive intervals rather than a new powered quality decision. Read the experimental record →