# 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**

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

**Two-Stage Models**

12. **MLP + Error Predictor** - Base MLP with unified error correction
13. **Multi-Head + Multi-Error Predictor** - Multi-Head with 4 independent error predictors
14. **Multi-Head + Multi-Error Predictor (Large Dataset)** - Large dataset variant
15. **Multi-MLP + Multi-Error Predictor** - 4 independent branches with 4 independent error predictors
16. **Multi-MLP + Multi-Error Predictor (Large Dataset)** - Large dataset variant
17. **Multi-ResNet + Multi-Error Predictor (Large Dataset)** - Deep ResNet-style branches
18. **Transformer + Error Predictor (Large Dataset)** - Transformer base with error correction
19. **Multi-Head + Cross-Attention Error Predictor** - Cross-attention mechanism for error correction
20. **Multi-Head + 3-Stage Error Predictor** - Three cascaded error correction stages

**Alternative Code Prediction Models**

21. **Multi-MLP (Classification Code)** - 3 regression branches + 1 classification branch (10 logits for hue code)
22. **Multi-MLP (Classification Code) + Multi-Error Predictor** - Classification Code base with 3-branch error predictor (hue, value, chroma)
23. **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)
24. **Multi-MLP (Classification Code) (Large Dataset)** - Classification Code trained on 1.4M samples
25. **Multi-MLP (Classification Code) + Code-Aware Multi-Error Predictor (Large Dataset)** - Large dataset variant of code-aware error predictor
26. **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**

1.  **Dense xyY Grid** (~500K samples)
    - Regular grid in valid xyY space (MacAdam limits for Illuminant C)
    - Captures general input distribution
2.  **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)
3.  **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:

1.  Base model learns the general xyY to Munsell mapping
2.  Error predictor learns residual corrections specific to each component
3.  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:

1.  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)
2.  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     |

![Gamma sweep results](_static/gamma_sweep_plot.png)

_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.

![Training data density](training_data_large_hue_density.png)

_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**

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

**Two-Stage Models**

4.  **Multi-MLP + Error Predictor** - Base Multi-MLP with unified error correction
5.  **Multi-MLP + Multi-Error Predictor** - 3 independent error predictors (BEST)
6.  **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:

```python
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:

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

**Affected Models**: Large dataset error predictors using BatchNorm.

**Workarounds**:

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

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

**References**:

- [BatchNorm non-trainable exception](https://github.com/pytorch/pytorch/issues/98602)
- [ONNX export incorrect on MPS](https://github.com/pytorch/pytorch/issues/83230)
- [MPS kernel bugs](https://elanapearl.github.io/blog/2025/the-bug-that-taught-me-pytorch/)