This document explains the XAI (Explainable AI) analyses in the Convergence MAE notebook.
The Masked Autoencoder learns compressed representations (embeddings) of atmospheric flight data. These analyses help understand what the model learned and whether it captures meaningful physics rather than artifacts.
Question answered: Which embedding dimensions encode which physical quantities?
Method: Computes Pearson correlation between each of the 128/256 embedding dimensions and physical variables (temperature, pressure, humidity, etc.).
How to interpret:
- High correlation (|r| > 0.5) → That dimension explicitly encodes that variable
- Multiple dimensions correlating with one variable → Distributed representation
- No correlations → Model may use nonlinear combinations
Good result: Several dimensions strongly correlate with known physics (altitude, temperature).
Bad result: No significant correlations → Embeddings may not capture interpretable features.
Question answered: What atmospheric regimes does the model distinguish?
Method:
- Clusters embeddings using KMeans (k=6)
- Characterizes each cluster by mean values of physical variables
- Auto-generates labels like "High-Alt / Cold / Dry"
How to interpret:
- Distinct cluster profiles → Model separates different atmospheric conditions
- Clusters matching known regimes (boundary layer, free troposphere, clouds) → Physically meaningful
Good result: Clusters correspond to recognizable flight conditions.
Bad result: All clusters look similar, or clusters separate by flight ID instead of physics.
Question answered: Is the learned structure physical or artifactual?
Method: Projects high-dimensional embeddings to 2D using UMAP, then colors points by:
- Cluster assignment
- Flight ID (artifact check)
- Physical variables (altitude, temperature, etc.)
How to interpret:
- Smooth gradients colored by physics → Model learned physical structure
- Sharp boundaries by flight ID → Model memorized flight-specific patterns (bad)
- Silhouette scores: higher for clusters than flight IDs = good
Good result: UMAP structure correlates with physical variables, not flight identity.
Bad result: Points cluster by which flight they came from.
Question answered: Are physical concepts linearly extractable from embeddings?
Method: Trains simple logistic regression classifiers on embeddings to predict binary labels:
- High vs low altitude
- Warm vs cold
- Moist vs dry
- In-cloud vs clear (if available)
How to interpret:
- Accuracy >> 50% → Concept is linearly encoded (easy to extract)
- Accuracy ≈ 50% → Concept not directly encoded (may require nonlinear decoder)
- Compare to baseline (majority class) to assess improvement
Good result: 80-95% accuracy on physical state classification.
Bad result: Near-chance accuracy suggests embeddings don't capture that concept.
Question answered: What variable relationships did the model learn?
Method: Extracts attention weights from the transformer encoder's first layer, aggregates by variable pairs to show which variables "attend to" each other.
How to interpret:
- High attention between related variables (ATX↔THETA, MR↔RHUM) → Learned physical relationships
- Attention patterns matching known thermodynamics → Model discovered real correlations
- Uniform attention → Model treats all variables equally (less interpretable)
Good result: Thermodynamically related variables show elevated mutual attention.
Bad result: Random or uniform attention patterns.
Question answered: Where does the model struggle? What might be anomalies?
Method:
- Computes per-variable reconstruction MSE
- Correlates total error with flight conditions
- Identifies high-error windows as potential anomalies
How to interpret:
- Per-variable error: Hard-to-reconstruct variables may have complex dynamics or poor correlations with others
- Error vs conditions: If error correlates with altitude/temperature, model struggles in specific regimes
- Anomalies: Windows with unusually high error may represent rare/interesting atmospheric states
Good result: Most variables reconstruct well; anomalies correspond to known interesting events.
Bad result: Uniformly high error, or error correlates with flight ID (artifact).
Consolidates findings from all analyses into a single report with key takeaways.
| Check | Good Sign | Warning Sign |
|---|---|---|
| Correlations | Strong r with physics | No significant correlations |
| Clusters | Match atmospheric regimes | Cluster by flight ID |
| UMAP | Smooth physical gradients | Sharp flight boundaries |
| Probing | >70% accuracy | Near chance (50%) |
| Attention | Related vars attend together | Uniform attention |
| Errors | Low, condition-dependent | High, flight-dependent |
Autoencoders compress, they don't inherently explain. These XAI techniques bridge the gap by:
- Mapping learned representations to known quantities
- Testing whether human-interpretable concepts are encoded
- Visualizing structure for sanity checks
For causal explanations or feature importance on predictions, consider:
- SHAP / Integrated Gradients on downstream tasks
- Concept bottleneck models
- Attention-guided feature attribution