Improve armem prediction interface

Closed by !244 (merged)