ePC inference solver for FabricPC - #47
Conversation
cgoemaere
left a comment
cgoemaere
left a comment
There was a problem hiding this comment.
Hi Matthew, I added my comments, making sure to address your 5 points (and a bit more). Overall, it looks like a solid plan and I had mostly minor remarks.
As I mentioned in my email, I'm currently working on a cleaner solution to port ePC's advantages to arbitrary graphs. In the meantime, your unrolling approach sounds like a good start for the problem. I'll get back to you when I have more results for my idea (I'll try to benchmark it against the unrolling method to see how it compares).
Feel free to ask for more feedback at any point; I'd be happy to see this PR come through!
e82c701 to
77ea096
Compare
matthewbehrend
changed the title
|
Original design was running forward twice because the node contract bundled projection, error pairing, and energy computation in the node forward method. Split the three node operations. Architecture is cleaner and no redundant calls. |
|
rebased on 0.4.0 |
e70098f to
dddd215
Compare
[epc_convergence.html](https://github.com/user-attachments/files/31688133/epc_convergence.html)
Attached charts for ePC settling.
sPC and ePC comparisons were on the ResNet-18 demo /examples/resnet18_cifar10_demo.py
100 epochs training:
A few observations from the charts:
Resnet18/CIFAR-10 Results (2 epochs):epochs/arm: 2 | trials: 5 | sPC: 120 steps @ eta 0.1 python examples/epc_spc_resnet18_compare.py --mode sweep --n_trials 5 --epc_eta 0.1 python examples/epc_spc_resnet18_compare.py --mode sweep --n_trials 5 --epc_eta 0.03 python examples/epc_spc_resnet18_compare.py --mode sweep --n_trials 5 --epc_eta 0.01 python examples/epc_spc_resnet18_compare.py --mode sweep --n_trials 5 --epc_eta 0.001 python examples/epc_spc_resnet18_compare.py --mode sweep --n_trials 5 --epc_eta 0.0001 ============================================================ Resnet18/CIFAR-10 Results (100 epochs):Trainer: pc | Inference: epc (eta 0.001, 1 steps) | Activation: gelu | Epochs: 100 | LR: 0.001 | Augment: True Training time: 1102.0s (11.0s per epoch) Trainer: pc | Inference: epc (eta 0.001, 2 steps) | Activation: gelu | Epochs: 100 | LR: 0.001 | Augment: True Training time: 1546.5s (15.5s per epoch) (instability with increasing steps) Training time: 2547.2s (25.5s per epoch) |
|
Hi Matthew The figure looks exactly right! sPC struggles to push energy through the network, with a noticeable delay for deeper layers. By contrast, ePC optimizes everything at once. And larger eta gets faster convergence to the same equilibrium. This is exactly as it should be. Everything you mention adds up. Here's my view per observation bullet:
I had a quick look at the code of inference_epc.py. It looks correct, but of course, it's hard to judge properly. I'd be more confident with a test that involves a linear oracle. I have one here or you could also use the formula from Innocenti (Theorem 1 / Equation 5). If the ePC inference solver returns the exact optimal equilibrium state / energy for a bunch of different linear networks, that's a strong signal that the implementation is correct. Same goes for the sPC inference solver, of course. As for a test for the oracle itself (because nobody likes a wrong oracle), I believe some hard-coded numbers on tiny models can do the trick. (or the code should be so readable that it's obviously true; doable with Innocenti's Theorem 1). Hope this helps! |
dec1d8d to
f864577
Compare
…method with unroll degree argument. Enumerate ePC process implementation with framing like the existing backpropagation design pattern following topological sequence.
…itialization for `in_degree == 0` nodes to ensure correct state propagation. Remove unnecessary state forcing in solvers to support diverse node types.
…efactor to scope.
Template methods (inference_step, update_latents, forward_value_and_grad) dispatched by re-resolving their own class via type(structure.config["inference"]), so any composed solver would re-dispatch to the composition object instead of itself. They are now classmethods dispatching on cls. run_inference becomes an instance method reading self.config, wrapped by new begin_segment/finalize_state classmethod hooks (default identity) and a segments() instance method returning ((self, infer_steps),) for per-step consumers. The tracking history variants iterate segments() instead of reading config["infer_steps"] directly, so composed schedules are tracked segment by segment with metric stacks concatenated along the step axis. Single-solver output is unchanged. Tests calling the static run_inference form migrate to the instance call; conftest's with_inference gains an optional prebuilt inference argument. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_topological_sort gains an unroll argument. DAGs keep the exact Kahn/BFS order. A cyclic graph now raises GraphCycleError unless graph(..., unroll=U) is given, replacing the print-warning that returned a partial order omitting cycle members and everything downstream. With unroll=U, strongly connected components (iterative Tarjan) collapse to a condensation DAG ordered by the same Kahn seeding rules; each nontrivial component's members are emitted U times in BFS order from its entry nodes. x->a<->b->y at U=2 yields (x, a, b, a, b, y). GraphStructure gains a schedule field holding the full visit schedule; node_order becomes first_occurrence_order(schedule) (equal on DAGs). The pytree aux tuple is extended in matching positional order. Consumer migrations: FeedforwardStateInit pass 2 walks structure.schedule, so cyclic graphs gain true feedforward initialization through cycles; initialize_graph_state gains a shared post-pass assigning z_mu <- z_latent (cast to z_mu's float dtype) for every in_degree == 0 node, fixing the z_mu = 0 / error = 0 inconsistency all three initializers left on source nodes; muPC's compute_mupc_scalings and _count_skip_connections_depth raise on duplicate node_order entries (they model one energy term per merge node and must never receive the unrolled schedule). Cyclic call sites (tests' _build_cycle, examples/mnist_cyclic_graph.py) pass unroll explicitly, and the cyclic-graph section of the building-models guide gains a complete graph(..., unroll=U) example. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
forward() fused three stages behind the sPC dataflow direction: predict (z_mu from params and in-edge inputs), pair (error = z_latent - z_mu, copy-pasted into all 18 node bodies), and score (the energy functional plus in-forward custom terms). sPC consumes z_latent -> error; the error-parameterized solver needs error -> z_latent, and needs z_mu before z_latent exists. Node authors now implement predict(params, inputs, state, node_info) -> (z_mu, aux) and optionally override energy(params, inputs, state, aux, node_info) for custom terms. NodeBase owns the pair (pair_error / pair_latent — one volume-preserving bijection shared by both solver directions) and the assembly templates forward, forward_with_aux, and forward_from_error (the ePC derive direction: z_latent = z_mu + error, with the clamp deciding the free side). These are audited as non-override-points by the new tests/test_node_contract.py. energy_functional is deleted; its body is the default energy(). Source semantics get one owner: the template's in_degree == 0 guard (previously inlined in the solver branch, where IdentityNode.forward would crash if ever called on a source). The unclamped-readout forcing branch in forward_and_latent_grads is deleted: it zeroed error, energy, and latent_grad — contradicting the method's own contract and discarding in-forward energy terms, so a Hopfield readout could never settle onto its attractor as an output node. Unclamped readouts now take the ordinary autodiff path; eval accuracy is unchanged (predictions read z_mu), reported eval energy now includes readout energy previously zeroed. Migration is complete across the 13 library forward() bodies (Linear's _forward_with_preact is deleted — predict returns pre_activation as aux, and LinearExplicitGrad's analytic overrides call forward_with_aux, with source semantics delegated to the base short-circuit), StorkeyHopfield's accumulate_hopfield_energy becomes its energy() override with today's op order preserved, and the five external custom nodes in examples and tests. The custom-nodes guide is rewritten around the two-method contract, the node API reference gains the contract and aux pattern/anti-pattern (aux is snapshotted at predict time — under ePC before z_latent is derived — so z_latent-dependent energy terms read state.z_latent inside energy()), and CHANGELOG records the breaking change. Full suite passes with no expectation edits on pre-existing tests: sPC outputs are bit-identical. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
EPCInference relaxes the prediction errors and derives the latents: each inference step derives every node's state along structure.schedule (z_latent = z_mu + error via forward_from_error, the clamp deciding the free side), takes one jax.value_and_grad of the total in_degree > 0 energy with respect to the error pytree, accumulates the gradient into latent_grad, and steps every unclamped node's error (compute_new_error, with decay). finalize_state runs one detached derive pass so the returned state satisfies z_latent = z_mu + error with energies at the final point, feeding the existing local weight-gradient path unchanged. Because the error gradient is taken through the full network's transfer function, eta_infer defaults to 1e-3 and is tuned like a weight learning rate; sPC's local rates overshoot along the global gradient. Clamped nodes never enter the relaxed pytree, which also keeps int-dtype token sources out of AD. muPC's scale_inputs applies inside the differentiated forward; the per-hop gradient preconditioners (jacobian_gain, self_grad_scale) condition sPC's one-hop updates and are not replicated in the global reverse pass. Tests pin: error = 0 <-> feedforward init at any unroll degree; gradient correctness against the closed form and a hand-rolled jax.grad; energy descent; shared equilibria and weight gradients with InferenceSGD on a strictly convex DAG including an unclamped top-down prior; the forward_from_error branches (CrossEntropy-clamped output, Gaussian readout staying at zero error, Hopfield readout receiving its attractor gradient, int-token EmbeddingNode); cyclic warm-start semantics under the unrolled schedule; muPC input scaling; insertion-order independence of one-step gradients; and the z_latent = z_mu + error invariant of the finalized state. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
InferenceSchedule folds the graph state through its component solvers in order — e.g. a few cheap global ePC steps to near-equilibrium, then sPC refinement on the true arbitrary-graph energy, warm-started from ePC's solution. Each solver receives z_latent, z_mu, and error exactly as the previous segment (or the initializer) left them; each applies its own begin_segment/finalize_state. segments() flattens nested schedules for per-step consumers (the tracking module iterates it since the dispatch refactor); inference_step and compute_new_latent raise, since a schedule has no single per-step rule. Exports EPCInference and InferenceSchedule from fabricpc.core. The inference API guide documents both solvers (including the tune-eta-like-a-weight-learning-rate guidance and the unrolled-energy semantics on cyclic graphs), and the predictive-coding guide introduces the error parameterization and schedule composition beside the state-based inner loop. Tests pin: segment flattening including nesting; single-solver schedule bit-identical to the plain solver; ePC-then-sPC bit-identical to manual sequential calls (the boundary passes state as-is); energy non-increase across the handoff; execution inside jax.jit(train_step); tracking parity (one metric row per step across segments, final state matching run_inference); and the raising stubs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ject build_resnet18 takes a required inference: InferenceBase in place of the infer_steps/eta_infer kwargs and the hardcoded InferenceSGDNormClip, so one builder serves both solvers and composed schedules; _create_mupc_model and run_single_mupc migrate, with CLI behavior unchanged (run_single_mupc constructs the same InferenceSGDNormClip from its arguments). New examples/epc_spc_resnet18_compare.py (importlib load of the demo builder, per PC_backprop_compare): - --mode sweep: one PlannedMultiContrastExperiment with an arm per ePC step count T1 plus the sPC baseline, empty contrast family (the runner supplies the paired trial loop). All arms train the same epochs, so each arm is one (wall-clock, accuracy) point; per-trial accuracy at equal wall-clock interpolates the ePC points at sPC's time, and wall-clock to equal accuracy takes the smallest-T1 arm reaching sPC's accuracy. Both are tested with paired_ttest and cohens_d across trials. Two-panel plotly chart (accuracy and wall-clock vs T1, log x, sPC reference line with SE band) written to epc_step_sweep.html, with png behind a kaleido import guard. - --mode convergence: single seed, no training; identical params and initial state for both solvers on one test batch via run_inference_with_history. Reports per-node energy-vs-step (one line per node colored by schedule depth, side-by-side panels — a global curve can read as sPC near-convergence while deep nodes have received no signal), the E* criterion (ePC steps to reach sPC's final total energy), and measured post-warmup per-step wall-clock for both solvers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The measured resnet18/CIFAR-10 convergence (epc_spc_resnet18_compare.py --mode convergence, recorded in the script docstring) showed 1e-3 needing 105 steps to reach sPC-120's final total energy versus 12 at 1e-2 and 5 at 3e-2: the weight-learning-rate starting point was an order of magnitude too conservative. The constructor docstring and the inference API guide carry the measurements; the compare script's --epc_eta now falls back to the constructor default instead of --lr so the two cannot drift. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Introduce pytest markers for long-running convergence tests in `pyproject.toml`. - Update `state_initializer.py` to initialize error states correctly. - Modify `graph_construction.py` to improve validation of the `unroll` parameter. - Add a utility function for total energy calculation in `conftest.py`. - Update documentation to clarify the behavior of readout nodes and inference methods.
…aining checkpoints.
…acy results of the ePC inference rate in the design file
…ests for new node contract, replacing assertion energy >= 0 with a tolerance check.
This new function compiles state initialization and inference into a single XLA program, ensuring consistent energy initialization across calls. It replaces the previous separate inference method to prevent discrepancies in cuDNN conv algorithm selection.
…osed inference solvers. Document ePC steps > 1 required for PC solution.
…a_max and discussion of stability issues
…nostics fabricpc/utils/linear_pc_oracle.py assembles the energy of a linear-Gaussian DAG as a quadratic E = 1/2 ||A z_free - c||^2 from params, edges, muPC forward scales, biases, and precisions, and solves it by least squares. It never calls node or solver code, so it is an independent reference for EPCInference and the state-based solvers. theorem1_energy gives the closed-form chain energy of Innocenti et al. 2024 (Theorem 1), extended to per-node precisions and biases. The diagnostics expose the latent and error Hessians, the stability bound 2/lambda_max, the excited spectrum, and the per-mode relaxed fraction 1 - (1 - eta*lambda)^T that sets the backprop-like and near-equilibrium regimes of ePC. The DAG check tests every edge against node_order: a cycle unrolled once visits each member once, so the schedule length is not a DAG test. tests/test_linear_pc_oracle.py pins the oracle on hand-computed scalar numbers, Theorem 1 against least squares on chains with biases, precisions, and muPC scaling, the precision-weighted error pull-back, the explicit error-Hessian form, the eigenvalue floor, and the validator's rejections (tanh, CrossEntropy, flatten_input, StorkeyHopfield, cycles at unroll 1 and 2). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… tests EPCInference.error_energy returns the total energy as a function of the relaxed errors together with the current error pytree; forward_value_and_grad differentiates it (same ops, same order, so gradients are unchanged) and the same closure now feeds Hessian-vector products and power iteration. linear_pc_oracle gains make_top_epsilon_eigenvalue / top_epsilon_eigenvalue, power iteration on that HVP that measures lambda_max of the error Hessian on any graph, so the stability bound 2/lambda_max is available beyond the linear oracle. regime_label(lambda_max) names the (eta, steps) regime from the fastest excited mode's relaxed fraction 1 - |1 - eta*lambda_max|^T: backprop-like, partially relaxed, near equilibrium, or unstable above eta*lambda_max = 2. It requires the measured lambda_max: an eta*T-only label reads the slowest mode at the unit-precision floor and misreads the resnet18 sweep by about ten times. The class docstring replaces the infer_steps > 1 caution with the mechanism (one step leaves epsilon = -eta * backprop activation gradient exactly; weight gradients match backprop to first order in eta*lambda_max, exactly for layers fed only by clamps) and records the measured resnet18 outcome: the defaults collapsed at epoch 20 of a 100-epoch run while one and two steps survived. Tests: EPCInference and InferenceSGD (muPC chain included) reach the oracle's equilibrium on twelve graphs; the 0.95/1.05 stability bracket pins the scale of both gradient implementations; the HVP and power iteration match the oracle's H_eps; ePC's epsilon gradient at zero equals the backprop activation gradient; one-step errors equal -eta * that gradient; one-step local weight gradients equal eta * backprop on hidden layers and backprop on the output, with the O(eta^2) remainder measured to scale linearly in eta (batch 1, so summed and mean gradients coincide); the regime-label bands. conftest gains inject_biases. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…analysis Answers the reviewer's bullets on the resnet18 ePC-vs-sPC figure with exact results from the linear oracle and the same diagnostics on nonlinear graphs through EPCInference.error_energy. Four CPU sections run in about twenty seconds: - backprop_regime: 1-step ePC weight gradients approach eta * backprop (hidden) and backprop (output) at first order in eta, exactly on a layer fed only by the clamp; one Adam step removes the eta scaling; a one-eigenvalue fit of the recorded 2-epoch sweep through the relaxed fraction 1 - (1 - eta*lambda)^T gives lambda_eff = 12; the formula matches the solver to four decimals on a linear chain. - equilibrium_profile: per-layer equilibrium energies versus depth and weight scale (11.7 decades of spread at depth 20, std 0.5); the sPC transient is top-heavy for hundreds of steps and reaches the oracle after thousands; ePC moves every layer from the first update. - convergence_spectra: H_z and excited H_eps spectra versus depth; sPC needs 30k steps at depth 20 where ePC needs 75; measured contractions agree. - stability: lambda_max(H_eps) versus weight scale and depth (the bound 2/lambda_max shrinks as weights grow, the mechanism behind late collapse); power iteration matches the oracle to 6e-8; a gelu MLP bracket. GPU options: --resnet18 measures lambda_max at init on the demo's muPC resnet18 (16.4 on one 64-sample batch, a factor 1.4 from the sweep fit) and labels every recorded sweep cell; --track_lambda_max N trains the demo graph with its optimizer (optionally on the 100-epoch schedule via --schedule_epochs) and logs eta*lambda_max on a fixed probe batch beside train energy and test accuracy, writing a CSV per cell. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…oint The one-line caution "use infer_steps > 1" keyed on the wrong criterion. The regime parameter is eta_infer * infer_steps * lambda_max(H_eps): one step from zero error leaves the errors at -eta_infer times the backprop activation gradient, each excited mode then relaxes by 1 - (1 - eta*lambda)^T, and eta*lambda_max < 2 is required for stability at every step count, one step included. The inference guide gains a backprop-regime paragraph with the rule, the measured resnet18 numbers (lambda_max 16.4 at init, sweep-fitted lambda_eff 12, the defaults collapsing at epoch 20 of a 100-epoch run while 1 and 2 steps survived), and tuning rows stated in terms of the bound; the composed-schedule example moves off (1e-2, 5), the cell that collapsed at every step count. The predictive-coding guide and the troubleshooting FAQ name the ePC route to backprop. The resnet18 demo prints EPCInference.regime_label at init from lambda_max_at_init (power iteration on one test batch) and its docstring records the six 100-epoch outcomes, relabels the --infer_steps 1 run as the backprop-equivalent regime beside the backprop reference, and replaces the "fewer collapses" and "deeper layers learn" claims with the mechanism. The compare script's per-arm report adds a regime column and its docstring carries the sweep interpretation (the 38.8% figure is ePC's own small-eta*T limit; no backprop arm was run). The archived ePC design doc gets the interpretation under its tables and a note that its 1e-2 default was superseded. CHANGELOG entries for the oracle, error_energy, regime_label, and the analysis script. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
--resnet18 and --track_lambda_max now derive graph, train, and eval keys with the demo's three-way split of PRNGKey(--seed), default 42, and the loaders use the same seed, so the tracked run is the demo's first trial (the 100-epoch run that collapsed at epoch 20) with probes added, and the measured lambda_max belongs to the graph the demo trains. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…_max > 2 scripts/epc_analysis.py --track_lambda_max 50 --num_epochs 30 --schedule_epochs 100 --augment (seed 42, the demo's first trial on the 100-epoch schedule, probes on a fixed 64-sample test batch every 50 updates): - eta 1e-3, 5 steps (the defaults): reproduces the log's 54.76% at epoch 10, peaks at 56.36% at epoch 12; lambda_max grows 16 -> 51 (epoch 10) -> 130 (12) -> 470 (13) -> 3500 (14, eta*lambda_max first above 2 at update 2700) -> 12500, then a dead network (lambda_max = 1) and chance accuracy at epoch 15. Accuracy started falling at epoch 13 as eta*lambda_max of 0.2-0.5 took the run out of the backprop regime. - eta 1e-2, 1 step: lambda_max 15 -> 40 by epoch 5 -> 220 in epoch 6 (crossing at update 1150), chance at epoch 7. One step cannot iterate, but with eta*lambda_max > 2 the step lands each mode farther from equilibrium than it started. lambda_max grew about threefold per epoch once training was under way, so the bound measured at init is a starting point, not a guarantee. The demo docstring, the inference guide, the CHANGELOG, and the dev plan record the outcome. The script gains a two-panel chart of lambda_max (log, against 2/eta) and per-epoch accuracy, rendered at the end of each tracked cell with --plot or from existing CSVs with --plot_track. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Regime reads f, not lambda_min,exc carry Ritz values and gradient weights relative breakdown guard
…cept the pieces that need the `IterContext` trainer or the GPU control runs. Nothing is committed. **What landed** - `fabricpc/core/epsilon_spectrum.py`: `lanczos_extremes`, `EpsilonSpectrum` (with `from_modes` and `host`), `weighted_relaxed_fraction`, `make_epsilon_spectrum`, `epsilon_spectrum`. Exported from `fabricpc.core`. - `EPCInference.regime(spectrum) -> Regime` with the rewritten class docstring. `regime_label`, `top_epsilon_eigenvalue`, and `make_top_epsilon_eigenvalue` are deleted and every caller migrated. The oracle is NumPy-only and gained `gradient_weights`, an exact `weighted_relaxed_fraction`, and the `stability_bound` guard. - `fabricpc/training/regime_probe.py`: `RegimeProbe` and `read_regime_csv`, exported from `fabricpc.training`. - Demo: `spectrum_at_init`, the settings lines, `--track_regime`, `--schedule_epochs`, the probe wired into both callbacks. Compare script and `scripts/epc_analysis.py` migrated; the tracking loop and its eight arguments are gone; `--plot_track` reads the metadata columns and renders four panels. - Tests: `test_epsilon_spectrum.py` and `test_regime_probe.py` are new; the oracle and ePC tests are migrated. Guide 12, 16, 03, CHANGELOG, the reviewer-response plan, and the report are updated. **Verification.** Full suite: 631 passed, 7 skipped. The CPU analysis script runs in 19 s. On the local RTX 3090 the demo's one-epoch run prints the regime line, `--resnet18` runs in 31 s, and `--plot_track` renders both a PC and a backprop CSV. **Two findings that change the picture** 1. A defect in the plan's estimator surfaced on the depth-5 chain: with ten excited modes in float32 the breakdown guard never fired and λ_min came out at the unexcited floor (1.000 instead of 1.726). The floor ghosts carry weight below 1e-13. I added an eps(dtype) weight floor for the extremes, which is the same cutoff the guard expresses through β. Recorded in the plan under Design 1 and pinned by a regression test. 2. The ResNet-18 spectrum at init contradicts the plan's compact-band premise. λ_max is 16.45 as before, but λ_min is −0.42 (indefinite at init, 1.2% of the gradient on negative curvature), and f̄ for the defaults is 0.010 against f_max 0.080. The 2-epoch sweep accuracy follows f_max, not the f̄ band: cells at η = 0.01, T = 3 to 5 read backprop-like while accuracy has left the backprop value, and the plateau cells at η = 0.03 read partially relaxed. I kept the plan's f̄ band, exposed f_max beside f̄ in `str(regime)`, the sweep table, and the report, and wrote the finding into report Section 5.8. My recommendation is a two-sided band: "backprop-like" requires f_max < 0.1, "near PC equilibrium" requires f̄ > 0.9. That is your call. **Smaller deviations from the plan.** `Regime` carries `eta` and `steps`. The sweep letter gains an `r` suffix for the reversal flag. `--power_iters` became `--lanczos_iters`. **Pending, and why** - The `IterContext` PR: two probe tests in `TestWithTrain` skip until `fabricpc.training.IterContext` exists, and the demo's `--track_regime` path is untested. The probe's row logic is tested directly against `make_train_step` contexts. - The four GPU control runs (deliverable I) and the slots that depend on them: the demo docstring's control-run table, report Section 5.9, and the growth phases. The plan's commit split still applies: steps 1 and 2 (the spectrum module and oracle edits; the `Regime` migration) are ready now, and the probe, demo tracking, and script reader belong after the rebase. Consolidated review iterations records to the branch design file epc_inference_solver.md
28db7c5 to
ab7bbdb
Compare
…net18 test runs and documented analysis.
2 participants
Hi Cedric, please review this draft design document. It's a rough cut and I'm sending it for your early feedback. The PR is a draft and is not going to be merged as-is.
docs/dev_plans_archive/epc_inference_solver.mdplans an ePC solver (arXiv 2505.20137) alongside FabricPC's existing state-based PC, plus a composable inference schedule and a graph-topology scheduler so ePC also accepts cyclic graphs. No code yet; the point is to check the formulation before I build it.Where your read would help most:
The reparameterization mapped onto arbitrary graphs. Per node i, ε_i is the relaxed variable and the latent is derived as z_latent_i := z_mu_i + ε_i in schedule order, with one
jax.value_and_gradover the whole ε pytree. Does the node partition in Formulation miss a case in particular unclamped source nodes (in-degree 0, top-down priors), where no inputs are available to make a projection and I relax z_latent itself instead?Cyclic graphs by unrolling with tied errors. A cycle member visited U times recomputes z_mu from the latest source latents but reuses the same ε. Is that the faithful ePC treatment of a recurrent block, or does each traversal need its own ε?
Composition with state-based PC. The plan runs a few ePC steps, then state-based refinement on the full-graph energy including back edges. Is that refinement pass buying anything ePC on the unrolled graph does not already give?
muPC interaction. Forward input scaling is applied inside the differentiated forward; the per-hop gradient preconditioners are deliberately dropped, on the reasoning that they exist to counter the per-layer signal decay ePC removes. I would like a second opinion on that.
Benchmark design (Component 6): ePC at 10 inference steps against state-based PC at 120 on resnet18/CIFAR-10, as both an equal-epoch and a roughly equal-wall-clock comparison. Is that the comparison you would want to see?
Comments inline on the document are ideal. Corrections on anything I have misread in the paper are very welcome.