In their work on code reviews Lo Heander identifies two opposite forces. We need understanding, seeing the code as a part of the system, and that requires a lot of context. But we also need clarity, being able to separate the part of the code that is under review, and that is better done with less context.
Unit tests define the collaborators and the purpose of a piece of code. When reading code with unit tests, we can read the tests first, and focus on how the unit interacts with the system. The review of the computational code, like the body of a function, can then focus on the implementation.
Putting this in the framework created by Heander, we solve the conflict between clarity and understanding by separation of concerns. Each test provides context for purpose and interaction for that unit. At that time we are not interested in the implementation, just the interfaces, and can remove that part of the cognitive noise. When we have enough understanding of the purpose of the code we can move to the implementation. That should then be a pretty small unit, with well defined collaborators, that we are already familiar with from the test. We have the external context in place and can focus on implementation details.
Separating concerns like that lightens the cognitive load, and might be a part of solving the conflict between clarity and understanding.
Test driving you code in small units, does not only help with cognitive load in the moment, but also helps everyone that will read that piece of code in the future.