Enable more efficient batching over multiple (co)tangent vectors

[jpbrodrick89]

Summary

Provide the ability to provide multiple (co)tangent vectors and efficiently calculate the corresponding jvp's/vjp's with a single HTTP request to the Tesseract.

Why is this needed?

It is very common to require a jvp or vjp with respect to multiple (co)tangent vectors. For example, when calculating a Jacobian to perform a Newton solve. While the jacobian endpoint already exists, this is never called by our open source pipeline builder tesseract-jax; instead either vector_jacobian_product (jax.jacobian = jax.jacrev) or jacobian_vector_product (jax.jacfwd) is called multiple times sequentially. This is inefficient for multiple reasons:

• Parallelisation/SIMD is not possible as the calls are all made synchronously.
• We suffer any latency due to HTTP requests, etc. multiple times
• We lose the ability of taking advantage of any tricks such as would be provided jax.vmap to avoid unnecessary re-computation of primals on the forward pass.

While allowing async calls with pmap and num_workers > 1 should be possible and fix the first two issues and re-writing the batching rule to instead call jax.jacobian and calculate a matrix-matrix product could work, enabling jacobian_vector_product and vector_jacobian_product would be more efficient in most cases (e.g. if no. of tangent vectors is > 1 and <<n).

Usage example

My proposal would be to allow jacobian_vector_product and vector_jacobian_product to allow (at least one) additional dimension on the (co)tangent vector inputs. We may decide to insist on the additional dimensions being added to all elements, but if not they could be expanded to. The Tesseract would look at the shape of the (co)tangent vector and work out whether an additional dimensions has been added and if so then call a batched version of itself.

[dionhaefner]

Which one is the real missing behavior here – batching over JVP and VJP, or invoking jacobian from existing clients to compute full Jacobians?

[jpbrodrick89]

I would really argue the former, it would cover all bases and perhaps even make the latter redundant. If we instead had a batching rule that just read off from Jacobian this would be inefficient in almost any instance that didn't require exactly the entire dense Jacobian matrix, particularly when one can take advantage of sparse Jacobian techniques such as the colouring method used for tridiagonal Jacobians in my recent lineax PR Our Jacobian endpoint is required (by validation) to always instantiate an entire matrix (for the given jac in/outputs) and load it into memory even if the matrix is sparse.

[johnbcoughlin]

+1 to Jon. If I understand things correctly, it's not really up to us what all existing clients do -- this includes jax.jacobian when the Tesseract forms just one component of the function whose jacobian we want.

[dionhaefner]

Can we come up with a concrete example (MVE) where this feature would be required / a massive speed-up? I agree with the high-level reasoning here but without a concrete case it's easy to get fooled, especially when arguing with performance (that is better measured than derived by making a lot of assumptions).

[jpbrodrick89]

Working assumption

We want to support a diverse suite of auto-diff powered workflows that goes beyond gradient-based optimization and sensitivity calculation, including but not limited to root finding and least squares. Such workflows should be of comparable performance to Tesseract-free approaches for at least tens of parameters.

Potential solutions

I take the assumption that we can only pick one of the below options and stick with it, I believe leaving this particular choice up to the user (e.g. vmap_method="from_jacobian") introduces unnecessary complexity/a potential footgun.

1. Current approach: call jacobian_vector_product/vector_jacobian_product sequentially or in parallel. I will not distinguish between these two—despite the latter taking less wall time it will still use a similar amount of resources and require a very large num_workers to reap full benefit.
2. Always call jacobian endpoint when batching jvp's/vjp's and then perform the necessary matrix multiplication.
3. Allow jacobian_vector_product/vector_jacobian_product to natively support batched (co)tangent vectors.

Example: Poor performance of current approach (1) for least squares fitting

I will use an example from my discourse forum tutorial on Gaussian radial basis function fitting comparing the most effective least squares approach using the Tesseract jacobian endpoint directly or using tesseract-jax with jax.jacfwd.

Edit `optimization_routine.py` by replacing everything after `print("Launching Tesseract")`

from scipy.optimize import least_squares

from tesseract import Tesseract
from tesseract_jax import apply_tesseract

# Launch Tesseract
with Tesseract.from_image(image="rbf_fitting") as tess:
    def y_pred_wrapped(weights):
        inputs["weights"] = weights
        return tess.apply(inputs)["y_pred"] - y_target

    def jac_wrapped(weights):
        inputs["weights"] = weights
        return (
            tess.jacobian(inputs, ["weights"], ["y_pred"])["y_pred"]["weights"]
        )

    # Call once to compile the function
    y_pred_wrapped(weights_0)
    jac_wrapped(weights_0)

    print(
        "Starting Least Squares optimization process using Tesseract `jacobian` endpoint."
    )
    start = timer()
    ls = least_squares(
        y_pred_wrapped, weights_0, method="lm", jac=jac_wrapped, x_scale="jac"
    )
    end = timer()
    print(
        f"Completed in {end - start:.3f} seconds, with {ls.nfev} calls to apply "
        f"and {ls.njev} calls to jacobian."
    )

    # tesseract-jax doesn't like primitive float
    inputs["scale"] = jax.numpy.array(0.05)

    @jax.jit
    def y_pred_wrapped(weights):
        inputs["weights"] = weights
        return apply_tesseract(tess, inputs)["y_pred"] - y_target
    
    @jax.jit
    def jac_wrapped(weights):
        return jax.jacobian(y_pred_wrapped)(weights)

    # Call once to compile the function
    y_pred_wrapped(weights_0)
    jac_wrapped(weights_0)
    
    print(
        "Starting Least Squares optimization process using `tesseract-jax` with `jax.jacfwd`."
    )
    start = timer()
    ls = least_squares(
        y_pred_wrapped, weights_0, method="lm", jac=jac_wrapped, x_scale="jac"
    )
    end = timer()
    print(
        f"Completed in {end - start:.3f} seconds, with {ls.nfev} calls to apply "
        f"and {ls.njev} calls to jacobian."
    )

This results in the following:

Launching Tesseract
Starting Least Squares optimization process using Tesseract `jacobian` endpoint.
Completed in 0.066 seconds, with 7 calls to apply and 4 calls to jacobian.
Starting Least Squares optimization process using `tesseract-jax` with `jax.jacfwd`.
Completed in 1.068 seconds, with 7 calls to apply and 4 calls to jacobian.

For a very simple Tesseract with a 20x50 Jacobian we observe a 16x slowdown with tesseract-jax. It is assumed but not tested that native batching of jacobian_vector_product should be just as performant as the jacobian endpoint (as this is what jax is doing under the hood).

Issues with always using the jacobian endpoint when batching

Some of these have already been discussed above (unnecessary instantiation and piping of entire Jacobian, inability to benefit from coloring methods which can cause orders of magnitude speed up for large array sizes). I am not going to try and implement a Tesseract-specific example of these just yet (as I'm not sure what you're looking for), but a few more relevant points:

• jax.linearize also requires a batching rule but may only be called with a small number of vectors which is used for Gauss-Newton products in conjugate-gradients method.
• matrix multiplication also comes with added costs that could be treated more efficiently under the hood, e.g. with tracing.

[dionhaefner]

Yea sounds reasonable. Since (co)tangents are simply reduced over the Jacobian we can always guarantee that naive batching works.

To me the biggest question is how to handle this change in a way that is ideally non-breaking. Hence my suggestion to implement "batching rule of vjp == jacobian * cotan", which can be handled entirely client-side. If we want native auto-batching over tangents things get more difficult to implement in a backwards-compatible way.

[jpbrodrick89]

Yes, agree backwards-compatibility is tricky. I think the solution is to make this an opt-in feature in some way. While we could try and check for global variables in tesseract_api.py (e.g. TESSERACT_SUPPORTS_BATCHING=True), currently the only way we use to check whether a tesseract_api.py supports different behaviours is by checking availability of endpoints. Therefore, I propose the addition of two private endpoints:

• jacobian_vector_products/batched_jacobian_vector_product
• vector_jacobian_products/batched_vector_jacobian_product
  The user would then be able to call jacobian_vector_product (no plural) with a batched tangent vector and after concluding this is indeed batched based on its shape during validation (in schema_generation) we call jacobian_vector_proudcts if it is implemented or throw a ValidationError/resort to a default if not. I think defaulting to tan * jacobian if jacobian is implemented or a simple Python for loop/map if not would make most sense.

I appreciate this adds some more complexity to the user, but that is likely an inevitable consequence of allowing such a features. This design would mean the feature can be ignored if not relevant to a users need.

If we can sketch up a design we're both happy with by the end of next week I'd be happy to sign myself up for implementing this as part of some mission during the next quest.

Is it possible to support arbitrary batch dimensions or do we really have to limit ourselves to one batch dimension here?

[dionhaefner]

For posterity, here are a few things we discussed offline:

• This feature is easier to add when there's a concrete need for it – so if anyone is waiting for this functionality to land, please let us know!
• We will likely want to implement this in a backwards-compatible way, e.g. via dedicated batch endpoints or a flag power-users can set to signal that their implementation supports batching.
• We currently don't enforce shapes on tangents / cotangents, so technically this is already implementable by hand. This is more undefined behavior than a feature, but could be useful to build a prototype and demonstrate that this functionality is desirable in real-world use cases.