Modular Jets for Supervised Pipelines: Diagnosing Mirage vs Identifiability

TL;DR

<3-5 sentence high-level summary> Modular Jets addresses the gap between risk-based evaluation and understanding a model's internal modular decomposition in supervised pipelines. By treating the input space as a task manifold and instrumenting module taps, the authors estimate empirical jets that capture local linear responses to structured perturbations. They define mirage vs identifiability regimes and prove a jet-identifiability result for linear two-module pipelines, complemented by a practical MoJet algorithm for empirical jet estimation and diagnostics. Across synthetic and real data, jets reveal internal differences that risk metrics miss, providing a scalable, post-hoc auditing tool with potential to guide evaluation design and interpretability in larger systems.

Abstract

Classical supervised learning evaluates models primarily via predictive risk on hold-out data. Such evaluations quantify how well a function behaves on a distribution, but they do not address whether the internal decomposition of a model is uniquely determined by the data and evaluation design. In this paper, we introduce \emph{Modular Jets} for regression and classification pipelines. Given a task manifold (input space), a modular decomposition, and access to module-level representations, we estimate empirical jets, which are local linear response maps that describe how each module reacts to small structured perturbations of the input. We propose an empirical notion of \emph{mirage} regimes, where multiple distinct modular decompositions induce indistinguishable jets and thus remain observationally equivalent, and contrast this with an \emph{identifiable} regime, where the observed jets single out a decomposition up to natural symmetries. In the setting of two-module linear regression pipelines we prove a jet-identifiability theorem. Under mild rank assumptions and access to module-level jets, the internal factorisation is uniquely determined, whereas risk-only evaluation admits a large family of mirage decompositions that implement the same input-to-output map. We then present an algorithm (MoJet) for empirical jet estimation and mirage diagnostics, and illustrate the framework using linear and deep regression as well as pipeline classification.

Figures (8)

• Figure 1: Linear regression sanity check. Each point corresponds to one coordinate of the jet-estimated gradient $\hat{\beta}_{\mathrm{jet}}$ (averaged over base inputs) versus the corresponding ground-truth coefficient $\beta^\star$. The diagonal line indicates perfect agreement. The close clustering around the diagonal confirms that empirical jets recover the true linear parameters up to small numerical error.
• Figure 2: Two-module deep regressor. Left: numerical ranks of bottleneck jets for Model A and Model B under coarse (left cluster) and structured (right cluster) perturbations. Model A consistently has rank $3$, while Model B has rank close to $1$ (coarse) and between $1$ and $2$ (structured). Right: jet similarity between the bottleneck modules of Model A and Model B under coarse and structured perturbations. Average similarity increases from about $0.68$ (coarse) to about $0.84$ (structured), indicating stronger alignment along task-relevant directions.
• Figure 3: Pipeline classification (PCA and logistic regression vs. dense classifier). Left: numerical ranks of module 1 jets under coarse and structured perturbations. The PCA module is always rank $3$; the dense model is rank $20$ under coarse noise but collapses to rank $3$ when perturbations are restricted to the principal subspace. Right: jet similarity between the PCA module and the dense model's identity module. Similarity is low under coarse perturbations ($\approx 0.44$) but essentially $1.00$ under aligned perturbations, indicating that the two pipelines behave identically along task-relevant directions but differ sharply in their treatment of nuisance directions.
• Figure 4: Digits classification (PCA and logistic regression vs. one-hidden-layer MLP). Left: numerical ranks of module-1 jets under coarse and aligned perturbations for both models. The PCA module has rank $10$ by construction; the MLP hidden layer exhibits higher rank under coarse perturbations (median $18$) but collapses to rank $10$ under aligned perturbations. Right: jet similarity between the PCA module and the MLP hidden layer. Similarity is moderate under coarse perturbations (mean $0.649996$) but essentially $1.0$ under aligned perturbations (mean $0.999986$), indicating that the two pipelines behave almost identically along task-relevant directions while differing in their response to nuisance directions.
• Figure 6: JetSim as a function of subspace dimensionality $k$ for the PCA and MLP jets on the digits task. The vertical dashed line indicates the $k=54$ chosen by the $95\%$ variance rule. JetSim continues to increase smoothly beyond the spectral "elbow", and remains high for $k$ in the range $[40,60]$, showing that our similarity conclusions are not sensitive to the exact cutoff.

Theorems & Definitions (4)

• theorem 1: Jet-based identifiability in linear two-module pipelines
• proof
• corollary 1: Consistency of linear jets under ideal perturbations
• proof