Representation Geometry as a Diagnostic for Out-of-Distribution Robustness

TL;DR

This work presents TorRicc, a geometry-based post-hoc diagnostic that uses class-conditional mutual k-NN graphs constructed from in-distribution embeddings to derive global spectral complexity via a torsion proxy and local regularity via Ollivier--Ricci curvature. By combining these metrics into GeoScore, the authors show that lower torsion and higher curvature reliably predict better OOD performance across architectures, shifts, and checkpoints, enabling near-oracle unsupervised checkpoint selection under distribution shift. The framework is validated on CIFAR-10 and Tiny-ImageNet-C, with extensive ablations and sanity checks demonstrating sensitivity to task-aligned geometry rather than superficial statistics. Overall, representation geometry provides a principled, label-free diagnostic for robustness monitoring and model selection in shifting environments, with broad practical implications for deployment under distribution shift.

Abstract

Robust generalization under distribution shift remains difficult to monitor and optimize in the absence of target-domain labels, as models with similar in-distribution accuracy can exhibit markedly different out-of-distribution (OOD) performance. While prior work has focused on training-time regularization and low-order representation statistics, little is known about whether the geometric structure of learned embeddings provides reliable post-hoc signals of robustness. We propose a geometry-based diagnostic framework that constructs class-conditional mutual k-nearest-neighbor graphs from in-distribution embeddings and extracts two complementary invariants: a global spectral complexity proxy based on the reduced log-determinant of the normalized Laplacian, and a local smoothness measure based on Ollivier--Ricci curvature. Across multiple architectures, training regimes, and corruption benchmarks, we find that lower spectral complexity and higher mean curvature consistently predict stronger OOD accuracy across checkpoints. Controlled perturbations and topological analyses further show that these signals reflect meaningful representation structure rather than superficial embedding statistics. Our results demonstrate that representation geometry enables interpretable, label-free robustness diagnosis and supports reliable unsupervised checkpoint selection under distribution shift.

Figures (10)

• Figure 1: Overview of the proposed TorRicc pipeline for post-hoc robustness diagnosis. Given in-distribution embeddings from a trained checkpoint, we construct class-conditional mutual $k$-NN graphs and extract complementary global (torsion-inspired spectral complexity) and local (Ollivier--Ricci curvature) invariants. These geometry-based signals are fused into a lightweight score that enables unsupervised checkpoint ranking and robustness monitoring under distribution shift.
• Figure 2: Torsion proxy vs. OOD accuracy on CIFAR-10.1. Lower torsion correlates with higher robustness.
• Figure 3: Mean Ollivier--Ricci curvature vs. OOD accuracy. Higher curvature correlates with better robustness.
• Figure 4: Correlation vs. CIFAR-10-C severity.
• Figure 5: Control study under structure-breaking interventions.