At the end of Transformers: Age of Extinction, Optimus flies off with the seed to seek out the Creators.
The question is, where did he get the ability to fly?
During the first movie, he couldn't fly. During the second movie, Jetfire (aka Skyfire, modern name) sacrifices himself to give Optimus a flight upgrade. During the third movie, NEST had engineered flight tech into his trailer, allowing him to use a jetpack later on. During the fourth movie, it's a bit uncertain. One theory suggests that the 'Sword of Judgement' had to do with the sudden flight capability:
Optimus Prime has been able to fly since the end of Transformers (2) : Revenge of The Fallen through, Transformers (3) : Dark of the Moon to, Transformers (4) : Age of Extinction.
He achieved flight power from Jetfire, a Decepticon who had defected to the Autobots. Jetfire sacrifices himself towards the end of Transformers 2 and relinquishes some of his parts and weapons to Optimus. Ratchet and Jolt fit these to Optimus, giving him wings and afterburners for sustainable flight.
During the entire movie of "Transformers 4," Optimus is constantly upgrading. If you noticed, "The Sword of Judgement" upgraded him as soon as he grabbed it. Watch his arms and neck and back throughout the movie. It's constantly becoming better, or "upgraded" as you might say. He develops the ability to fly before he actually does it. Not only that, but he has knowledge of it. Before the final battle in the movie he is fully upgraded, and therefore able to fly. His nanostructure has to evolve in order to adapt to the sword's level of training and physical handling. Or it turns him into a knight and gives him special abilities that he wouldn't have had before grabbing it
Any justification is wrong because there is no rational explanation.
You could suggest that it is an evolution of the abilities he gained since Revenge of the Fallen: that he originally flew using parts from Jetfire, then developed the technology into his Jet Wing form in Dark of the Moon, then incorporated the technology into his feet. But if he previously required an entire jetpack to fly within our atmosphere, how could the technology evolve so quickly into compact space travel rocket boots within the five years he was constantly on the run?
Many people think that the Sword of Judgement gave him "upgrades." Those clicks you saw in his forearm weren't a transformation of any sort; they were his arm mechanisms locking in place to accommodate a heavy task.
Or maybe Cade Yeager added the upgrade when he repaired him? Except of course he didn't. There's no way he could do more than several years with Wheeljack and NEST within a few hours.
At the end of the day, it makes no sense. There are numerous points in the film where this ability would have come in handy, but he suddenly uses it at the end. A lot of abilities in the recent two films are unexplained; both Optimus and Megatron have alternate modes, and Bumblebee can suddenly regenerate. Don't overthink anything.