learning-munsell / docs /learning_munsell.md

Initial commit.

3c7db92 2 months ago

30.3 kB

Learning Munsell

Technical documentation covering performance benchmarks, training methodology, architecture design, and experimental findings.

Overview

This project implements ML models for bidirectional conversion between CIE xyY colorspace values and Munsell specifications:

xyY to Munsell (from_xyY): 33 models, best Delta-E 0.51
Munsell to xyY (to_xyY): 6 models, best Delta-E 0.48

Delta-E Interpretation

< 1.0: Not perceptible by human eye
1-2: Perceptible through close observation
2-10: Perceptible at a glance
> 10: Colors are perceived as completely different

Our best models achieve Delta-E 0.48-0.51, meaning the difference between ML prediction and iterative algorithm is not perceptible by the human eye.

xyY to Munsell (from_xyY)

Performance Benchmarks

33 models compared using all 2,734 REAL Munsell colors:

Model	Delta-E	Speed (ms)
Colour Library (Baseline)	0.00	116.3
Multi-MLP + Multi-Error Predictor	0.51	0.061
Multi-ResNet + Multi-Error Predictor (Large Dataset)	0.52	0.096
Transformer + Error Predictor (Large Dataset)	0.52	0.163
Multi-Head + Multi-Error Predictor (Large Dataset)	0.53	0.046
Multi-MLP (Classification Code) + Code-Aware Multi-Error Predictor (Large Dataset)	0.53	0.118
Multi-MLP (Classification Code) + Code-Aware Multi-Error Predictor	0.53	0.052
Multi-MLP (Classification Code) + Multi-Error Predictor	0.53	0.050
MLP + Error Predictor	0.53	0.036
Multi-Head + Multi-Error Predictor	0.54	0.057
Multi-ResNet (Large Dataset)	0.56	0.047
Multi-MLP + Multi-Error Predictor (Large Dataset)	0.56	0.060
Multi-Head + Cross-Attention Error Predictor	0.57	0.032
Multi-Head (Large Dataset)	0.57	0.012
FT-Transformer	0.70	0.067
Unified MLP	0.71	0.074
Transformer (Large Dataset)	0.73	0.120
Mixture of Experts	0.74	0.021
Multi-MLP	0.91	0.027
Multi-MLP (Classification Code)	0.92	0.026
Multi-MLP (Classification Code) (Large Dataset)	0.89	0.066
MLP + Self-Attention	0.88	0.185
Deep + Wide	1.18	0.078
MLP (Base Only)	1.30	0.009
Multi-MLP (Hue Angle sin/cos)	1.78	0.021

Note: The Colour library baseline had 171 convergence failures out of 2,734 samples (6.3% failure rate). See the full HTML report for all 33 models.

Best Models:

Best Accuracy: Multi-MLP + Multi-Error Predictor — Delta-E 0.51, 1,905x faster
Fastest: MLP Base Only (0.009 ms/sample) — 13,365x faster than Colour library
Best Balance: MLP + Error Predictor — 3,196x faster with Delta-E 0.53

Model Architectures

33 models were systematically evaluated:

Single-Stage Models

MLP (Base Only) - Simple MLP network, 3 inputs to 4 outputs
Unified MLP - Single large MLP with shared features
Multi-Head - Shared encoder with 4 independent decoder heads
Multi-Head (Large Dataset) - Multi-Head trained on 1.4M samples
Multi-MLP - 4 completely independent MLP branches (one per output)
Multi-MLP (Large Dataset) - Multi-MLP trained on 1.4M samples
MLP + Self-Attention - MLP with attention mechanism for feature weighting
Deep + Wide - Combined deep and wide network paths
Mixture of Experts - Gating network selecting specialized expert networks
Transformer (Large Dataset) - Feature Tokenizer Transformer for tabular data
FT-Transformer - Feature Tokenizer Transformer (standard size)

Two-Stage Models

MLP + Error Predictor - Base MLP with unified error correction
Multi-Head + Multi-Error Predictor - Multi-Head with 4 independent error predictors
Multi-Head + Multi-Error Predictor (Large Dataset) - Large dataset variant
Multi-MLP + Multi-Error Predictor - 4 independent branches with 4 independent error predictors
Multi-MLP + Multi-Error Predictor (Large Dataset) - Large dataset variant
Multi-ResNet + Multi-Error Predictor (Large Dataset) - Deep ResNet-style branches
Transformer + Error Predictor (Large Dataset) - Transformer base with error correction
Multi-Head + Cross-Attention Error Predictor - Cross-attention mechanism for error correction
Multi-Head + 3-Stage Error Predictor - Three cascaded error correction stages

Alternative Code Prediction Models

Multi-MLP (Classification Code) - 3 regression branches + 1 classification branch (10 logits for hue code)
Multi-MLP (Classification Code) + Multi-Error Predictor - Classification Code base with 3-branch error predictor (hue, value, chroma)
Multi-MLP (Classification Code) + Code-Aware Multi-Error Predictor - Classification Code base with code-aware 3-branch error predictor (input: xyY + regression + code one-hot)
Multi-MLP (Classification Code) (Large Dataset) - Classification Code trained on 1.4M samples
Multi-MLP (Classification Code) + Code-Aware Multi-Error Predictor (Large Dataset) - Large dataset variant of code-aware error predictor
Multi-MLP (Hue Angle sin/cos) - Encodes full Munsell angle as sin/cos pair, eliminates separate code branch The Multi-MLP + Multi-Error Predictor architecture achieved the best results with Delta-E 0.51.

Training Methodology

Data Generation

Dense xyY Grid (~500K samples)
- Regular grid in valid xyY space (MacAdam limits for Illuminant C)
- Captures general input distribution
Boundary Refinement (~700K samples)
- Adaptive dense sampling near Munsell gamut boundaries
- Uses maximum_chroma_from_renotation to detect edges
- Focuses on regions where iterative algorithm is most complex
- Includes Y/GY/G hue regions with high value/chroma (challenging areas)
Forward Augmentation (~200K samples)
- Dense Munsell space sampling via munsell_specification_to_xyY
- Ensures coverage of known valid colors

Total: ~1.4M samples for large dataset training.

Loss Functions

Two loss function approaches were tested:

Precision-Focused Loss (Default):

total_loss = 1.0 * MSE + 0.5 * MAE + 0.3 * log_penalty + 0.5 * huber_loss

MSE: Standard mean squared error
MAE: Mean absolute error
Log penalty: Heavily penalizes small errors (pushes toward high precision)
Huber loss: Small delta (0.01) for precision on small errors

Pure MSE Loss (Optimized config):

total_loss = MSE

Interestingly, the precision-focused loss achieved better Delta-E despite higher validation MSE, suggesting the custom weighting better correlates with perceptual accuracy.

Design Rationale

Two-Stage Architecture

The error predictor stage corrects systematic biases in the base model:

Base model learns the general xyY to Munsell mapping
Error predictor learns residual corrections specific to each component
Combined prediction: final = base_prediction + error_correction

This decomposition allows each stage to specialize and reduces the complexity each network must learn.

Independent Branch Design

Munsell components have different characteristics:

Hue: Circular (0-10, wrapping), most complex
Value: Linear (0-10), easiest to predict
Chroma: Highly variable range depending on hue/value
Code: Discrete hue sector (0-9)

Shared encoders force compromises between these different prediction tasks. Independent branches allow full specialization.

Architecture Details

MLP (Base Only)

Simple feedforward network predicting all 4 outputs simultaneously:

Input (3) ──► Linear Layers ──► Output (4: hue, value, chroma, code)

Smallest model (~2.3 MB ONNX)
Fastest inference (0.009 ms)
Baseline for comparison

Unified MLP

Single large MLP with shared internal features:

Input (3) ──► 128 ──► 256 ──► 512 ──► 256 ──► 128 ──► Output (4)

Shared representations across all outputs
Moderate size, good speed

Multi-Head MLP

Shared encoder with specialized decoder heads:

Input (3) ──► SHARED ENCODER (3→128→256→512) ──┬──► Hue Head (512→256→128→1)
                                               ├──► Value Head (512→256→128→1)
                                               ├──► Chroma Head (512→384→256→128→1)
                                               └──► Code Head (512→256→128→1)

Shared encoder learns common color space features
4 specialized decoder heads branch from shared representation
Parameter efficient (encoder weights shared)
Fast inference (encoder computed once)

Multi-MLP

Fully independent branches with no weight sharing:

Input (3) ──► Hue Branch    (3→128→256→512→256→128→1)
Input (3) ──► Value Branch  (3→128→256→512→256→128→1)
Input (3) ──► Chroma Branch (3→256→512→1024→512→256→1)  [2x wider]
Input (3) ──► Code Branch   (3→128→256→512→256→128→1)

4 completely independent MLPs
Each branch learns its own features from scratch
Chroma branch is wider (2x) to handle its complexity
Better accuracy than Multi-Head on large dataset (Delta-E 0.52 vs 0.56 with error predictors)

Multi-ResNet

Deep branches with residual-style connections:

Input (3) ──► Hue Branch    (3→256→512→512→512→256→1)    [6 layers]
Input (3) ──► Value Branch  (3→256→512→512→512→256→1)    [6 layers]
Input (3) ──► Chroma Branch (3→512→1024→1024→1024→512→1) [6 layers, 2x wider]
Input (3) ──► Code Branch   (3→256→512→512→512→256→1)    [6 layers]

Deeper architecture than Multi-MLP
BatchNorm + SiLU activation
Strong accuracy when combined with error predictor (Delta-E 0.52)
Largest model (~14MB base, ~28MB with error predictor)

Deep + Wide

Combined deep and wide network paths:

Input (3) ──┬──► Deep Path (multiple layers) ──┬──► Concat ──► Output (4)
            └──► Wide Path (direct connection) ─┘

Deep path captures complex patterns
Wide path preserves direct input information
Good for mixed linear/nonlinear relationships

MLP + Self-Attention

MLP with attention mechanism for feature weighting:

Input (3) ──► MLP ──► Self-Attention ──► Output (4)

Attention weights learn feature importance
Slower due to attention computation (0.173 ms)
Did not improve over simpler MLPs

Mixture of Experts

Gating network selecting specialized expert networks:

Input (3) ──► Gating Network ──► Weighted sum of Expert outputs ──► Output (4)

Multiple expert networks specialize in different input regions
Gating network learns which expert to use
More complex but did not outperform Multi-MLP

FT-Transformer

Feature Tokenizer Transformer for tabular data:

Input (3) ──► Feature Tokenizer ──► Transformer Blocks ──► Output (4)

Each input feature tokenized separately
Self-attention across feature tokens
Good for tabular data with feature interactions
Slower inference due to attention computation

Error Predictor (Two-Stage)

Second-stage network that corrects base model errors:

Stage 1: Input (3) ──► Base Model ──► Base Prediction (4)
Stage 2: [Input (3), Base Prediction (4)] ──► Error Predictor ──► Error Correction (4)
Final:   Base Prediction + Error Correction = Final Output

Learns residual corrections for each component
Can have unified (1 network) or multi (4 networks) error predictors
Consistently improves accuracy across all base architectures
Best results: Multi-MLP + Multi-Error Predictor (Delta-E 0.51)

Loss-Metric Mismatch

An important finding: optimizing MSE does not optimize Delta-E.

The Optuna hyperparameter search minimized validation MSE, but the best MSE configuration did not achieve the best Delta-E. This is because:

MSE treats all component errors equally
Delta-E (CIE2000) weights errors based on human perception
The precision-focused loss with custom weights better approximates perceptual importance

Weighted Boundary Loss (Experimental)

Analysis of model errors revealed systematic underperformance on Y/GY/G hues (Yellow/Green-Yellow/Green) with high value and chroma. The weighted boundary loss approach was explored to address this by:

Applying 3x loss weight to samples in challenging regions:
- Hue: 0.18-0.35 (normalized range covering Y/YG/G)
- Value > 0.7 (high brightness)
- Chroma > 0.5 (high saturation)
Adding boundary penalty to prevent predictions exceeding Munsell gamut limits

Finding: The large dataset approach (~1.4M samples with dense boundary sampling) naturally provides sufficient coverage of these challenging regions. Both the weighted boundary loss model and the standard dataset models achieve similar results, making explicit loss weighting optional. The best overall model is Multi-MLP + Multi-Error Predictor with Delta-E 0.51.

Hue Code Boundary Analysis

The best model (Multi-MLP + Multi-Error Predictor, Delta-E 0.51) treats the hue code as a continuous regression target. After rounding to the nearest integer, 13.1% of predictions (335/2,563 real Munsell colours) have incorrect hue codes — mostly off-by-one at hue family boundaries (e.g., model predicts code 9.198, rounds to 9 instead of correct 10).

While the forward direction (Munsell to xyY) is essentially perfect, these code errors produce large visible colour shifts in round-trip (xyY to Munsell to xyY) comparisons because adjacent hue codes represent different hue families (e.g., P vs PB).

Two alternative architectures were designed and evaluated to address this:

Approach A: Classification Head for Code (MultiMLPClassCodeToMunsell)

3 regression branches (hue, value, chroma) identical to standard Multi-MLP
1 classification branch outputting 10 logits (one per hue code 1-10)
Loss: weighted MSE on regression targets + cross-entropy on code class
Inference: code = argmax(logits) + 1
Rationale: Discrete classification avoids the rounding boundary problem entirely

Architecture:

Input (3) ──► Hue Branch    (3→128→256→512→256→128→1)
Input (3) ──► Value Branch  (3→128→256→512→256→128→1)
Input (3) ──► Chroma Branch (3→256→512→1024→512→256→1)  [2x wider]
Input (3) ──► Code Branch   (3→128→256→512→256→128→10)  [10 logits]

Approach B: Hue Angle sin/cos Encoding (MultiMLPHueAngleToMunsell)

Encodes the full Munsell angle (hue + (code - 1) * 10) as sin/cos pair
Eliminates separate hue and code branches entirely
Loss: MSE on [sin, cos, value_norm, chroma_norm]
Inference: angle = atan2(sin, cos) * 100 / (2*pi), then decompose to hue + code
Rationale: Continuous circular encoding naturally handles wrap-around

Architecture:

Input (3) ──► Hue Angle Branch (3→128→256→512→256→128→2)  [sin, cos]
Input (3) ──► Value Branch     (3→128→256→512→256→128→1)
Input (3) ──► Chroma Branch    (3→256→512→1024→512→256→1)  [2x wider]

Results

Model	Delta-E	Code Accuracy	Code MAE	Speed (ms)
Multi-MLP + Multi-Error Predictor	0.51	86.9%	-	0.058
Classification Code + Code-Aware Multi-Error Predictor	0.53	100.0%	0.0004	0.052
Classification Code + Multi-Error Predictor	0.53	100.0%	0.0004	0.050
Multi-MLP (Classification Code)	0.92	100.0%	0.0004	0.026
Multi-MLP (Hue Angle sin/cos)	1.78	96.3%	-	0.021

Key Findings:

Classification Code + Code-Aware Multi-Error Predictor achieves the best of both worlds: perfect code accuracy (100%) with Delta-E 0.53 — competitive with the overall best model (0.51). The code-aware 3-branch error predictor receives the one-hot encoded classified hue code (10 dims) alongside xyY input and regression predictions (input_dim=16), allowing each branch to learn hue-family-specific corrections. This marginally outperforms the standard multi-error predictor (Delta-E 0.5255 vs 0.5278).
Classification Code (base only) achieves perfect code accuracy but Delta-E 0.92 without error correction.
Hue Angle sin/cos improves code accuracy to 96.3% (vs 86.9% baseline) but has the worst Delta-E (1.78).

Experimental Findings

The following experiments were conducted but did not improve results:

Circular Hue Encoding

Two approaches to encoding hue as sin/cos were tested, motivated by systematic hue regression errors at hue family boundaries (e.g., 10B, 10RP, 10R) where the scalar hue regression suffers from the 0/10 discontinuity:

Swatch	Centore	ONNX (Large)	Max RGB diff
10B 4/6	1.1PB 4.0/5.9	0.3PB 4.0/5.9	0.157
10RP 4/6	0.8R 4.0/5.8	2.7RP 4.1/5.7	0.098
10R 9/6	0.7YR 9.1/6.4	6.5YR 9.1/6.5	0.081
10RP 3/6	0.6R 3.0/5.8	4.8RP 3.1/5.8	0.072

Approach A — Full Angle sin/cos (MultiMLPHueAngleToMunsell): Encodes the full Munsell angle (hue + (code - 1) * 10, range 0-100) as a sin/cos pair, eliminating the separate code branch entirely.

Approach B — Within-Family sin/cos + Classification Code (MultiMLPClassCodeCircularToMunsell): Retains the perfect classification code head while encoding only the within-family hue (0-10) as sin/cos, hypothesizing that the boundary discontinuity could be smoothed without losing code accuracy.

Results:

Model	Delta-E	Hue MAE	Code Acc.
Classification Code + Code-Aware Error Pred (Large)	0.53	0.045	100.0%
Full Angle sin/cos (Approach A)	1.78	-	96.3%
Within-Family sin/cos + Code (Approach B, Large)	1.71	0.817	100.0%
Approach B + Code-Aware Error Predictor (Large)	1.88	1.140	100.0%

Both circular approaches performed significantly worse than scalar hue regression (Delta-E 1.71-1.88 vs 0.53).

Key Takeaway: Hue within a Munsell family is not circular. Within a single family, hue 0 and hue 10 are genuinely different points (opposite edges of the family). The sin/cos encoding wraps them as if they were the same, confusing the network and degrading hue prediction by ~18x (MAE 0.817 vs 0.045). The circularity only exists across the full 0-100 Munsell angle (where 100 = 0), but Approach A showed that encoding the full angle is also ineffective — the atan2 recovery introduces its own ambiguities across the 10 hue families. The boundary errors in the scalar regression model (max RGB diff 0.157 on 4 out of 2,734 swatches) are a better trade-off than the global accuracy degradation from circular encoding.

Delta-E Training

Training with differentiable Delta-E CIE2000 loss via round-trip through the Munsell-to-xyY approximator.

Hypothesis: Perceptual Delta-E loss might outperform MSE-trained models.

Implementation: JAX/Flax model with combined MSE + Delta-E loss. Requires lower learning rate (1e-4 vs 3e-4) for stability; higher rates cause NaN gradients.

Results: While Delta-E is comparable, hue accuracy is ~10x worse:

Metric (Normalized MAE)	Delta-E Model	MSE Model
Hue MAE	0.30	0.03
Value MAE	0.002	0.004
Chroma MAE	0.007	0.008
Code MAE	0.07	0.01
Delta-E (perceptual)	0.52	0.50

Key Takeaway: Perceptual similarity != specification accuracy. The MSE model's slightly better Delta-E (0.50 vs 0.52) comes at the cost of ~10x worse hue accuracy, making it unsuitable for specification prediction. Delta-E is too permissive for hue, allowing the model to find "shortcuts" that minimize perceptual difference without correctly predicting the Munsell specification.

Classical Interpolation

Classical interpolation methods were tested on 4,995 reference Munsell colors (80% train / 20% test split). ML evaluated on 2,734 REAL Munsell colors.

Results (Validation MAE):

Component	RBF	KD-Tree	Delaunay	ML (Best)
Hue	1.40	1.40	1.29	0.03
Value	0.01	0.10	0.02	0.05
Chroma	0.22	0.99	0.35	0.11
Code	0.33	0.28	0.28	0.00

Key Insight: The reference dataset (4,995 colors) is too sparse for 3D xyY interpolation. Classical methods fail on hue prediction (MAE ~1.3-1.4), while ML achieves 47x better hue accuracy and 2-3x better chroma/code accuracy.

Circular Hue Loss

Circular distance metrics for hue prediction, accounting for cyclic nature (0-10 wraps).

Results: The circular loss model performed 21x worse on hue MAE (5.14 vs 0.24). Combined with the circular encoding findings above, this provides strong evidence that within-family Munsell hue should be treated as a bounded linear quantity, not a circular one.

Key Takeaway: Mathematical correctness != training effectiveness. The circular distance creates gradient discontinuities that harm optimization.

REAL-Only Refinement

Fine-tuning using only REAL Munsell colors (2,734) instead of ALL colors (4,995).

Results: Essentially identical performance (Delta-E 1.5233 vs 1.5191).

Key Takeaway: Data quality is not the bottleneck. Both REAL and extrapolated colors are sufficiently accurate.

Gamma Normalization

Gamma correction to the Y (luminance) channel during normalization.

Results: No consistent improvement across gamma values 1.0-3.0:

Gamma	Median ΔE (± std)
1.0 (baseline)	0.730 ± 0.054
2.5 (best)	0.683 ± 0.132

Key Takeaway: Gamma normalization does not provide consistent improvement. Standard deviations overlap - differences are within noise.

Uniform Random Sampling

The default training data generator perturbs around the 4,995 base colours from MUNSELL_COLOURS_ALL (hue prefixes 2.5, 5, 7.5, 10), creating islands of coverage with gaps between prefixes and zero samples below hue 1.0. A uniform random sampling approach was tested to fill these gaps.

Implementation: Sample hue uniformly in [0, 10], value in [1, 9], chroma in [0, 50], and code uniformly from {1, ..., 10}. Invalid specifications are discarded. 2M valid samples generated.

Results:

Model	Delta-E	Code Acc.	Hue MAE
Class. Code + Error Pred (perturbation, 500K)	0.53	100.0%	0.048
Class. Code + Error Pred (uniform, 2M)	2.23	82.2%	1.186

Key Takeaway: Uniform random sampling produces significantly worse models (Delta-E 2.23 vs 0.53, code accuracy 82.2% vs 100%). The perturbation-based approach concentrates samples around actual Munsell colours where the interpolation function has structure, while uniform sampling wastes density on regions of the space where many hue/value/chroma combinations are out of gamut and the accepted samples are too sparse to learn the mapping accurately. The coverage gaps at hue boundaries in the perturbation-based data are a minor visual artefact (affecting ~4 out of 2,734 swatches) that does not justify the dramatic accuracy loss from uniform sampling.

Munsell to xyY (to_xyY)

Performance Benchmarks

6 models compared using all 2,734 REAL Munsell colors:

Model	Delta-E	Speed (ms)
Colour Library (Baseline)	0.00	1.40
Multi-MLP + Multi-Error Predictor	0.48	0.066
Multi-Head + Multi-Error Predictor	0.51	0.054
Simple MLP	0.52	0.002
Multi-MLP	0.57	0.028
Multi-Head	0.60	0.013
Multi-MLP + Error Predictor	0.61	0.060

Best Models:

Best Accuracy: Multi-MLP + Multi-Error Predictor — Delta-E 0.48, 21x faster
Fastest: Simple MLP (0.002 ms/sample) — 669x faster than Colour library

Model Architectures

6 models were evaluated for the Munsell to xyY direction:

Single-Stage Models

Simple MLP - Basic MLP network, 4 inputs to 3 outputs
Multi-Head - Shared encoder with 3 independent decoder heads (x, y, Y)
Multi-MLP - 3 completely independent MLP branches

Two-Stage Models

Multi-MLP + Error Predictor - Base Multi-MLP with unified error correction
Multi-MLP + Multi-Error Predictor - 3 independent error predictors (BEST)
Multi-Head + Multi-Error Predictor - Multi-Head with error correction

The Multi-MLP + Multi-Error Predictor architecture achieved the best results with Delta-E 0.48.

Differentiable Approximator

A small MLP (68K parameters) trained to approximate the Munsell to xyY conversion for use in differentiable Delta-E loss:

Architecture: 4 -> 128 -> 256 -> 128 -> 3 with LayerNorm + SiLU
Accuracy: MAE ~0.0006 for x, y, and Y components
Output formats: PyTorch (.pth), ONNX, and JAX-compatible weights (.npz)

This enables differentiable Munsell to xyY conversion, which was previously only possible through non-differentiable lookup tables.

Shared Infrastructure

Hyperparameter Optimization

Optuna was used for systematic hyperparameter search over:

Learning rate (1e-4 to 1e-3)
Batch size (256, 512, 1024)
Dropout rate (0.0 to 0.2)
Chroma branch width multiplier (1.0 to 2.0)
Loss function weights (MSE, Huber)

Key finding: No dropout (0.0) consistently performed better across all models in both conversion directions, contrary to typical deep learning recommendations for regularization.

Training Infrastructure

Optimizer: AdamW with weight decay
Scheduler: ReduceLROnPlateau (patience=10, factor=0.5)
Early stopping: Patience=20 epochs
Checkpointing: Best model saved based on validation loss
Logging: MLflow for experiment tracking

JAX Delta-E Implementation

Located in learning_munsell/losses/jax_delta_e.py:

Differentiable xyY -> XYZ -> Lab color space conversions
Full CIE 2000 Delta-E implementation with gradient support
JIT-compiled functions for performance

Usage:

from learning_munsell.losses import delta_E_loss, delta_E_CIE2000

# Compute perceptual loss between predicted and target xyY
loss = delta_E_loss(pred_xyY, target_xyY)

Limitations

BatchNorm Instability on MPS

Models using BatchNorm1d layers exhibit numerical instability when trained on Apple Silicon GPUs via the MPS backend:

Validation loss spikes during training
Occasional extreme outputs during inference (e.g., 20M instead of ~0.1)
Non-reproducible behavior

Affected Models: Large dataset error predictors using BatchNorm.

Workarounds:

Use CPU for training
Replace BatchNorm with LayerNorm
Use smaller models (300K samples vs 2M)
Skip error predictor stage for affected models

The recommended production model (multi_mlp_multi_error_predictor.onnx) does not exhibit this instability.

References: