| --- |
| license: mit |
| library_name: pytorch |
| tags: [symbolic-mathematics, jepa, representation-learning, symbolic-regression, function-encoder] |
| --- |
| |
| # Func-JEPA |
|
|
| Pretrained weights of *Func-JEPA: Learning a Sobolev Space of Functions as Latent Vectors* (Alejandro Moreno Arcas, |
| 2026). Code: https://github.com/amorarc/Func-JEPA. Data: https://huggingface.co/datasets/amorarc/func-jepa-data. |
|
|
| Func-JEPA encodes the symbolic form of a univariate function into a vector of R^256 whose norm is the exact H¹ norm of |
| the function and whose dot products approximate H¹ inner products. Differentiation is a learned latent operator. Small |
| modules on the frozen encoder decode latents to formulas, compose and integrate in latent space, and map sampled values |
| to latents. |
|
|
| ## Usage |
|
|
| ```bash |
| pip install git+https://github.com/amorarc/Func-JEPA |
| ``` |
|
|
| ```python |
| import numpy as np |
| from fjepa import FuncJEPA |
| |
| m = FuncJEPA.from_pretrained("amorarc/func-jepa") |
| z = m.encode("4*sin(7*x) - 3*cos(7*x)") |
| m.norm(z) # 12.87, the H1 norm of f on [-pi, pi] |
| m.decode(m.derivative(z))[0] # 21*sin(7*x) + 28*cos(7*x) |
| m.regress(32 * np.cos(6 * m.nodes) + 8) # 32*cos(6*x) + 8, from 64 samples |
| ``` |
|
|
| ## Files |
|
|
| | Folder | Parameters | Role | |
| |---|---|---| |
| | `encoder_pi` | 35.1M | encoder + derivative operator on `[-π, π]` (all results of the paper) | |
| | `encoder_05` | 35.1M | the same model trained on `[0, 5]` | |
| | `decoder` | 17.7M | latent → prefix tokens | |
| | `compose` | 52.9M | composition operators: sin, cos, log, sqrt, square, cube, quadratic | |
| | `integrator` | 7.6M | latent of `f` → latent of the zero-mean antiderivative | |
| | `value_encoder` | 7.4M | 64 samples → latent direction | |
| | `decoder_lc`, `value_encoder_lc` | 17.7M, 7.4M | variants fine-tuned for the Lample–Charton ODE benchmark | |
|
|
| Each folder holds `model.safetensors` and `config.json`; `SHA256.json` lists the hashes. |
|
|
| ## Training |
|
|
| Synthetic data only: 2.0M functions from 17 parameterised families and short random expression trees, generated with |
| SymPy. Objective: H¹ angle matching, linearity, and a JEPA term predicting the latent of `f'` from that of `f`. One |
| training run of 40k steps on a single 8 GB GPU (6.6 hours); the modules were trained afterwards on the frozen encoder. |
|
|
| ## Results (from the paper, single run) |
|
|
| - Latent cosines and dot products correlate 0.993 and 0.998 with their H¹ counterparts. |
| - The derivative operator retrieves `f'` among 3,956 candidates for 82.0% of test functions. |
| - Exact symbolic regression from 64 samples: 76% / 76% / 62% at noise 0 / 0.1% / 1% on 50 test functions. |
| - Symbolic integration on the paper's test set: 91.7% verified by SymPy (full solver, which needs the training data). |
|
|
| ## Limitations |
|
|
| - Latent operators are approximate (median relative error of the derivative: 7%). Predicted latents decode to the |
| right shape more often than to exact constants; refit and verify the candidates. |
| - Weak outside the pretraining families: 20–25% on the integration test sets of Lample and Charton and 16–25% on |
| their ODE test sets, where task-specific models exceed 90%. |
| - The symbolic-regression advantage is measured on functions of the pretraining families. |
| - Univariate functions on a fixed interval, 37-token vocabulary without special functions, at most 126 tokens. |
| - Single training run: variability across seeds was not measured. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{morenoarcas2026funcjepa, |
| title = {Func-JEPA: Learning a Sobolev Space of Functions as Latent Vectors}, |
| author = {Moreno Arcas, Alejandro}, |
| year = {2026} |
| } |
| ``` |
|
|
