Exit gates usage at inference timeHow does the exit gate reduce inference compute?
I read the paper and inspected the provided inference code, but I could not determine how the exit gate currently reduces actual inference compute or latency.
As far as I understand, the implementation computes all four recurrent steps for every generated token and only afterward selects the hidden state/logits corresponding to the predicted exit step. In that case, the gate performs predictive early exit, but not computational early exit: all four loops have already been executed.
True token-wise early stopping also appears to create a KV-cache issue. For example, if token n exits after loop 2 but token n+1 continues to loop 4, then loop-4 attention for token n+1 requires loop-4 keys and values for token n, which were never computed.
There are several possible approaches to this problem. For example, Confident Adaptive Language Modeling discusses early exit in autoregressive Transformers:
https://arxiv.org/pdf/2207.07061
Could you clarify the intended inference procedure?
In particular:
- Is there a separate implementation that actually stops computation once the exit threshold is reached?
- If different tokens exit at different recurrent steps, how are the missing deeper-loop KV-cache entries handled?
It would be very helpful if you could provide the exact inference algorithm or point to the relevant implementation.