Generates eight 3D Matplotlib animations visualizing the geometric transformations of Layer Normalization and RMSNorm. Illustrates mean centering (plane projection), variance normalization (hypersphere projection), affine shifts (gamma/beta), epsilon regularization, and RMSNorm's "directional collapse" failure mode in unstable networks.

RmsNormGeometry.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation, PillowWriter
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import sys

# ==========================================
# 1. SHARED DATA GENERATION (Done Once)
# ==========================================
print("Generating dataset...")
np.random.seed(42)
n_dims = 3

# Create two distinct clusters of points to demonstrate Epsilon's effect
n_large = 10
n_small = 10
n_points = n_large + n_small
radius = np.sqrt(n_dims)

# A. Original points
# Large variance points (Spread significantly increased so sigma >> epsilon)
pts_large = np.random.rand(n_large, n_dims) * 25 - 12.5
# Small variance points (Tightly clustered)
pts_small = np.random.rand(n_small, n_dims) * 1.5 + np.array([4, 4, 4])
raw_points = np.vstack([pts_large, pts_small])

# B. Phase 1: Mean Centering (Projection to Plane x+y+z=0)
row_means = np.mean(raw_points, axis=1, keepdims=True)
centered_points = raw_points - row_means

# C. Phase 2: Standardization (Projection to Sphere of radius sqrt(n))
row_stds = np.std(raw_points, axis=1, keepdims=True)
standardized_points = centered_points / row_stds

# D. RMSNorm points
rms_vals = np.sqrt(np.mean(raw_points**2, axis=1, keepdims=True))
rms_points = raw_points / rms_vals

# ── Case 2 specific data (independent of pts_small used in epsilon animation) ──
# Each point gets its own scalar mean drawn uniformly from [3,8] or [-8,-3],
# with small per-component Gaussian noise (std=0.3) so mu/sigma ~ 10-27 for every
# point. This makes the raw positions visibly spread across the space while still
# deep in the collapse regime — the diversity of starting positions makes the
# convergence to two antipodal poles the more striking.
np.random.seed(7)
n_pos = 8
n_neg = 8
std_scale = 0.3  # small std keeps mu/sigma >> 1 for all points

means_pos = np.random.uniform(3.0, 8.0, n_pos)    # diverse positive scalar means
means_neg = np.random.uniform(-8.0, -3.0, n_neg)   # diverse negative scalar means

# Each row: [mu_i + noise, mu_i + noise, mu_i + noise] — components share a mean
# but differ slightly, so std across components is ~std_scale, not zero
pts_pos = np.column_stack([means_pos + np.random.randn(n_pos) * std_scale,
                            means_pos + np.random.randn(n_pos) * std_scale,
                            means_pos + np.random.randn(n_pos) * std_scale])
pts_neg = np.column_stack([means_neg + np.random.randn(n_neg) * std_scale,
                            means_neg + np.random.randn(n_neg) * std_scale,
                            means_neg + np.random.randn(n_neg) * std_scale])

rms_pos = pts_pos / np.sqrt(np.mean(pts_pos**2, axis=1, keepdims=True))
rms_neg = pts_neg / np.sqrt(np.mean(pts_neg**2, axis=1, keepdims=True))

ln_pos = (pts_pos - pts_pos.mean(axis=1, keepdims=True)) / pts_pos.std(axis=1, keepdims=True)
ln_neg = (pts_neg - pts_neg.mean(axis=1, keepdims=True)) / pts_neg.std(axis=1, keepdims=True)

# Pole axis endpoints (scaled to radius for drawing)
pole_axis = radius * np.array([1, 1, 1]) / np.sqrt(3)

# Colors: warm (orange) for positive mean, cool (slate blue) for negative mean
COLOR_POS = 'darkorange'
COLOR_NEG = 'mediumslateblue'

# Shared animation parameters
total_frames = 120
fps = 20

# Reusable function for wireframe sphere
u, v = np.mgrid[0:2*np.pi:30j, 0:np.pi:15j]
xs = radius * np.cos(u) * np.sin(v)
ys = radius * np.sin(u) * np.sin(v)
zs = radius * np.cos(v)

# Reusable LayerNorm Equator
v1 = np.array([1, -1, 0])
v1 = v1 / np.linalg.norm(v1)
n = np.array([1, 1, 1])
n = n / np.linalg.norm(n)
v2 = np.cross(n, v1)
theta_circle = np.linspace(0, 2 * np.pi, 100)
circle_pts = np.array([radius * np.cos(t) * v1 + radius * np.sin(t) * v2 for t in theta_circle])

# ==========================================
# 2. ANIMATION 1: MEAN CENTERING
# ==========================================
print("\nPreparing Animation 1: Mean Centering...")
fig1 = plt.figure(figsize=(18, 8), dpi=100)
ax1a = fig1.add_subplot(121, projection='3d')
ax1b = fig1.add_subplot(122, projection='3d')

bound1_anim = np.max(np.abs(raw_points)) + 1
# Since we are adding raw_points to the right subplot, we use bound1_anim for both
bound1_final = bound1_anim

# Setup Axes
for ax, bound, title in zip([ax1a, ax1b], [bound1_anim, bound1_final],
                            ["Animation: Projection to Orthogonal Complement", "Final View: Centered Data"]):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.1, color='dodgerblue', zorder=0)

fig1.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
# Changed view to make the projection more visible
ax1a.view_init(elev=20, azim=150)

# Left Subplot: Animation Elements
t = np.linspace(-bound1_anim, bound1_anim, 50)
ax1a.plot(t, t, t, color='purple', linewidth=2, linestyle='--', zorder=1)
ax1a.scatter(raw_points[:,0], raw_points[:,1], raw_points[:,2], c='royalblue', s=50, edgecolors='black')

scat1a = ax1a.scatter(raw_points[:,0], raw_points[:,1], raw_points[:,2], c='crimson', s=50, edgecolors='black')
paths1 = [ax1a.plot([raw_points[i,0], raw_points[i,0]],
                    [raw_points[i,1], raw_points[i,1]],
                    [raw_points[i,2], raw_points[i,2]], color='black', linestyle=':', alpha=0.6)[0] for i in range(n_points)]

# Right Subplot: Static Final Data with full details
ax1b.plot(t, t, t, color='purple', linewidth=2, linestyle='--', zorder=1)
ax1b.scatter(raw_points[:,0], raw_points[:,1], raw_points[:,2], c='royalblue', s=50, edgecolors='black')
ax1b.scatter(centered_points[:,0], centered_points[:,1], centered_points[:,2], c='crimson', s=50, edgecolors='black')
for i in range(n_points):
    ax1b.plot([raw_points[i,0], centered_points[i,0]],
              [raw_points[i,1], centered_points[i,1]],
              [raw_points[i,2], centered_points[i,2]], color='black', linestyle=':', alpha=0.6)

leg1 = [mlines.Line2D([], [], color='royalblue', marker='o', markeredgecolor='black', linestyle='None', label='Raw Data'),
        mlines.Line2D([], [], color='crimson', marker='o', markeredgecolor='black', linestyle='None', label='Centered Data'),
        mlines.Line2D([], [], color='purple', linestyle='--', linewidth=2, label='Mean Axis (Vector of 1s)'),
        mpatches.Patch(color='dodgerblue', alpha=0.3, label='Orthogonal Plane')]
ax1a.legend(handles=leg1, loc='upper left')

def update1(frame):
    s = np.sin((frame / total_frames) * np.pi/2)
    cur_pts = raw_points + s * (centered_points - raw_points)

    scat1a._offsets3d = (cur_pts[:,0], cur_pts[:,1], cur_pts[:,2])
    for i in range(n_points):
        paths1[i].set_data([raw_points[i,0], cur_pts[i,0]], [raw_points[i,1], cur_pts[i,1]])
        paths1[i].set_3d_properties([raw_points[i,2], cur_pts[i,2]])

    ax1b.view_init(elev=20 + 10 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 1 Frame {frame+1}/{total_frames}...')
    return fig1,

ani1 = FuncAnimation(fig1, update1, frames=total_frames, interval=50, blit=False)
ani1.save('01_mean_centering.gif', writer=PillowWriter(fps=fps))
plt.close(fig1)

# ==========================================
# 3. ANIMATION 2: STANDARDIZATION
# ==========================================
print("\n\nPreparing Animation 2: Standardization...")
fig2 = plt.figure(figsize=(18, 8), dpi=100)
ax2a = fig2.add_subplot(121, projection='3d')
ax2b = fig2.add_subplot(122, projection='3d')

bound2_anim = 5.0
bound2_final = 2.5 # Tighter zoom for the sphere

for ax, bound, title in zip([ax2a, ax2b], [bound2_anim, bound2_final],
                            [f"Animation: Projection to Hypersphere (r=√{n_dims})", "Final View: Standardized Data"]):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.1, color='dodgerblue', zorder=0)
    ax.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
    ax.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)

fig2.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
ax2a.view_init(elev=20, azim=45)

# Left Subplot (Animation)
ax2a.scatter(centered_points[:,0], centered_points[:,1], centered_points[:,2], c='forestgreen', s=50, edgecolors='black')
scat2a = ax2a.scatter(centered_points[:,0], centered_points[:,1], centered_points[:,2], c='gold', s=70, edgecolors='black', zorder=6)
paths2 = [ax2a.plot([centered_points[i,0], centered_points[i,0]],
                    [centered_points[i,1], centered_points[i,1]],
                    [centered_points[i,2], centered_points[i,2]], color='black', linestyle=':', alpha=0.6)[0] for i in range(n_points)]

# Right Subplot (Final)
ax2b.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='gold', s=70, edgecolors='black', zorder=6)

leg2 = [mlines.Line2D([], [], color='forestgreen', marker='o', linestyle='None', label='Centered Data'),
        mlines.Line2D([], [], color='gold', marker='o', linestyle='None', label='Standardized Data'),
        mlines.Line2D([], [], color='gray', alpha=0.5, label='Hypersphere'),
        mlines.Line2D([], [], color='gold', linewidth=3.0, label='LayerNorm Equator')]
ax2a.legend(handles=leg2, loc='upper left')

def update2(frame):
    s = np.sin((frame / total_frames) * np.pi/2)
    cur_pts = centered_points + s * (standardized_points - centered_points)

    scat2a._offsets3d = (cur_pts[:,0], cur_pts[:,1], cur_pts[:,2])
    for i in range(n_points):
        paths2[i].set_data([centered_points[i,0], cur_pts[i,0]], [centered_points[i,1], cur_pts[i,1]])
        paths2[i].set_3d_properties([centered_points[i,2], cur_pts[i,2]])

    ax2b.view_init(elev=15 + 15 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 2 Frame {frame+1}/{total_frames}...')
    return fig2,

ani2 = FuncAnimation(fig2, update2, frames=total_frames, interval=50, blit=False)
ani2.save('02_standardization.gif', writer=PillowWriter(fps=fps))
plt.close(fig2)

# ==========================================
# 4. ANIMATION 3: LAYERNORM PARAMS (GAMMA)
# ==========================================
print("\n\nPreparing Animation 3: LayerNorm Scaling (Gamma)...")
fig3 = plt.figure(figsize=(18, 8), dpi=100)
ax3a = fig3.add_subplot(121, projection='3d')
ax3b = fig3.add_subplot(122, projection='3d')

target_gamma = np.array([2.0, 0.4, 1.5])
bound3_anim = 4.5
bound3_final = 3.5

for ax, bound, title in zip([ax3a, ax3b], [bound3_anim, bound3_final],
                            ["Animation: Scaling ($\gamma$)", "Final View: Ellipsoid"]):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.1, color='dodgerblue', zorder=0)

fig3.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
ax3a.view_init(elev=30, azim=35)

# Left Subplot
ax3a.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
ax3a.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)
ax3a.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='gold', s=50, edgecolors='black', zorder=4)

scat3a = ax3a.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='magenta', s=70, edgecolors='black', zorder=6)
paths3 = [ax3a.plot([standardized_points[i,0], standardized_points[i,0]],
                    [standardized_points[i,1], standardized_points[i,1]],
                    [standardized_points[i,2], standardized_points[i,2]], color='black', linestyle=':', alpha=0.6)[0] for i in range(n_points)]
gamma_wire_a = ax3a.plot_wireframe(xs, ys, zs, color='magenta', alpha=0.2, linewidth=0.5)

# Right Subplot (Static geometry)
g_pts_final = standardized_points * target_gamma

# Add 'Before' state to right subplot
ax3b.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
ax3b.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)
ax3b.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='gold', s=50, edgecolors='black', zorder=4)
for i in range(n_points):
    ax3b.plot([standardized_points[i,0], g_pts_final[i,0]],
              [standardized_points[i,1], g_pts_final[i,1]],
              [standardized_points[i,2], g_pts_final[i,2]], color='black', linestyle=':', alpha=0.4)

# Add 'After' state to right subplot
ax3b.plot_wireframe(xs * target_gamma[0], ys * target_gamma[1], zs * target_gamma[2], color='magenta', alpha=0.2, linewidth=0.5)
ax3b.scatter(g_pts_final[:,0], g_pts_final[:,1], g_pts_final[:,2], c='magenta', s=70, edgecolors='black', zorder=6)

leg3 = [mlines.Line2D([], [], color='gold', marker='o', linestyle='None', label='Standardized Data'),
        mlines.Line2D([], [], color='magenta', marker='o', linestyle='None', label='Scaled Data ($\gamma$)'),
        mlines.Line2D([], [], color='magenta', alpha=0.5, label='Scaled Ellipsoid')]
ax3a.legend(handles=leg3, loc='upper left')

def update3(frame):
    global gamma_wire_a
    s = np.sin((frame / total_frames) * np.pi/2)
    current_gamma = 1.0 + s * (target_gamma - 1.0)
    g_pts = standardized_points * current_gamma

    scat3a._offsets3d = (g_pts[:,0], g_pts[:,1], g_pts[:,2])
    for i in range(n_points):
        paths3[i].set_data([standardized_points[i,0], g_pts[i,0]], [standardized_points[i,1], g_pts[i,1]])
        paths3[i].set_3d_properties([standardized_points[i,2], g_pts[i,2]])

    gamma_wire_a.remove()
    gamma_wire_a = ax3a.plot_wireframe(xs * current_gamma[0], ys * current_gamma[1], zs * current_gamma[2], color='magenta', alpha=0.2, linewidth=0.5)

    ax3b.view_init(elev=20 + 10 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 3 Frame {frame+1}/{total_frames}...')
    return fig3,

ani3 = FuncAnimation(fig3, update3, frames=total_frames, interval=50, blit=False)
ani3.save('03_layernorm_gamma.gif', writer=PillowWriter(fps=fps))
plt.close(fig3)

# ==========================================
# 5. ANIMATION 4: LAYERNORM PARAMS (BETA)
# ==========================================
print("\n\nPreparing Animation 4: LayerNorm Shifting (Beta)...")
fig4 = plt.figure(figsize=(18, 8), dpi=100)
ax4a = fig4.add_subplot(121, projection='3d')
ax4b = fig4.add_subplot(122, projection='3d')

target_beta = np.array([3.0, 2.0, -2.0])
bound4_anim = 5.5

for ax, bound, title in zip([ax4a, ax4b], [bound4_anim, bound4_anim],
                            ["Animation: Shifting ($\\beta$)", "Final View: Shifted Region"]):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.1, color='dodgerblue', zorder=0)

fig4.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
ax4a.view_init(elev=20, azim=150)

# Left Subplot
ax4a.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
ax4a.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)
ax4a.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='gold', s=50, edgecolors='black', zorder=4)

scat4a = ax4a.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='darkorange', s=70, edgecolors='black', zorder=6)
paths4 = [ax4a.plot([standardized_points[i,0], standardized_points[i,0]],
                    [standardized_points[i,1], standardized_points[i,1]],
                    [standardized_points[i,2], standardized_points[i,2]], color='black', linestyle=':', alpha=0.6)[0] for i in range(n_points)]
beta_wire_a = ax4a.plot_wireframe(xs, ys, zs, color='darkorange', alpha=0.2, linewidth=0.5)

# Right Subplot (Static geometry in shifted location)
b_pts_final = standardized_points + target_beta

# Add 'Before' state to right subplot
ax4b.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
ax4b.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)
ax4b.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='gold', s=50, edgecolors='black', zorder=4)
for i in range(n_points):
    ax4b.plot([standardized_points[i,0], b_pts_final[i,0]],
              [standardized_points[i,1], b_pts_final[i,1]],
              [standardized_points[i,2], b_pts_final[i,2]], color='black', linestyle=':', alpha=0.4)

# Add 'After' state to right subplot
ax4b.plot_wireframe(xs + target_beta[0], ys + target_beta[1], zs + target_beta[2], color='darkorange', alpha=0.2, linewidth=0.5)
ax4b.scatter(b_pts_final[:,0], b_pts_final[:,1], b_pts_final[:,2], c='darkorange', s=70, edgecolors='black', zorder=6)

leg4 = [mlines.Line2D([], [], color='gold', marker='o', linestyle='None', label='Standardized Data'),
        mlines.Line2D([], [], color='darkorange', marker='o', linestyle='None', label='Shifted Data ($\\beta$)')]
ax4a.legend(handles=leg4, loc='upper left')

def update4(frame):
    global beta_wire_a
    s = np.sin((frame / total_frames) * np.pi/2)
    current_beta = s * target_beta
    b_pts = standardized_points + current_beta

    scat4a._offsets3d = (b_pts[:,0], b_pts[:,1], b_pts[:,2])
    for i in range(n_points):
        paths4[i].set_data([standardized_points[i,0], b_pts[i,0]], [standardized_points[i,1], b_pts[i,1]])
        paths4[i].set_3d_properties([standardized_points[i,2], b_pts[i,2]])

    beta_wire_a.remove()
    beta_wire_a = ax4a.plot_wireframe(xs + current_beta[0], ys + current_beta[1], zs + current_beta[2], color='darkorange', alpha=0.2, linewidth=0.5)

    ax4b.view_init(elev=20 + 10 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 4 Frame {frame+1}/{total_frames}...')
    return fig4,

ani4 = FuncAnimation(fig4, update4, frames=total_frames, interval=50, blit=False)
ani4.save('04_layernorm_beta.gif', writer=PillowWriter(fps=fps))
plt.close(fig4)

# ==========================================
# 6. ANIMATION 5: EPSILON REGULARIZATION
# ==========================================
print("\n\nPreparing Animation 5: Epsilon Regularization...")
fig5 = plt.figure(figsize=(18, 8), dpi=100)
ax5a = fig5.add_subplot(121, projection='3d')
ax5b = fig5.add_subplot(122, projection='3d')

bound5_anim = 4.0
bound5_final = 2.5
target_epsilon = 1.0
final_eps_pts = centered_points / np.sqrt(row_stds**2 + target_epsilon)

for ax, bound, title in zip([ax5a, ax5b], [bound5_anim, bound5_final],
                            ["Animation: Regularization ($\epsilon$)", "Final View: Regularized Data"]):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.05, color='dodgerblue', zorder=0)
    ax.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
    ax.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)

fig5.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
ax5a.view_init(elev=25, azim=45)

# Left Subplot
ax5a.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='black', marker='x', s=50, zorder=4)
ax5a.scatter(final_eps_pts[:n_large,0], final_eps_pts[:n_large,1], final_eps_pts[:n_large,2], facecolors='none', edgecolors='royalblue', s=100, linewidth=2, zorder=5)
ax5a.scatter(final_eps_pts[n_large:,0], final_eps_pts[n_large:,1], final_eps_pts[n_large:,2], facecolors='none', edgecolors='crimson', s=100, linewidth=2, zorder=5)

scat5a_l = ax5a.scatter(centered_points[:n_large,0], centered_points[:n_large,1], centered_points[:n_large,2], c='royalblue', s=60, edgecolors='black', zorder=6)
scat5a_s = ax5a.scatter(centered_points[n_large:,0], centered_points[n_large:,1], centered_points[n_large:,2], c='crimson', s=60, edgecolors='black', zorder=6)

for i in range(n_points):
    ax5a.plot([centered_points[i,0], standardized_points[i,0]],
              [centered_points[i,1], standardized_points[i,1]],
              [centered_points[i,2], standardized_points[i,2]], color='black', linestyle=':', alpha=0.3, zorder=3)

# Right Subplot (Static Final state)
ax5b.scatter(standardized_points[:,0], standardized_points[:,1], standardized_points[:,2], c='black', marker='x', s=50, zorder=4)
ax5b.scatter(final_eps_pts[:n_large,0], final_eps_pts[:n_large,1], final_eps_pts[:n_large,2], c='royalblue', s=60, edgecolors='black', zorder=6)
ax5b.scatter(final_eps_pts[n_large:,0], final_eps_pts[n_large:,1], final_eps_pts[n_large:,2], c='crimson', s=60, edgecolors='black', zorder=6)

leg5 = [mlines.Line2D([], [], color='black', marker='x', linestyle='None', label='Ideal Target ($\epsilon=0$)'),
        mlines.Line2D([], [], color='royalblue', marker='o', linestyle='None', label='Final: Large Var'),
        mlines.Line2D([], [], color='crimson', marker='o', linestyle='None', label='Final: Small Var')]
ax5a.legend(handles=leg5, loc='upper left')

def update5(frame):
    s = np.sin((frame / total_frames) * np.pi/2)
    cur_pts = centered_points + s * (final_eps_pts - centered_points)

    scat5a_l._offsets3d = (cur_pts[:n_large,0], cur_pts[:n_large,1], cur_pts[:n_large,2])
    scat5a_s._offsets3d = (cur_pts[n_large:,0], cur_pts[n_large:,1], cur_pts[n_large:,2])

    ax5b.view_init(elev=20 + 10 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 5 Frame {frame+1}/{total_frames}...')
    return fig5,

ani5 = FuncAnimation(fig5, update5, frames=total_frames, interval=50, blit=False)
ani5.save('05_layernorm_epsilon.gif', writer=PillowWriter(fps=fps))
plt.close(fig5)

# ==========================================
# 7. ANIMATION 6: RMSNORM (CASE 1: Sigma >> Mean)
# ==========================================
print("\n\nPreparing Animation 6: RMSNorm (Case 1: Sigma >> Mean)...")
fig6 = plt.figure(figsize=(18, 8), dpi=100)
ax6a = fig6.add_subplot(121, projection='3d')
ax6b = fig6.add_subplot(122, projection='3d')

raw_l = raw_points[:n_large]
rms_l = rms_points[:n_large]
ln_l = standardized_points[:n_large]

# Bounding Box Logic: Left subplot needs to fit huge initial spread. Right fits the sphere.
bound6_anim = 4
bound6_final = 2.5

# Setup Left
ax6a.set_xlim([-bound6_anim, bound6_anim]); ax6a.set_ylim([-bound6_anim, bound6_anim]); ax6a.set_zlim([-bound6_anim, bound6_anim])
# Setup Right
ax6b.set_xlim([-bound6_final, bound6_final]); ax6b.set_ylim([-bound6_final, bound6_final]); ax6b.set_zlim([-bound6_final, bound6_final])

for ax, title in zip([ax6a, ax6b], ["Animation: RMSNorm Case 1 ($\sigma \gg \mu$)", "Final View: RMSNorm points"]):
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    # Get bound dynamically to draw the plane
    b = ax.get_xlim()[1]
    xx, yy = np.meshgrid(np.linspace(-b, b, 10), np.linspace(-b, b, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.05, color='dodgerblue', zorder=0)
    ax.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
    ax.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)

fig6.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)
ax6a.view_init(elev=20, azim=140)

# Left Subplot
ax6a.scatter(ln_l[:,0], ln_l[:,1], ln_l[:,2], c='black', marker='x', s=50, zorder=4)
ax6a.scatter(rms_l[:,0], rms_l[:,1], rms_l[:,2], facecolors='none', edgecolors='royalblue', s=100, linewidth=2, zorder=5)
for i in range(n_large):
    ax6a.plot([rms_l[i,0], ln_l[i,0]], [rms_l[i,1], ln_l[i,1]], [rms_l[i,2], ln_l[i,2]], color='gray', linestyle='-', linewidth=1.5, alpha=0.5, zorder=2)

scat6a = ax6a.scatter(raw_l[:,0], raw_l[:,1], raw_l[:,2], c='royalblue', s=60, edgecolors='black', zorder=6)
paths6 = [ax6a.plot([raw_l[i,0], raw_l[i,0]],
                    [raw_l[i,1], raw_l[i,1]],
                    [raw_l[i,2], raw_l[i,2]], color='black', linestyle=':', alpha=0.4, zorder=3)[0] for i in range(n_large)]

# Right Subplot (Static Final)
ax6b.scatter(ln_l[:,0], ln_l[:,1], ln_l[:,2], c='black', marker='x', s=50, zorder=4)
ax6b.scatter(rms_l[:,0], rms_l[:,1], rms_l[:,2], c='royalblue', s=60, edgecolors='black', zorder=6)
for i in range(n_large):
    ax6b.plot([rms_l[i,0], ln_l[i,0]], [rms_l[i,1], ln_l[i,1]], [rms_l[i,2], ln_l[i,2]], color='gray', linestyle='-', linewidth=1.5, alpha=0.5, zorder=2)

leg6 = [mlines.Line2D([], [], color='black', marker='x', linestyle='None', label='Ideal LN Target'),
        mlines.Line2D([], [], color='royalblue', marker='o', linestyle='None', label='RMSNorm Position'),
        mlines.Line2D([], [], color='gray', linestyle='-', linewidth=1.5, label='RMS vs LN Gap')]
ax6a.legend(handles=leg6, loc='upper left')

def update6(frame):
    s = np.sin((frame / total_frames) * np.pi/2)
    cur_pts = raw_l + s * (rms_l - raw_l)

    scat6a._offsets3d = (cur_pts[:,0], cur_pts[:,1], cur_pts[:,2])
    for i in range(n_large):
        paths6[i].set_data([raw_l[i,0], cur_pts[i,0]], [raw_l[i,1], cur_pts[i,1]])
        paths6[i].set_3d_properties([raw_l[i,2], cur_pts[i,2]])

    ax6b.view_init(elev=15 + 5 * np.sin(frame * 2 * np.pi / total_frames), azim=frame * (360 / total_frames))
    sys.stdout.write(f'\rRendering Anim 6 Frame {frame+1}/{total_frames}...')
    return fig6,

ani6 = FuncAnimation(fig6, update6, frames=total_frames, interval=50, blit=False)
ani6.save('06_rmsnorm_case1.gif', writer=PillowWriter(fps=fps))
plt.close(fig6)

# ==========================================
# 8. ANIMATION 7: RMSNORM (CASE 2: sign(mu)*gamma collapse)
# ==========================================
print("\n\nPreparing Animation 7: RMSNorm (Case 2: sign(mu)*gamma collapse)...")
fig7 = plt.figure(figsize=(18, 8), dpi=100)
ax7a = fig7.add_subplot(121, projection='3d')
ax7b = fig7.add_subplot(122, projection='3d')

# Left panel needs to cover the full spread of diverse means (up to ~8.9)
bound7_anim = 10.0
bound7_final = 2.5

for ax, bound, title in zip(
    [ax7a, ax7b],
    [bound7_anim, bound7_final],
    ["Animation: RMSNorm Case 2 — $\mu \gg \sigma$", "Final View: Directional Collapse to Poles"]
):
    ax.set_xlim([-bound, bound]); ax.set_ylim([-bound, bound]); ax.set_zlim([-bound, bound])
    ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
    ax.set_title(title, fontsize=14)
    xx, yy = np.meshgrid(np.linspace(-bound, bound, 10), np.linspace(-bound, bound, 10))
    ax.plot_surface(xx, yy, -xx-yy, alpha=0.05, color='dodgerblue', zorder=0)
    ax.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
    ax.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)
    # Pole axis: the [1,1,1] direction along which collapse occurs
    ax.plot([-pole_axis[0], pole_axis[0]],
            [-pole_axis[1], pole_axis[1]],
            [-pole_axis[2], pole_axis[2]],
            color='dimgray', linewidth=1.5, linestyle='--', alpha=0.7, zorder=1)
    # Mark the two poles explicitly
    ax.scatter(*pole_axis,  marker='*', s=200, c='darkorange',      zorder=7, edgecolors='black')
    ax.scatter(*(-pole_axis), marker='*', s=200, c='mediumslateblue', zorder=7, edgecolors='black')

fig7.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95, wspace=0.0)

# View: [1,1,1] axis roughly in-plane so both poles are clearly visible
ax7a.view_init(elev=80, azim=115)
ax7b.view_init(elev=28, azim=45)

# ── Left Subplot (Animation) ──
# Show LN targets for reference
ax7a.scatter(ln_pos[:,0], ln_pos[:,1], ln_pos[:,2], c='black', marker='x', s=50, zorder=4)
ax7a.scatter(ln_neg[:,0], ln_neg[:,1], ln_neg[:,2], c='black', marker='x', s=50, zorder=4)
# Show final RMSNorm destinations (hollow rings at the poles)
ax7a.scatter(rms_pos[:,0], rms_pos[:,1], rms_pos[:,2],
             facecolors='none', edgecolors=COLOR_POS, s=120, linewidth=2, zorder=5)
ax7a.scatter(rms_neg[:,0], rms_neg[:,1], rms_neg[:,2],
             facecolors='none', edgecolors=COLOR_NEG, s=120, linewidth=2, zorder=5)

# Animated raw points (start position)
scat7a_pos = ax7a.scatter(pts_pos[:,0], pts_pos[:,1], pts_pos[:,2],
                           c=COLOR_POS, s=70, edgecolors='black', zorder=6)
scat7a_neg = ax7a.scatter(pts_neg[:,0], pts_neg[:,1], pts_neg[:,2],
                           c=COLOR_NEG, s=70, edgecolors='black', zorder=6)
paths7_pos = [ax7a.plot([pts_pos[i,0]]*2, [pts_pos[i,1]]*2, [pts_pos[i,2]]*2,
                         color=COLOR_POS, linestyle=':', alpha=0.5, zorder=3)[0] for i in range(n_pos)]
paths7_neg = [ax7a.plot([pts_neg[i,0]]*2, [pts_neg[i,1]]*2, [pts_neg[i,2]]*2,
                         color=COLOR_NEG, linestyle=':', alpha=0.5, zorder=3)[0] for i in range(n_neg)]

# ── Right Subplot (Static Final) ──
ax7b.scatter(ln_pos[:,0], ln_pos[:,1], ln_pos[:,2], c='black', marker='x', s=50, zorder=4)
ax7b.scatter(ln_neg[:,0], ln_neg[:,1], ln_neg[:,2], c='black', marker='x', s=50, zorder=4)
ax7b.scatter(rms_pos[:,0], rms_pos[:,1], rms_pos[:,2],
             c=COLOR_POS, s=70, edgecolors='black', zorder=6)
ax7b.scatter(rms_neg[:,0], rms_neg[:,1], rms_neg[:,2],
             c=COLOR_NEG, s=70, edgecolors='black', zorder=6)
for i in range(n_pos):
    ax7b.plot([rms_pos[i,0], ln_pos[i,0]], [rms_pos[i,1], ln_pos[i,1]],
              [rms_pos[i,2], ln_pos[i,2]], color='gray', linestyle='-', linewidth=1.2, alpha=0.5, zorder=2)
for i in range(n_neg):
    ax7b.plot([rms_neg[i,0], ln_neg[i,0]], [rms_neg[i,1], ln_neg[i,1]],
              [rms_neg[i,2], ln_neg[i,2]], color='gray', linestyle='-', linewidth=1.2, alpha=0.5, zorder=2)

leg7 = [mlines.Line2D([], [], color='black',    marker='x', linestyle='None', label='Ideal LN Target'),
        mlines.Line2D([], [], color=COLOR_POS,  marker='o', linestyle='None', label='RMSNorm ($\mu > 0$) → +pole'),
        mlines.Line2D([], [], color=COLOR_NEG,  marker='o', linestyle='None', label='RMSNorm ($\mu < 0$) → −pole'),
        mlines.Line2D([], [], color='black',    marker='*', linestyle='None', markersize=10, label='Collapse poles'),
        mlines.Line2D([], [], color='gray',     linestyle='-',  linewidth=1.5, label='RMS vs LN Gap')]
ax7a.legend(handles=leg7, loc='upper left', fontsize=9)

def update7(frame):
    s = np.sin((frame / total_frames) * np.pi / 2)
    cur_pos = pts_pos + s * (rms_pos - pts_pos)
    cur_neg = pts_neg + s * (rms_neg - pts_neg)

    scat7a_pos._offsets3d = (cur_pos[:,0], cur_pos[:,1], cur_pos[:,2])
    scat7a_neg._offsets3d = (cur_neg[:,0], cur_neg[:,1], cur_neg[:,2])
    for i in range(n_pos):
        paths7_pos[i].set_data([pts_pos[i,0], cur_pos[i,0]], [pts_pos[i,1], cur_pos[i,1]])
        paths7_pos[i].set_3d_properties([pts_pos[i,2], cur_pos[i,2]])
    for i in range(n_neg):
        paths7_neg[i].set_data([pts_neg[i,0], cur_neg[i,0]], [pts_neg[i,1], cur_neg[i,1]])
        paths7_neg[i].set_3d_properties([pts_neg[i,2], cur_neg[i,2]])

    # Gentle oscillating rotation to show the 3D structure without losing the pole axis view
    ax7b.view_init(elev=28 + 8 * np.sin(frame * 2 * np.pi / total_frames),
                   azim=45 + frame * (180 / total_frames))
    sys.stdout.write(f'\rRendering Anim 7 Frame {frame+1}/{total_frames}...')
    return fig7,

ani7 = FuncAnimation(fig7, update7, frames=total_frames, interval=50, blit=False)
ani7.save('07_rmsnorm_case2.gif', writer=PillowWriter(fps=fps))
plt.close(fig7)

# ==========================================
# 9. ANIMATION 8: RMSNORM COMPARISON (Both Cases)
# ==========================================
print("\n\nPreparing Animation 8: RMSNorm Comparison (Both Cases)...")
fig8 = plt.figure(figsize=(10, 8), dpi=100)
ax8 = fig8.add_subplot(111, projection='3d')

bound8 = 2.5
ax8.set_xlim([-bound8, bound8]); ax8.set_ylim([-bound8, bound8]); ax8.set_zlim([-bound8, bound8])
ax8.set_xlabel('X'); ax8.set_ylabel('Y'); ax8.set_zlabel('Z')
ax8.set_title("RMSNorm Comparison: $\sigma \gg \mu$ vs $\mu \gg \sigma$", fontsize=14)
fig8.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.95)

# Background
xx8, yy8 = np.meshgrid(np.linspace(-bound8, bound8, 10), np.linspace(-bound8, bound8, 10))
ax8.plot_surface(xx8, yy8, -xx8-yy8, alpha=0.05, color='dodgerblue', zorder=0)
ax8.plot_wireframe(xs, ys, zs, color='gray', alpha=0.15, linewidth=0.5)
ax8.plot(circle_pts[:, 0], circle_pts[:, 1], circle_pts[:, 2], color='gold', linewidth=3.0, zorder=5)

# Pole axis and markers
ax8.plot([-pole_axis[0], pole_axis[0]], [-pole_axis[1], pole_axis[1]], [-pole_axis[2], pole_axis[2]],
         color='dimgray', linewidth=1.5, linestyle='--', alpha=0.7, zorder=1)
ax8.scatter(*pole_axis,   marker='*', s=200, c='darkorange',      zorder=7, edgecolors='black')
ax8.scatter(*(-pole_axis), marker='*', s=200, c='mediumslateblue', zorder=7, edgecolors='black')

# Ideal LN targets (all as black x)
ax8.scatter(ln_l[:,0],   ln_l[:,1],   ln_l[:,2],   c='black', marker='x', s=50, zorder=4)
ax8.scatter(ln_pos[:,0], ln_pos[:,1], ln_pos[:,2], c='black', marker='x', s=50, zorder=4)
ax8.scatter(ln_neg[:,0], ln_neg[:,1], ln_neg[:,2], c='black', marker='x', s=50, zorder=4)

# Case 1: healthy (royalblue, spread around equator)
ax8.scatter(rms_l[:,0], rms_l[:,1], rms_l[:,2], c='royalblue', s=60, edgecolors='black', zorder=6)
# Case 2: unstable, positive mean (darkorange → +pole)
ax8.scatter(rms_pos[:,0], rms_pos[:,1], rms_pos[:,2], c=COLOR_POS, s=60, edgecolors='black', zorder=6)
# Case 2: unstable, negative mean (mediumslateblue → -pole)
ax8.scatter(rms_neg[:,0], rms_neg[:,1], rms_neg[:,2], c=COLOR_NEG, s=60, edgecolors='black', zorder=6)

# Gap lines
for i in range(n_large):
    ax8.plot([rms_l[i,0], ln_l[i,0]], [rms_l[i,1], ln_l[i,1]], [rms_l[i,2], ln_l[i,2]],
             color='gray', linestyle='-', linewidth=1.2, alpha=0.4, zorder=2)
for i in range(n_pos):
    ax8.plot([rms_pos[i,0], ln_pos[i,0]], [rms_pos[i,1], ln_pos[i,1]], [rms_pos[i,2], ln_pos[i,2]],
             color='gray', linestyle='-', linewidth=1.2, alpha=0.4, zorder=2)
for i in range(n_neg):
    ax8.plot([rms_neg[i,0], ln_neg[i,0]], [rms_neg[i,1], ln_neg[i,1]], [rms_neg[i,2], ln_neg[i,2]],
             color='gray', linestyle='-', linewidth=1.2, alpha=0.4, zorder=2)

leg8 = [mlines.Line2D([], [], color='black',     marker='x', linestyle='None', label='Ideal LN Target'),
        mlines.Line2D([], [], color='royalblue',  marker='o', linestyle='None', label='RMSNorm ($\sigma \gg \mu$) — healthy'),
        mlines.Line2D([], [], color=COLOR_POS,    marker='o', linestyle='None', label='RMSNorm ($\mu > 0$) → +pole'),
        mlines.Line2D([], [], color=COLOR_NEG,    marker='o', linestyle='None', label='RMSNorm ($\mu < 0$) → −pole'),
        mlines.Line2D([], [], color='gray',       linestyle='-', linewidth=1.5, label='RMS vs LN Gap')]
ax8.legend(handles=leg8, loc='upper left', fontsize=9)

def update8(frame):
    azim = 45 + frame * (360 / total_frames)
    elev = 28 + 12 * np.sin(frame * 2 * np.pi / total_frames)
    ax8.view_init(elev=elev, azim=azim)
    if frame % 10 == 0 or frame == total_frames - 1:
        sys.stdout.write(f'\rRendering Anim 8 Frame {frame+1}/{total_frames}...')
        sys.stdout.flush()
    return fig8,

ani8 = FuncAnimation(fig8, update8, frames=total_frames, interval=50, blit=False)
ani8.save('08_rmsnorm_comparison.gif', writer=PillowWriter(fps=fps))
plt.close(fig8)
print("\nSaved '08_rmsnorm_comparison.gif'")

print("\n\nAll visualizations successfully generated!")