Dynamic Cost-Error Dimensions per Batch

[tklosek0]

❓ Questions and Help

Hello,

first of all, I wanted to thank you all for pushing this repo constantly!

I'm currently working with Theseus and wondered if/or how it is possible to work with batch-dependent cost-function dimensions.
I try to explain my issue, by defining:

example_cost = th.AutoDiffCostFunction(
            optim_vars=[opt_var],
            aux_vars=[aux_var],
            err_fn=self.example_err_fn, 
            dim=fixed_dim,
            cost_weight=self.example_weight,
            name="example_cost")
self.obj_func.add(example_cost)

Now assume fixed_dim=64 as an upper bound generally that is known in advance.
I now propagate the training sample with batch_size=32 through the AutoDiffCostFunction to calculate the example_cost.
The output of the AutoDiffCostFunction must be (,fixed_dim).

My use case currently is, that the dimension of the training sample inside the Batch is not constant.
Here for an example of a shuffled Batch:

Sample 0 = [64]
Sample 1 = [8]
Sample 2 = [16]
.
.
Sample 31 = [32]

My real fixed_dim is a function of the sample_dimension. An easy example would be, that the fixed_dim is exactly the dimension of the sample. Right now I have to pad the samples to the largest dimension inside the batch since otherwise I would violate the dimension constraint given by the AutoDiffCostFunction.
I wonder if it would be possible to define Sample wise err_fn-Dimension for a Batch like this:

err_fn_dim for Sample 0 = 64
err_fn_dim for Sample 1 = 8
err_fn_dim for Sample 2 = 16
.
.
err_fn_dim for Sample = 32

This would not only speed up the computation but also give the users another degree of freedom.

Would appreciate hearing from you!

[luisenp]

Hi @tklosek0. Thanks for your interest and kind comments!

This feature is tricky to add, because the dimension of the cost functions is used to put together a batched linear system of equations, and we currently don't have a good way to solve systems with different structure in the same batch. For example, for our dense solvers we need to explicitly compute J^t * J where J is the jacobians matrix of shape (batch_size, sum_cost_dims, sum_var_dims). If the cost function dimension depends on the sample, then this matrix cannot be easily be put together. For sparse solvers the computation is different, but similar problems exist.

We have discussed internally supporting sparsity within cost functions themselves, which might be a way to get around this issue. But this is only in preliminary discussions at this point.

[luisenp]

One potential alternative you could try would be to have separate cost functions for each possible error dimension you anticipate, and "route" samples to the corresponding cost function (making sure you keep correct batch indices). We have an internal mechanism to skip cost function computation if their weights are set to zero for all batch indices. So, if in each batch you are able to set zero weight to several of the possible dimensions, you may be able to save some of the cost of building the linear system (which tends to be a significant part of the total cost).

It's a bit messy, but might be worth a shot.