Loss Landscapes of LLMs: The Map Beneath Gradient Descent

Shrijith Venkatramana, developer of the AI code review tool LiveReview, explains the concept of loss landscapes in large language models, describing how gradient descent navigates high-dimensional parameter spaces and why the geometry of trained LLMs is more complex than a simple bowl-shaped minimum. The post references research by Ian Goodfellow and colleagues showing that neural network optimization paths may lack the expected obstacles, offering insights into learning rates, initialization, and model merging.

Hello, I'm Shrijith Venkatramana, and I'm building LiveReview — a blast-radius aware AI code review built for your business-critical systems. Star us to help devs discover the project, give it a try, and share your feedback to help improve the product. There is a strange fact about training a large language model: A model with hundreds of billions of parameters is trained by repeatedly nudging a point in an unimaginably high-dimensional space downhill. That sentence sounds almost absurd. Imagine a landscape where every coordinate is a model weight. With 70 billion parameters, your "position" is a vector with 70 billion coordinates. The training objective assigns one scalar value to that position. Gradient descent looks at the local slope and says: Move this way. Then it does it again. And again. And again. The resulting object is the loss landscape . For developers, loss landscapes are more than a mathematical curiosity. They provide a useful mental model for understanding why learning rates explode, why initialization matters, why some architectures train dramatically better than others, why independently trained models can sometimes be merged or connected, and why the geometry of a trained LLM is much stranger than the familiar picture of a ball rolling into a single bowl. The most interesting part is that the naive picture of "find the lowest valley" is increasingly misleading. Start with an ordinary function: y = x - 3 ^2 Plot it and you get a bowl. The minimum is at: x = 3 loss = 0 Now imagine two parameters: L w1, w2 You can plot this as a 3D surface. Every point w1, w2 corresponds to one model, and its height corresponds to the loss. Neural networks simply take this idea to an absurd scale. For a model with parameters theta = theta 1, theta 2, ..., theta N the training objective is: L theta where N might be billions. So the actual landscape has billions of dimensions. You cannot draw it. But the mathematical object is perfectly well-defined. For an autoregressive language model, a simplified training loss is cross-entropy: L theta = - 1/T sum t log p theta x t | x <t The model predicts the next token, we compare that distribution with the actual next token, and average the negative log probabilities. Training asks us to solve approximately: min theta L theta The remarkable engineering achievement of modern deep learning is that this apparently ridiculous optimization problem is actually tractable. And that realization was itself historically important. In 2014, Ian Goodfellow, Oriol Vinyals, and Andrew Saxe investigated the optimization behavior of neural networks and found something contrary to the prevailing intuition: for several networks they examined, the loss along a straight path from initialization toward the trained solution did not exhibit the giant obstacles that one might expect from a highly non-convex function. Their observation helped motivate a different view of neural-network optimization: perhaps these landscapes are difficult in high-dimensional ways, but not necessarily because SGD is constantly trapped in horrible local minima. Google Research https://research.google/pubs/qualitatively-characterizing-neural-network-optimization-problems/?utm source=chatgpt.com That distinction becomes crucial once we get to LLMs. Suppose you're standing somewhere in the landscape. The gradient is: grad L theta It points in the direction of steepest increase in loss. So gradient descent takes: theta new = theta - eta grad L theta where eta is the learning rate. The simplest mental model is skiing downhill. But there is an important subtlety. The gradient tells you about local slope , not the shape of the entire mountain range. Consider: / \ / \ Two places could have exactly the same gradient while having radically different curvature around them. That leads naturally to the Hessian: H = d^2 L / d theta^2 Conceptually, the Hessian tells you how the slope itself changes. In one dimension: L x = x^2 has: dL/dx = 2x d2L/dx2 = 2 The curvature is positive everywhere. For a neural network, the Hessian is an enormous matrix: N x N for N parameters. For a 70-billion-parameter model, explicitly constructing it would be laughably impractical. Instead, practitioners often reason about its eigenvalues or quantities related to them. Very roughly: php large positive eigenvalue - steep direction small eigenvalue - flat direction negative eigenvalue - locally downhill in some direction This already gives us a much better picture of a trained model. It isn't merely sitting at a point. It is sitting somewhere inside a complicated geometry containing directions that are extraordinarily stiff and directions that barely matter. There is a famous old intuition about non-convex optimization: There must be countless terrible local minima, and SGD somehow has to avoid them. Modern neural-network research complicated this picture. Goodfellow, Vinyals, and Saxe found surprisingly unobstructed paths between initialization and solutions for networks they studied. Later work became even more striking. In 2018, Felix Draxler and collaborators and, independently, Timur Garipov and collaborators showed that different independently trained neural networks could be connected through low-loss curves in parameter space. Rather than finding isolated little valleys separated by huge mountains, one could often find a continuous path between solutions that stayed at low loss. Proceedings of Machine Learning Research https://proceedings.mlr.press/v80/draxler18a.html?utm source=chatgpt.com Imagine this: traditional intuition: \ / \ / \ / minimum minimum ^ high barrier ^ versus: actual high-dimensional geometry: / \ / \ A B There can be many distinct parameter vectors that all implement excellent solutions, with relatively easy paths connecting them. This matters enormously for language models. Two independently trained models can have completely different parameter vectors: theta A = theta B while implementing broadly similar functions. And there is another complication: neural-network parameterizations contain huge amounts of redundancy. For example, hidden units can sometimes be permuted without changing the function represented by the network. Rescaling symmetries and other parameterization effects create additional equivalent or near-equivalent representations. So the question "Where is the optimum?" may be much less meaningful than: "What does the region of good solutions look like?" That is a much more interesting question. Suppose near a trained solution theta , we perturb the parameters by a small vector delta . A second-order Taylor approximation gives: L theta + delta ≈ L theta + grad L theta ^T delta + 1/2 delta^T H delta At a well-trained solution: grad L theta ≈ 0 so approximately: L theta + delta ≈ L theta + 1/2 delta^T H delta Now imagine diagonalizing the Hessian. Then the loss increase can approximately be thought of as: Delta L ≈ 1/2 sum i lambda i delta i^2 where lambda i is the curvature along direction i . Consider three directions: lambda 1 = 1000 lambda 2 = 1 lambda 3 = 0.000001 Move by the same amount in each direction. The first direction produces a huge loss change. The second produces a moderate change. The third essentially does nothing. This is the intuition behind flat directions . A billion-dimensional model can therefore have a tiny collection of very sensitive directions embedded inside an enormous space of comparatively forgiving directions. That is one reason the parameter count alone tells us almost nothing about how difficult optimization is. Suppose a model has 10^11 parameters. Even if only one part in a million corresponded to strongly curved directions, that would still be: 10^11 / 10^6 = 10^5 or roughly 100,000 highly sensitive dimensions. And that leaves roughly 99,999,900,000 other directions. This is why "the model has billions of parameters" does not imply that optimization is equivalently difficult in billions of independent ways. The geometry is highly anisotropic. Learning rate schedules suddenly become much less mysterious when viewed geometrically. Consider the simplest quadratic: L x = 1/2 lambda x^2 Gradient descent gives: x new = x - eta lambda x or: x new = 1 - eta lambda x For this to converge rather than explode, roughly: |1 - eta lambda| < 1 which implies: 0 < eta < 2/lambda So the maximum stable learning rate depends on curvature. Now replace the single lambda with the largest Hessian eigenvalue: lambda max and you get the rough intuition: eta must be small enough for the stiffest direction This explains a frustrating phenomenon engineers routinely encounter. You can have a model where most directions are beautifully flat, yet one pathological direction is enormously steep. The optimizer cannot simply say: "Most of the landscape is flat, so let's take huge steps." The steep direction gets to veto that decision. This is one reason optimization systems spend so much effort on learning-rate schedules, warmup, normalization, optimizer state, gradient clipping, and parameterization. They are, in various ways, attempts to make movement through the landscape numerically manageable. Suppose early training contains badly scaled gradients. Jumping immediately to the final learning rate can move the parameters an enormous distance through the landscape before the model has settled into a useful region. Warmup effectively says: start cautiously ↓ observe the geometry through gradients ↓ increase step size It is not literally measuring the Hessian at every step, but geometrically it is doing something compatible with the idea that optimization dynamics change dramatically during training. You will often hear: Flat minima generalize better. There is a real phenomenon behind this statement, but the slogan is too simplistic. Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein popularized practical visualization techniques for neural-network loss landscapes. Their 2018 work showed how architecture and optimization choices affect the observed geometry and introduced filter normalization to make visual comparisons more meaningful. ML Anthology https://mlanthology.org/neurips/2018/li2018neurips-visualizing/?utm source=chatgpt.com The core insight is intuitive. Suppose two solutions have identical training loss: Solution A: steep bowl Solution B: broad basin A small parameter perturbation may barely affect B but substantially hurt A. That sounds like B should be more robust. But there is a serious technical wrinkle: sharpness depends on parameterization and scale. Suppose we multiply one layer's weights by 10 and compensate by dividing another layer's weights by 10. The represented function can remain essentially unchanged while the raw parameter-space curvature changes. So saying: "This minimum has Hessian eigenvalue 500 and that one has eigenvalue 100" does not automatically tell you that the first function is less robust. You have to specify the geometry being measured. This is an important general lesson: Parameter space is not function space. Two parameter vectors that look wildly different can implement similar functions. Two parameter vectors that are close in Euclidean distance can sometimes implement meaningfully different functions. That distinction is particularly important for LLMs, because developers increasingly do operations directly on weights: fine-tuning LoRA weight interpolation model merging checkpoint averaging distillation continual pretraining All of these interact with parameter-space geometry. Now we can translate the geometry back into everyday LLM work. Suppose training suddenly does this: step 1000 loss = 3.8 step 1001 loss = 4.0 step 1002 loss = 5.7 step 1003 loss = NaN One useful interpretation is that optimization has entered a region where the chosen step size is incompatible with the local geometry. The cause could involve: learning rate gradient scale numerical precision activation statistics optimizer state data distribution normalization but the geometric symptom is simple: step too large relative to local curvature SGD does not observe the exact population gradient. It observes an estimate: g hat = g + noise A larger batch generally reduces the variance of this estimator. That means the optimizer experiences a different effective dynamical system. One way to visualize it: small batch: noisy path /\/\ /\/\ / large batch: smooth path /──────────── This noise is not necessarily undesirable. It can affect which parts of the landscape the optimizer visits and which solutions it eventually reaches. This is one reason optimization hyperparameters are not merely numerical plumbing. They can change the trajectory through the landscape itself. Residual connections are a particularly revealing example. A plain deep network can require every layer to learn a useful transformation. A residual block can instead learn approximately: f x = x + delta x where delta x is a correction. The identity path creates a much easier route for information and gradients. Li et al.'s loss-landscape experiments helped visualize the broader phenomenon: architectural choices can alter the geometry of optimization, not merely the number of parameters or FLOPs. ML Anthology https://mlanthology.org/neurips/2018/li2018neurips-visualizing/?utm source=chatgpt.com This is one reason the history of deep learning is partly the history of making the optimization landscape easier to traverse. Imagine training a frontier model costs: $50M and an optimization improvement reduces the required number of training steps by 10%. Very roughly: $50M 0.10 = $5M That is before considering engineering capacity, cluster availability, electricity, scheduling, opportunity cost, and failed runs. A seemingly abstract improvement to optimization geometry can therefore be worth millions of dollars. The economics of frontier training makes the landscape a systems problem. A better optimizer, initialization, normalization scheme, architecture, or learning-rate schedule is effectively a way of making the billion-dimensional terrain cheaper to cross. The most useful mental shift is this: Training an LLM is probably not best understood as searching for one magical global minimum. The picture is closer to finding a good region in an enormous, structured space of solutions. You can imagine: high loss /\ / \ / \ / \ / \ A===============================B low-loss region The "equals" line is not necessarily a straight interpolation. Garipov et al. demonstrated that low-loss curves could connect solutions that looked separated by barriers under naive linear interpolation. Draxler et al. likewise found essentially barrier-free paths between independently trained solutions in several settings. NeurIPS Papers https://papers.nips.cc/paper/2018/hash/be3087e74e9100d4bc4c6268cdbe8456-Abstract.html?utm source=chatgpt.com That gives us a striking reinterpretation of several modern LLM techniques. When you fine-tune a base model, you're moving through the landscape. When you train two different fine-tunes, you're landing at different places in the landscape. When you merge models, you're betting that useful solutions occupy sufficiently compatible regions of parameter space. When you average checkpoints, you're betting that nearby points lie within a useful basin. When you change the optimizer, you're changing the dynamics by which you travel. And when you scale the model, you are not merely adding more capacity. You are changing the dimensionality and geometry of the object being optimized. This is why loss landscapes are such a useful concept for developers: they connect seemingly unrelated engineering decisions into one underlying question. What kind of terrain are we asking gradient descent to navigate? The most important thing to take away is not a particular Hessian formula or visualization technique. It is the mental model. A neural network is a point in a gigantic parameter space. The loss function turns that space into a landscape. The gradient tells you which way is locally uphill. The optimizer chooses how aggressively to move. The Hessian describes local curvature. Architecture changes the terrain. Batch size changes the noise in your navigation. Learning rate determines whether your steps are cautious exploration or giant leaps. And surprisingly, good solutions may form broad, connected regions rather than isolated "perfect minima." Once you start seeing LLM training this way, a lot of seemingly arbitrary choices become geometrically legible. The next time a training run diverges, a fine-tune behaves unexpectedly, or two checkpoints refuse to combine nicely, ask yourself: What does the landscape around this model probably look like? And perhaps the more interesting frontier question is: As models scale from billions to trillions of parameters, what properties of their loss landscapes actually change—and which ones remain remarkably invariant? Your team's attention is limited, and the deluge of AI-generated code is making it harder to keep production stable while also shipping at high velocity. I'm building LiveReview , a blast-radius aware AI code review built for your business-critical systems. Instead of presenting every diff with equal emphasis, LiveReview scores each change by blast radius — how far its impact reaches through your call graph — so you can focus attention where it actually matters. Spend code review effort where business risk is highest — not spread evenly across every diff. Try LiveReview on your codebase: