There are alternatives to "undo", which are probably better suited to "undo last word" without making "what will undo undo" more mentally complicated. Having undo reverse the last edit, whether that's a word or a few words or adding an argument definition to a function, or adding an "if" block -- makes it very easy to keep in your head. "I want to undo that last thing I did".
The alternatives are:
- Control-W (delete the last word, stay in insert mode)
- "W" to move back a word (or "B" to move back a big word). So: <Esc>BC
There's apparently also a recipe for setting up vim to insert undo points whenever you hit space, so you could also enable that for the files you want to do that on.