jax-conformal-frame

name	jax-conformal-frame
description	Use this skill when constructing, training, evaluating, saving, loading, or documenting the JAX plus Equinox conformal coordinate frame field for disentanglement, a noise-conditional model that learns a local conformal frame J(x) = lambda(x) U(x) from data through score matching, integrability, and independence losses.

JAX Conformal Frame

This skill documents a self-contained JAX and Equinox implementation of the conformal coordinate frame field from the workshop paper "Conformal Coordinate Frames for Disentanglement". The method learns a local frame J(x) = lambda(x) U(x), where U(x) is orthogonal and lambda(x) > 0 is a single conformal factor shared across all directions. Each column of J(x) is a learned disentangled direction. Three objectives shape the frame. Noise conditional score matching fits the data distribution, a stochastic integrability loss makes the frame a coordinate frame of some diffeomorphism, and a stochastic independence loss enforces independent latent components.

The reference research code lives in the local_disentanglement package and is tangled with generax and local_coordinates. This skill vendors a fresh trimmed port under jax_conformal_frame/ that depends only on jax, equinox, optax, and diffrax. It ships four runnable examples and a five-file test suite.

The port covers the paper's headline configuration. The model is conformal and noise conditional, U(x, sigma) is the Q factor of a QR decomposition with sign correction, and the data are vectors. The general per-dimension independent mechanism analysis model, the image path, the matrix-exponential and Cayley orthogonal parameterizations, the c_ima independence variant, and the pullback-metric and exp-log-map geometry methods are out of scope.

cd ~/.claude/skills/jax-conformal-frame
uv sync --extra test
uv run python construct.py
uv run python train.py --steps 300 --dim 8
uv run python evaluate.py --dim 8 --steps 300
uv run python save_load.py
uv run python sanity_check.py        # long run, paper scale
uv run pytest tests

Module Map

The package is jax_conformal_frame. The top-level API is the config and the model. Energies, losses, and metrics are imported from their submodules.

• jax_conformal_frame.config.ConformalFrameConfig is the frozen dataclass that records the architecture hyperparameters and round-trips through JSON.
• jax_conformal_frame.embeddings.sinusoidal_embedding maps a scalar noise level to a fixed sinusoidal feature vector. It has no parameters.
• jax_conformal_frame.networks holds TimeConditionedMLP and VectorOrthogonalQR, the orthogonal matrix network.
• jax_conformal_frame.energies holds AbstractEnergy, Funnel, RotatedFunnel, and random_rotation_matrix.
• jax_conformal_frame.model.ConformalFrameModel is the model, with the training methods plus jacobian_flow and jacobian_flow_inv.
• jax_conformal_frame.losses holds the three losses, the NCSM block, the integrability block, and the independence block.
• jax_conformal_frame.metrics holds the funnel-axis alignment metric and the random-frame baseline.

Construct

ConformalFrameModel.init(cfg, key) builds the full model. Equinox stores parameters as a pytree of arrays. Static fields such as dim, sigma_min, and sigma_max are marked eqx.field(static=True). The config records architecture hyperparameters only. The loss weights and the learning rate are training choices and live as arguments to train.py.

cfg = ConformalFrameConfig(dim=8, hidden_size=128, n_layers=3)
model = ConformalFrameModel.init(cfg, jax.random.PRNGKey(0))

The model operates on a single point of shape (dim,). Batching is the caller's job through jax.vmap. See construct.py for a runnable version with shape assertions and an orthogonality check on U.

Train

train.py fits a model on a randomly rotated Neal's funnel. The data are whitened first. The whitened score is the energy score composed with the whitening map, which the independence loss needs. The optimizer is AdamW with a warmup cosine decay schedule and global gradient clipping. The single eqx.filter_jit sits on train_step, the outermost call, per the project's JIT discipline.

@eqx.filter_jit
def train_step(model, opt_state, batch, key, optimizer, score_fn,
               lambda_int, lambda_ind, sigma_min, sigma_max):
    def loss_fn(m):
        loss_score = ncsm_loss_batch_continuous(batch, m, sigma_min, sigma_max, k_score)
        loss_int = integrability_loss_batch(batch, m, k_int)
        loss_ind = independence_loss_batch(batch, m, score_fn, k_ind)
        total = loss_score + lambda_int * loss_int + lambda_ind * loss_ind
        return total, (loss_score, loss_int, loss_ind)
    (total, parts), grads = eqx.filter_value_and_grad(loss_fn, has_aux=True)(model)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    return eqx.apply_updates(model, updates), opt_state, total, parts

The independence loss is heavy tailed on Neal's funnel because the data score grows large in the funnel's low-variance tail. The total loss is therefore dominated by the independence term and does not decrease monotonically in a short run. train.py is a smoke test of the training machinery. It asserts that every loss term stays finite. The paper's results come from a long run, roughly 50000 steps with early stopping on the score loss. The default training weights are lambda_int = 1.0 and lambda_ind = 0.1.

Evaluate

evaluate.py trains a model briefly, then measures how well the learned frame recovers the funnel axis. Per-point alignment is the largest absolute inner product between a column of U(x) and the funnel direction. The script reports the learned alignment against the random-frame baseline and its large-dimension asymptote sqrt(2 log D / D).

The funnel axis is a direction in raw observation space. Training whitens the data, so the axis is rescaled by 1 / std before the alignment is computed. A short run recovers the axis well in low dimensions and weakly in high ones. The paper trains longer across several seeds and reports alignment from 0.87 to 0.99 in dimensions 4 through 64.

Sanity Check

sanity_check.py trains at the scale used in the paper, roughly 50000 steps with batch 512, warmup 1000, and the same AdamW schedule. It evaluates every 2500 steps, printing the funnel-axis alignment, the Fisher divergence, and the integrability loss. At the end it saves the final model and asserts that the best alignment over training clears both a floor (--min-alignment, default 0.8) and the random-frame baseline by a clear margin.

This is an end-to-end correctness check, one (dimension, seed) run rather than the paper's full sweep. It is a long run. On a GPU it finishes in minutes to under an hour. On a CPU it takes hours, so prefer a GPU.

uv run python sanity_check.py --dim 8 --seed 0
uv run python sanity_check.py --dim 16 --seed 3 --steps 50000

The quick train.py smoke test cannot show convergence, since its total loss is dominated by the heavy-tailed independence term. sanity_check.py is the script that confirms the method recovers the funnel axis.

Save And Load

save_load.py mirrors the jax-tarflow convention. Each saved model is one directory under MODEL_ROOT = Path("pretrained_models") with weights.eqx and metadata.json. The metadata serializes the ConformalFrameConfig so the skeleton can be rebuilt before the leaves are deserialized.

eqx.tree_serialise_leaves(model_dir / "weights.eqx", model)
metadata = {
    "name": name, "description": description,
    "saved_at": datetime.now(timezone.utc).isoformat(),
    "step": int(step), "loss": float(loss),
    "num_parameters": count_array_parameters(model),
    "config": asdict(cfg),
}

To load, reconstruct the skeleton from the saved config and then deserialize.

cfg = ConformalFrameConfig(**metadata["config"])
skeleton = ConformalFrameModel.init(cfg, key)
model = eqx.tree_deserialise_leaves(model_dir / "weights.eqx", skeleton)

No pretrained checkpoint ships with this skill. The example saves a freshly initialized tiny model and asserts an exact output roundtrip after reload.

Tests

Five test files under tests/ exercise the port with fast deterministic checks. There is no training test.

• test_orthogonality.py checks that VectorOrthogonalQR produces a genuine orthogonal matrix and that inverse is its exact transpose.
• test_conformal_log_s.py checks that the log loadings are uniform across the D coordinates and that score_model equals minus the gradient of neg_log_p.
• test_integrability_estimator.py builds the dense structure-equation violation from jax.jacobian of U and jax.grad of the negative log density, checks the Lie-bracket formula and antisymmetry, and verifies that the stochastic loss is unbiased.
• test_independence_estimator.py checks the masked-probe estimator against an explicit matrix and against the dense scaled Hessian of the model.
• test_jacobian_flow.py checks the diffrax-based jacobian_flow and the Jacobian column accessor.

Run the suite with uv run pytest tests.

Notes For Agents

• Keep JAX code in this style. Equinox modules with parameters as pytree leaves, and eqx.filter_jit only on the outermost train_step, so XLA compiles the whole graph at once.
• The matrix-vector methods U_mvp, U_mvp_inv, J_mvp, J_mvp_inv, and get_jacobian_column keep an unused key argument. The QR parameterization is deterministic. The argument exists so the vendored loss code matches the reference interface. Do not remove it.
• The conformal structure rests on log_s(x) = neg_log_p(x) / D broadcast to all D coordinates. The factor 1 / D is load bearing. Without it the Jacobian would not be lambda U.
• sigma is optional throughout. A value of None resolves to a scalar 0, which is the clean-data end of the noise range. The integrability and independence losses evaluate the frame on clean data, so only the NCSM loss uses a nonzero sigma.
• The QR sign correction Q = Q * sign(diag(R)) has a discontinuity where a diagonal entry of R crosses zero. It is the parameterization stated in the paper. Smooth alternatives such as the matrix exponential exist but are out of scope here.
• The stochastic integrability_loss returns an unbiased estimator of 2 * sum_{i<j} ||V_ij||^2, with the factor of 2 folded in. The independence_loss ships the masked-probe frobenius estimator only.
• Training whitens the data. Any direction expressed in raw observation space, such as the funnel axis, must be rescaled by 1 / std to be compared with the learned frame.