Path Independent Equilibrium Models Can Better Exploit Test-Time Computation
This work addresses the challenge of scalable inference for AI systems, particularly in handling harder or out-of-distribution data, though it is incremental as it builds on existing equilibrium models.
The paper tackles the problem of improving neural network performance on harder problem instances by leveraging test-time computation, showing that equilibrium models with path independence achieve better generalization, with experimental interventions confirming a strong correlation between path independence and accuracy on out-of-distribution samples.
Designing networks capable of attaining better performance with an increased inference budget is important to facilitate generalization to harder problem instances. Recent efforts have shown promising results in this direction by making use of depth-wise recurrent networks. We show that a broad class of architectures named equilibrium models display strong upwards generalization, and find that stronger performance on harder examples (which require more iterations of inference to get correct) strongly correlates with the path independence of the system -- its tendency to converge to the same steady-state behaviour regardless of initialization, given enough computation. Experimental interventions made to promote path independence result in improved generalization on harder problem instances, while those that penalize it degrade this ability. Path independence analyses are also useful on a per-example basis: for equilibrium models that have good in-distribution performance, path independence on out-of-distribution samples strongly correlates with accuracy. Our results help explain why equilibrium models are capable of strong upwards generalization and motivates future work that harnesses path independence as a general modelling principle to facilitate scalable test-time usage.