Unexpected behaviour of indexing?

[BurgerAndreas]

Hi,

Thanks again for this library!

1. Am I right in that n_states / indexing can be used to implement the sparse fixed-point correction of DEQ Optical Flow?

2. If yes, I am confused about the output in this example:

from torchdeq import get_deq

# Settings from `DEQ Optical Flow` paper
args = {
    "n_states": 2,
    "f_max_iter": 24,
}

deq = get_deq(args)

print('deq.indexing: ', deq.indexing)

Output: deq.indexing: [12, 12]
Expected output: [8, 16] (uniformly sample between 0 and 24)

Am I missinterpreting?

[BurgerAndreas]

I realised setting core="indexing" yields the more expected behaviour deq.indexing: [12, 24]

Full example:

from torchdeq import get_deq

# Settings from `DEQ Optical Flow` paper
args = {
    "core": "indexing",
    "n_states": 2,
    "f_max_iter": 24,
}

deq = get_deq(args)

print('deq.indexing: ', deq.indexing)

Question

Can you explain what the usecase (pros and cons) for core="indexing" and core="sliced" are?
From the documentation:

DEQIndexing and DEQSliced build different computational graphs in training but keep the same for test.

For DEQIndexing, it defines a computational graph with tracked gradients by indexing the internal solver states and applying the gradient function to the sampled states. This is equivalent to attaching the gradient function aside the full solver computational graph. The maximum number of DEQ function calls is defined by args.f_max_iter.

For DEQSliced, it slices the full solver steps into several smaller graphs (w/o grad). The gradient function will be applied to the returned state of each subgraph. Then a new fixed point solver will resume from the output of the gradient function. This is equivalent to inserting the gradient function into the full solver computational graph. The maximum number of DEQ function calls is defined by, for example, args.f_max_iter + args.n_states * args.grad.

[sablin39]

TL;DR The indexing built by DEQSliced is local index within each sliced trajectory and the indexing built by DEQIndexing is global index of the single full trajectory.

In DEQIndexing, indices are obtained from _conmute_f_iter(f_max_iter), returning uniformly picked index from a single trajectory of length f_max_iter. Then after at most f_max_iter of forward pass , the _solve_fixed_point method returns the sampled trajectory of given index. Each sample of the trajectory is directly assigned with a grad function like this. So the doc says DEQIndexing is called at most f_max_iter times. There's only one trajectory and the index is global index.

Example computation graph of DEQIndexing with $z^*=\tanh(Wz^*+(Ux+U_{bias}))$, n_states=4, f_max_iter=24, loss=sum((p - y).pow(2).mean() for p in preds) / max(1, len(z_list)), obtaining deq.indexing=[6,12,18,24]

---

In DEQSliced, indices are obtained from _compute_f_iter as a set of local index. This is where your [12,12] was calculated. The forward step becomes multiple sliced trajectory, a loop taking previous result as input. Your the sampled index are actually state_idx*index if you want a global view. Thus the forward process still runs at most n_states*len(indexing)=f_max_iter times.

Example computation graph of DEQSliced with $z^*=\tanh(Wz^*+(Ux+U_{bias}))$, n_states=4, f_max_iter=24, loss=sum((p - y).pow(2).mean() for p in preds) / max(1, len(z_list)), obtaining deq.indexing=[6,6,6,6]

---

Then let's go back to your question. I decompose it to two aspects you may intetested in.

• Computation graph:

  • DEQindexing builds a tall and thin graph that's great for throughput, which is suitable when both target and model are small, like language modeling examples deq-lm.
  • With a short and fat computation graph, DEQSliced is more suitable for compute-heavy tasks like image or video generations like deq-flow, and the depth of graph is bounded by n_state.

• Stability

  • By providing extra supervision within sliced trajectory, DEQSliced offers more chances of reaching equilibrium and backpropagates from multiple subgraph, which is beneficial for composite loss terms, and exploits dataset better during training.
  • However, when it comes to supervison over multiple sample points, the gradient of samples in DEQSliced can only propagates within the sliced trajectory, making it more robust but harder to converge.