mixle.reason.embedding module¶
A rate-adaptive common embedding whose active dimension scales with information content.
A fixed-width embedding wastes capacity on low-information inputs and truncates high-information
ones. This encoder learns a shared latent code with a variational per-coordinate posterior
q(z_k | x) = N(m_k(x), s_k(x)^2) and an ARD (automatic relevance determination) gate: a
coordinate whose posterior stays at its prior (KL(q || p) ~ 0) carries no information and is
inactive. The active dimension of an input is therefore #{k : KL(q(z_k|x) || p) > tau} –
it grows, per input and per modality, with the mutual information between the input and the latent.
Training is a rate–distortion (beta-VAE) objective: reconstruct the input subject to a rate budget
on the total KL. The beta knob sets bits-per-embedding; the data decides how those bits are
spent across coordinates, so a dense high-entropy input lights up more coordinates than a sparse one.
Because all inputs share one ordered coordinate system, the codes are comparable across modalities
(a common embedding), and can index a mixle.reason.CrossModalStore.
Torch is imported lazily; mixle.reason exposes this via a deferred attribute.
- class ScaledEmbedding(in_dim, max_dim=16, *, hidden=(64,), beta=1.0, kl_tau=1e-2, seed=0)[source]
Bases:
objectA beta-VAE-style common embedding with an ARD gate giving a data-dependent active dimension.
- Parameters:
in_dim (int) – input feature width.
max_dim (int) – the embedding’s maximum width (upper bound on active dimension).
hidden (tuple[int, ...]) – hidden widths shared by the encoder and decoder trunks.
beta (float) – rate weight in the ELBO (larger -> tighter rate budget -> fewer active dims).
kl_tau (float) – per-coordinate KL threshold (nats) above which a coordinate counts as active.
seed (int) – torch RNG seed.
- fit(X, *, epochs=400, lr=3e-3, weight_decay=0.0)[source]
Train the embedding on unlabeled inputs
X((n, in_dim)) by the beta-VAE ELBO.
- encode(X)[source]
The embedding means
(n, max_dim)– the common code (use with a store’s keys).
- coordinate_kl(X)[source]
Per-coordinate KL from the prior,
(n, max_dim)(nats) – how much each coord encodes.
- active_dim(X)[source]
Per-input active dimension: number of coordinates whose KL exceeds
kl_tau.