LLM по частям
← Все главы

22 · Версия материала 7

Не включайте затухание весов в моменты градиента

Соберите AdamW на основе градиентов именованных параметров и моментов с поправкой на смещение, примените затухание весов отдельной ветвью, а затем атомарно зафиксируйте все проверенные обновления.

Предскажите две поправки, которые будут вычтены

В конце главы 21 получены тензорные градиенты функции потерь, усреднённой по токенам. Перед шагом оптимизатора сопоставьте каждый градиент со стабильным именем параметра, по которому он вычислен. AdamW использует каждый именованный градиент, чтобы вычислить обновлённое значение соответствующего параметра. Оптимизатор сначала подготавливает полное именованное обновление, не записывая ни одного значения, а затем записывает все обновлённые значения в те же существующие листовые узлы параметров.

В рассматриваемом обновлении θ0\theta_0 — текущее значение параметра decoder.output.weight из группы с затуханием, а g1g_1 — накопленный градиент усреднённой по токенам функции потерь по этому же параметру:

θ0=[1,−2],g1=[0.2,−0.4].\theta_0=[1,-2], \qquad g_1=[0.2,-0.4].

AdamW хранит векторы моментов этого параметра под именем decoder.output.weight. Историю моментов определяет стабильное имя, а не положение параметра в списке.

Возьмите η=0.1\eta=0.1, β1=β2=0.5\beta_1=\beta_2=0.5, ε=0.1\varepsilon=0.1 и λ=0.1\lambda=0.1. В этом примере коэффициент масштабирования градиента α1=1\alpha_1=1, поэтому градиент, поступающий в AdamW, не меняется: g~1=g1\widetilde g_1=g_1. Отдельно предскажите адаптивную поправку и поправку затухания.

При нулевой инициализации моментов m1=[0.1,−0.2]m_1=[0.1,-0.2], а v1=[0.02,0.08]v_1=[0.02,0.08]. Оба знаменателя поправки равны 0.50.5, поэтому m^1=[0.2,−0.4]\hat m_1=[0.2,-0.4] и v^1=[0.04,0.16]\hat v_1=[0.04,0.16]. Следовательно:

ηm^1v^1+ε≈[0.066667,−0.08],ηλθ0=[0.01,−0.02].\eta\frac{\hat m_1}{\sqrt{\hat v_1}+\varepsilon} \approx[0.066667,-0.08], \qquad \eta\lambda\theta_0=[0.01,-0.02].

Вычтите обе поправки из старого значения и получите прогноз θ1≈[0.923333,−1.9]\theta_1\approx[0.923333,-1.9]. Знаки здесь важны: при затухании отрицательной координаты вычитается отрицательная поправка, поэтому координата приближается к нулю.

Скорректируйте моменты, а затухание примените в обход них

Оптимизатор получает исходный накопленный градиент gtg_t, который хранится у именованного параметра, и проверенный коэффициент 0≤αt≤10\leq\alpha_t\leq1. Для расчёта обновления он получает градиент после масштабирования:

g~t=αtgt.\widetilde g_t=\alpha_t g_t.

Один и тот же коэффициент применяется ко всем координатам полного набора параметров. Это не ещё одна скорость обучения: коэффициент изменяет только градиент, поступающий в моменты. Оптимизатор читает gtg_t, но не изменяет сохранённый градиент. После успешного шага параметр получает новое значение, однако остаётся тем же листовым узлом с тем же накопленным градиентом. AdamW использует этот градиент в расчёте и оставляет его в узле; в главе 23 цикл обучения явно вызывает zero_grad() после обновления и до следующего обратного прохода. При αt=1\alpha_t=1 каждая координата градиента поступает в расчёт без изменения, как в примере этой главы. В главе 33 коэффициент становится меньше единицы, когда общая норма градиента превышает заданный предел.

Рассмотрим также масштабирование с коэффициентом, отличным от единицы. Пусть исходный градиент равен [0.8,−0.4][0.8,-0.4], а αt=0.25\alpha_t=0.25. Тогда

g~t=0.25[0.8,−0.4]=[0.2,−0.1].\widetilde g_t=0.25[0.8,-0.4]=[0.2,-0.1].

В первый момент поступает [0.2,−0.1][0.2,-0.1], а во второй — квадраты координат [0.04,0.01][0.04,0.01]. Отдельная поправка затухания остаётся равной ηλθt−1\eta\lambda\theta_{t-1}.

Используя g~t\widetilde g_t, обновите поэлементно первый момент и второй нецентрированный момент:

mt=β1mt−1+(1−β1)g~t,vt=β2vt−1+(1−β2)g~t2.m_t=\beta_1m_{t-1}+(1-\beta_1)\widetilde g_t, \qquad v_t=\beta_2v_{t-1}+(1-\beta_2)\widetilde g_t^2.

Первый момент переносит вперёд недавнее направление градиента — в этом состоит интуиция метода импульса: совпадающие направления усиливают друг друга, а резкая смена направления смягчается. Второй нецентрированный момент отслеживает недавний квадрат величины градиента. Деление первого момента с поправкой на его среднеквадратичный масштаб адаптирует каждую координату по соотношению недавнего направления и величины; это не означает, что больший градиент всегда даёт меньший по модулю шаг.

Нулевая инициализация занижает первые значения моментов. Компенсируйте это занижение:

m^t=mt1−β1t,v^t=vt1−β2t.\hat m_t=\frac{m_t}{1-\beta_1^t}, \qquad \hat v_t=\frac{v_t}{1-\beta_2^t}.

Теперь примените полную формулу этой главы:

m^t=mt1−β1t,v^t=vt1−β2t,θt=(1−ηλ)θt−1−ηm^tv^t+ε\hat m_t=\frac{m_t}{1-\beta_1^t},\quad \hat v_t=\frac{v_t}{1-\beta_2^t},\quad \theta_t=(1-\eta\lambda)\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\varepsilon}

То же обновление можно представить как вычитание двух явных слагаемых: адаптивной поправки и ηλθt−1\eta\lambda\theta_{t-1}. Затухание весов, отделённое от градиентного обновления, означает, что второе слагаемое не входит ни в gtg_t, ни в g~t\widetilde g_t, ни в mtm_t или vtv_t. Поэтому изменение αt\alpha_t не масштабирует поправку затухания. Итоговое новое значение при этом по-прежнему учитывает адаптивную ветвь.

В этой формуле λ\lambda — эффективный коэффициент для текущего именованного параметра: он равен заданному затуханию для параметра из группы с затуханием и 00 для параметра из группы без затухания.

Карта групп реализует настраиваемое правило: она явно разбивает набор стабильных имён параметров на группы с затуханием и без него. Эти группы не пересекаются, а их объединение содержит каждое имя. В примере курса правило относит decoder.output.weight к группе с затуханием, а decoder.norm.scale — к группе без затухания; такое назначение не следует из формулы AdamW. Поэтому эффективный коэффициент λ\lambda для масштаба нормализации равен 00, и отдельная поправка затухания не стягивает обучаемый аффинный масштаб напрямую к нулю.

Различайте состояние шага и состояние параметра

  • θt−1\theta_{t-1} — один именованный тензор параметра до шага tt.
  • gtg_t — его исходный накопленный градиент функции потерь, усреднённой по токенам; он остаётся в существующем листовом узле, пока вызывающий код явно его не обнулит.
  • αt\alpha_t — проверенный коэффициент, общий для всех координат градиента; 0≤αt≤10\leq\alpha_t\leq1.
  • g~t=αtgt\widetilde g_t=\alpha_tg_t — градиент после масштабирования, который поступает в оба момента.
  • mtm_t и vtv_t — экспоненциально сглаживаемые оценки той же формы, хранящиеся по стабильным именам.
  • β1\beta_1 и β2\beta_2 задают сглаживание соответственно первого момента и второго нецентрированного момента.
  • m^t\hat m_t и v^t\hat v_t вносят поправку на смещение моментов к нулю в начале обучения.
  • η>0\eta>0 — скорость обучения, а λ≥0\lambda\geq0 — коэффициент затухания.
  • ε>0\varepsilon>0 стабилизирует знаменатель адаптивной ветви после извлечения квадратного корня.
  • θt\theta_t — обновлённое значение, которое записывается в тот же обучаемый листовой узел лишь после успешной подготовки всего набора.

При проверке с новым состоянием m0=v0=0m_0=v_0=0 и g1=0g_1=0, поэтому оба новых момента и адаптивная поправка равны нулю. Параметр из группы с затуханием всё равно уменьшается на ηλθ0\eta\lambda\theta_0, а параметр из группы без затухания не меняется. Если раньше градиент был ненулевым, равенство gt=0g_t=0 убирает лишь новый вклад: сохранённые моменты затухают, но всё ещё могут дать ненулевое адаптивное обновление. Принадлежность к группе управляет затуханием весов, а не памятью моментов.

От градиента по одному слову к декодерам, обучаемым с AdamW

Bengio et al., A Neural Probabilistic Language Model описывают более ранний этап развития языковых моделей. В нейронной языковой модели Бенжио и соавторов после предъявления одного слова из обучающего корпуса вместе с его контекстом выполняется прямое стохастическое обновление параметров. На каждом шаге единственная скалярная скорость обучения умножает градиент текущего примера; памяти о предыдущих градиентах по отдельным координатам нет. В обозначениях минимизации функции потерь основное действие выглядит просто:

θ←θ−ηg.\theta\leftarrow\theta-\eta g.

По сравнению с этим прямым правилом SGD метод импульса добавляет затухающую память о прошлых направлениях. Позднее Adam хранит две экспоненциально сглаженные оценки — первый момент и второй нецентрированный момент градиента — и вносит поправку на их начальное смещение к нулю; если член L2L_2 включён в градиент, эта пропорциональная параметру составляющая попадает в обе скользящие оценки. В AdamW поправка, стягивающая параметр к нулю, напротив, вынесена из градиента, по которому обновляются эти адаптивные моменты.

В кратких обозначениях минимизации функции потерь метод импульса использует ut=μut−1+gtu_t=\mu u_{t-1}+g_t и θt=θt−1−ηut\theta_t=\theta_{t-1}-\eta u_t. Если штраф L2L_2 включён в градиент, в моменты Adam поступает gt+λθt−1g_t+\lambda\theta_{t-1}. В AdamW пропорциональная параметру составляющая не входит в градиент, поступающий в моменты, а затухание весов применяется отдельно. В прямом примере этой главы αt=1\alpha_t=1, поэтому входом для моментов служит gtg_t; после ограничения нормы в главе 33 им служит g~t\widetilde g_t. Здесь utu_t — накопленная скорость обновления, а 0≤μ<10\leq\mu\lt1 — коэффициент её сглаживания.

Kingma and Ba, Adam: A Method for Stochastic Optimization обосновывают адаптивный этап: Adam Кингмы и Ба хранит две экспоненциально сглаженные оценки — первый момент и второй нецентрированный момент градиента — и компенсирует их начальное смещение к нулю. Деление на 1−β1t1-\beta_1^t и 1−β2t1-\beta_2^t нужно потому, что нулевая инициализация смещает первые оценки к нулю. Adam — оптимизатор общего назначения, а не архитектура LLM.

Loshchilov and Hutter, Decoupled Weight Decay Regularization показывают, что штраф L2L_2, добавленный к градиенту функции потерь, не равнозначен затуханию весов при использовании адаптивного оптимизатора. В AdamW Лошчилова и Хуттера пропорциональное параметру затухание вынесено из градиента, по которому обновляются адаптивные моменты.

В статье о LLaMA описано применение AdamW при предобучении декодерных языковых моделей с числом параметров от 77B до 6565B. Оптимизатор курса записывает каждое проверенное новое значение в существующий листовой узел именованного параметра; глава 33 добавляет скорость обучения из расписания, один проверенный множитель ограничения общей нормы градиента и явное обнуление градиентов. Смешанная точность и распределённое состояние оптимизатора остаются за рамками этой реализации.

Именно эта последовательность — от непосредственных градиентов предсказания следующего слова до адаптивных обновлений с отдельным затуханием, применяемых при предобучении декодеров, — определяет место AdamW на пути к современным LLM.

Touvron et al., LLaMA: Open and Efficient Foundation Language Models подтверждают современный пример: в статье о LLaMA описано применение AdamW при предобучении декодерных языковых моделей с числом параметров от 77B до 6565B. В их конфигурации используются β1=0.9\beta_1=0.9, β2=0.95\beta_2=0.95, затухание весов 0.10.1, ограничение нормы градиента, разогрев и косинусное снижение скорости обучения. В этой главе реализован механизм оптимизатора, а не вся конфигурация крупномасштабного обучения.

Исторический пример выполняет два шага обычного SGD, метода импульса, Adam со штрафом L2L_2, включённым в градиент, и AdamW для одной и той же скалярной последовательности градиентов функции потерь, а затем строит заданную траекторию при разной кривизне по координатам. Rust служит лишь исполняемой средой для этого сравнения, которое не зависит от языка программирования:

Внутри цикла, который строит траекторию, отдельный блок вызывает value() и получает текущее значение параметра AdamW без копирования, удерживая доступ только для чтения. Пока этот доступ действует, блок вычисляет обе координаты градиента квадратичной функции. При выходе из блока доступ завершается. Если бы заимствование для чтения осталось активным, step() не смог бы получить изменяемый доступ к тому же значению. После шага цикл явно вызывает zero_grad(), чтобы следующий обратный проход не добавлял новый градиент к уже использованному. Границы блоков определяют только время владения доступом; формулы оптимизатора и записанная траектория от этого не меняются.

Сопоставьте SGD, метод импульса, Adam со штрафом, включённым в градиент, и AdamW с отдельным затуханием rust/demos/ch22-adamw/src/lib.rs#historical-optimizer-road 
/// Four optimizer endpoints for one loss-gradient sequence.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct HistoricalUpdates {
    pub sgd: f64,
    pub momentum: f64,
    pub adam_l2: f64,
    pub adamw: f64,
}

pub fn historical_updates(parameter_value: f64, gradients: [f64; 2]) -> HistoricalUpdates {
    let mut sgd = parameter_value;
    for gradient in gradients {
        sgd -= LEARNING_RATE * gradient;
    }

    let mut momentum = parameter_value;
    let mut velocity = 0.0;
    for gradient in gradients {
        velocity = MOMENTUM_RATE * velocity + gradient;
        momentum -= LEARNING_RATE * velocity;
    }

    let adam_l2 = two_step_adaptive_update(parameter_value, gradients, true);
    let adamw = two_step_adaptive_update(parameter_value, gradients, false);
    HistoricalUpdates {
        sgd,
        momentum,
        adam_l2,
        adamw,
    }
}

fn two_step_adaptive_update(
    parameter_value: f64,
    gradients: [f64; 2],
    couple_l2_into_gradient: bool,
) -> f64 {
    let decoupled_decay = if couple_l2_into_gradient {
        0.0
    } else {
        WEIGHT_DECAY
    };
    let config = AdamWConfig::new(LEARNING_RATE, BETA1, BETA2, EPSILON, decoupled_decay)
        .expect("historical fixture configuration is valid");
    let parameters = vec![parameter("history.weight", &[1], &[parameter_value])];
    let mut optimizer = AdamW::new(config);
    for gradient in gradients {
        let current = parameters[0].tensor().value().as_slice()[0];
        let optimizer_gradient = if couple_l2_into_gradient {
            gradient + WEIGHT_DECAY * current
        } else {
            gradient
        };
        seed_gradient(&parameters[0], &[optimizer_gradient]);
        optimizer
            .step(&parameters)
            .expect("historical fixture update is finite");
        parameters[0]
            .tensor()
            .zero_grad()
            .expect("the historical fixture clears each used gradient");
    }
    parameters[0].tensor().value().as_slice()[0]
}

/// One exact point on the same anisotropic objective for both update rules.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TrajectoryPoint {
    pub step: usize,
    pub sgd: [f64; 2],
    pub adamw: [f64; 2],
}

/// Compares four updates on `q(x,y)=(x^2+4y^2)/2` from the same start.
pub fn anisotropic_trajectory() -> Vec<TrajectoryPoint> {
    const CURVATURE: [f64; 2] = [1.0, 4.0];
    const STEPS: usize = 4;

    let mut sgd = [1.0, 1.0];
    let adamw_parameter = parameter("trajectory.weight", &[2], &[1.0, 1.0]);
    let mut optimizer = AdamW::new(fixture_config());
    let mut points = vec![TrajectoryPoint {
        step: 0,
        sgd,
        adamw: [1.0, 1.0],
    }];

    for step in 1..=STEPS {
        let sgd_gradient = [CURVATURE[0] * sgd[0], CURVATURE[1] * sgd[1]];
        for axis in 0..2 {
            sgd[axis] -= LEARNING_RATE * sgd_gradient[axis];
        }

        let adamw_gradient = {
            let value = adamw_parameter.tensor().value();
            [
                CURVATURE[0] * value.as_slice()[0],
                CURVATURE[1] * value.as_slice()[1],
            ]
        };
        seed_gradient(&adamw_parameter, &adamw_gradient);
        optimizer
            .step(std::slice::from_ref(&adamw_parameter))
            .expect("bounded trajectory stays finite");
        adamw_parameter
            .tensor()
            .zero_grad()
            .expect("the trajectory clears each used gradient");
        let next = adamw_parameter.tensor().value();
        points.push(TrajectoryPoint {
            step,
            sgd,
            adamw: [next.as_slice()[0], next.as_slice()[1]],
        });
    }
    points
}

Подготовьте все именованные значения до записи в существующие узлы

Допустимые диапазоны скалярных параметров проверяются до создания состояния оптимизатора:

Проверьте скорость обучения, коэффициенты сглаживания моментов, стабилизатор и затухание rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-configuration 
/// The five scalar controls used by one fixed-learning-rate AdamW optimizer.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct AdamWConfig {
    learning_rate: f64,
    beta1: f64,
    beta2: f64,
    epsilon: f64,
    weight_decay: f64,
}

impl AdamWConfig {
    /// Validates the scalar controls for one AdamW configuration.
    pub fn new(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
        weight_decay: f64,
    ) -> Result<Self, AdamWError> {
        if !learning_rate.is_finite() || learning_rate <= 0.0 {
            return Err(AdamWError::InvalidLearningRate {
                value: learning_rate,
            });
        }
        if !beta1.is_finite() || !(0.0..1.0).contains(&beta1) {
            return Err(AdamWError::InvalidBeta1 { value: beta1 });
        }
        if !beta2.is_finite() || !(0.0..1.0).contains(&beta2) {
            return Err(AdamWError::InvalidBeta2 { value: beta2 });
        }
        if !epsilon.is_finite() || epsilon <= 0.0 {
            return Err(AdamWError::InvalidEpsilon { value: epsilon });
        }
        if !weight_decay.is_finite() || weight_decay < 0.0 {
            return Err(AdamWError::InvalidWeightDecay {
                value: weight_decay,
            });
        }
        Ok(Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
        })
    }

    pub const fn learning_rate(self) -> f64 {
        self.learning_rate
    }

    pub const fn beta1(self) -> f64 {
        self.beta1
    }

    pub const fn beta2(self) -> f64 {
        self.beta2
    }

    pub const fn epsilon(self) -> f64 {
        self.epsilon
    }

    pub const fn weight_decay(self) -> f64 {
        self.weight_decay
    }
}

Явная карта групп назначает каждое имя ровно один раз: decoder.output.weight входит в группу с затуханием, а decoder.norm.scale — в группу без затухания. Пустое назначение в целом, повторение имени внутри группы, пересечения, пропущенные имена и лишние назначения отклоняются до фиксации; при этом любая из двух групп по отдельности может быть пустой:

Разбейте стабильные имена на явные группы с затуханием и без затухания rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-parameter-groups 
/// The two explicit parameter groups used by the course's decay policy.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdamWGroup {
    Decay,
    NoDecay,
}

impl fmt::Display for AdamWGroup {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Decay => formatter.write_str("decay"),
            Self::NoDecay => formatter.write_str("no-decay"),
        }
    }
}

/// Exact stable-name assignments for decayed and decay-excluded parameters.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AdamWParameterGroups {
    decay: BTreeSet<String>,
    no_decay: BTreeSet<String>,
}

impl AdamWParameterGroups {
    pub fn new<D, N, DS, NS>(decay: D, no_decay: N) -> Result<Self, AdamWError>
    where
        D: IntoIterator<Item = DS>,
        N: IntoIterator<Item = NS>,
        DS: Into<String>,
        NS: Into<String>,
    {
        let decay = collect_group(decay, AdamWGroup::Decay)?;
        let no_decay = collect_group(no_decay, AdamWGroup::NoDecay)?;
        if decay.is_empty() && no_decay.is_empty() {
            return Err(AdamWError::EmptyParameterGroups);
        }
        if let Some(name) = decay.intersection(&no_decay).next() {
            return Err(AdamWError::ParameterInMultipleGroups {
                name: name.to_owned(),
            });
        }
        Ok(Self { decay, no_decay })
    }

    pub fn decayed_names(&self) -> impl ExactSizeIterator<Item = &str> {
        self.decay.iter().map(String::as_str)
    }

    pub fn excluded_names(&self) -> impl ExactSizeIterator<Item = &str> {
        self.no_decay.iter().map(String::as_str)
    }

    fn parameter_names(&self) -> Vec<String> {
        self.decay.union(&self.no_decay).cloned().collect()
    }

    fn decays(&self, name: &str) -> bool {
        self.decay.contains(name)
    }
}

fn collect_group<I, S>(names: I, group: AdamWGroup) -> Result<BTreeSet<String>, AdamWError>
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    let mut collected = BTreeSet::new();
    for name in names {
        let name = name.into();
        if name.is_empty() {
            return Err(AdamWError::EmptyGroupedParameterName { group });
        }
        if !collected.insert(name.clone()) {
            return Err(AdamWError::DuplicateGroupedParameter { group, name });
        }
    }
    Ok(collected)
}

Состояние моментов имеет ту же форму, что и соответствующий параметр, и хранится по его стабильному имени:

Свяжите каждую пару векторов моментов с одним стабильным именем rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-moment-state 
/// Name-keyed optimizer memory for one parameter tensor.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWMomentState {
    shape: Vec<usize>,
    first: Vec<f64>,
    second: Vec<f64>,
}

impl AdamWMomentState {
    fn zeros(shape: &[usize], elements: usize) -> Self {
        Self {
            shape: shape.to_vec(),
            first: vec![0.0; elements],
            second: vec![0.0; elements],
        }
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn first_moment(&self) -> &[f64] {
        &self.first
    }

    pub fn second_moment(&self) -> &[f64] {
        &self.second
    }
}

Обычные методы step и step_with_learning_rate выполняют полное атомарное обновление и возвращают только номер успешно зафиксированного шага. Ещё один обычный метод, step_with_learning_rate_and_gradient_scale, принимает скорость из расписания и αt\alpha_t, но тоже возвращает только номер шага. Четыре существовавших ранее метода — обычные и с трассировкой — работают в точности как при αt=1\alpha_t=1. Структура AdamWStep создаётся только при вызове step_with_trace или step_with_learning_rate_and_trace.

Все точки входа используют одну и ту же внутреннюю операцию подготовки и атомарной фиксации, а также один и тот же покоординатный расчёт AdamW. Трассировка записывает значения, получаемые в этом расчёте, а не вычисляет обновление повторно. Метод gradient() у записи трассировки возвращает именно градиент g~t\widetilde g_t, который поступил в оба момента, а не результат повторного расчёта в коде трассировки.

Разделите обычное выполнение AdamW и явно запрашиваемую трассировку rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-execution-and-trace-api 
    /// Reads the accumulated gradients and atomically updates every live leaf.
    ///
    /// All arithmetic, tensor construction, and optimizer-state changes are
    /// prepared first. An error leaves both the supplied parameters and this
    /// optimizer bit-identical. A successful commit preserves every parameter
    /// node and leaves its accumulated gradient for the caller to clear.
    /// The result is only the committed step number; use `step_with_trace` when
    /// the elementwise update vectors are needed for inspection.
    pub fn step(&mut self, parameters: &[NamedParameter]) -> Result<u64, AdamWError> {
        self.step_with_config(parameters, self.config, 1.0, NoAdamWTrace)
    }

    /// Applies the same transaction while recording every elementwise update.
    pub fn step_with_trace(
        &mut self,
        parameters: &[NamedParameter],
    ) -> Result<AdamWStep, AdamWError> {
        let observer = RecordAdamWTrace::with_capacity(parameters.len());
        self.step_with_config(parameters, self.config, 1.0, observer)
    }

    /// Applies one validated scheduled learning rate without resetting moments.
    ///
    /// The override belongs only to this update; `config()` keeps the optimizer's
    /// base rate. An invalid rate or any later preparation error leaves the
    /// parameters, moments, powers, and step counter unchanged.
    pub fn step_with_learning_rate(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
    ) -> Result<u64, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        self.step_with_config(parameters, step_config, 1.0, NoAdamWTrace)
    }

    /// Applies one scheduled rate and one validated global gradient scale.
    ///
    /// The scale multiplies only the gradient used by Adam's moments. The
    /// decoupled weight-decay branch continues to use the unscaled parameter.
    pub fn step_with_learning_rate_and_gradient_scale(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
        gradient_scale: f64,
    ) -> Result<u64, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        self.step_with_config(parameters, step_config, gradient_scale, NoAdamWTrace)
    }

    /// Applies a scheduled learning rate and records the complete update trace.
    pub fn step_with_learning_rate_and_trace(
        &mut self,
        parameters: &[NamedParameter],
        learning_rate: f64,
    ) -> Result<AdamWStep, AdamWError> {
        let step_config = self.config.with_learning_rate(learning_rate)?;
        let observer = RecordAdamWTrace::with_capacity(parameters.len());
        self.step_with_config(parameters, step_config, 1.0, observer)
    }

Только явно запрошенная трассировка сохраняет точные входные векторы, моменты, поправки и новые значения, разбираемые в этой главе:

Определите данные, сохраняемые только для шага AdamW с трассировкой rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-state-and-evidence 
/// Exact elementwise evidence prepared for one named parameter in a step.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWParameterUpdate {
    name: String,
    shape: Vec<usize>,
    before: Vec<f64>,
    gradient: Vec<f64>,
    decay_applied: bool,
    effective_weight_decay: f64,
    first_moment: Vec<f64>,
    second_moment: Vec<f64>,
    corrected_first_moment: Vec<f64>,
    corrected_second_moment: Vec<f64>,
    adaptive_direction: Vec<f64>,
    adaptive_delta: Vec<f64>,
    decay_delta: Vec<f64>,
    after: Vec<f64>,
}

impl AdamWParameterUpdate {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn before(&self) -> &[f64] {
        &self.before
    }

    /// The effective gradient used by the moment calculation.
    pub fn gradient(&self) -> &[f64] {
        &self.gradient
    }

    pub const fn decay_applied(&self) -> bool {
        self.decay_applied
    }

    pub const fn effective_weight_decay(&self) -> f64 {
        self.effective_weight_decay
    }

    pub fn first_moment(&self) -> &[f64] {
        &self.first_moment
    }

    pub fn second_moment(&self) -> &[f64] {
        &self.second_moment
    }

    pub fn corrected_first_moment(&self) -> &[f64] {
        &self.corrected_first_moment
    }

    pub fn corrected_second_moment(&self) -> &[f64] {
        &self.corrected_second_moment
    }

    pub fn adaptive_direction(&self) -> &[f64] {
        &self.adaptive_direction
    }

    pub fn adaptive_delta(&self) -> &[f64] {
        &self.adaptive_delta
    }

    pub fn decay_delta(&self) -> &[f64] {
        &self.decay_delta
    }

    pub fn after(&self) -> &[f64] {
        &self.after
    }
}

/// The optional trace for one committed multi-parameter update.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWStep {
    step: u64,
    learning_rate: f64,
    first_correction: f64,
    second_correction: f64,
    updates: Vec<AdamWParameterUpdate>,
}

impl AdamWStep {
    pub const fn step(&self) -> u64 {
        self.step
    }

    pub const fn learning_rate(&self) -> f64 {
        self.learning_rate
    }

    pub const fn first_correction(&self) -> f64 {
        self.first_correction
    }

    pub const fn second_correction(&self) -> f64 {
        self.second_correction
    }

    pub fn updates(&self) -> &[AdamWParameterUpdate] {
        &self.updates
    }
}

Трассировка необязательна, но требования к атомарности обновления сохраняются. Открытый метод сначала проверяет скорость обучения при создании конфигурации для очередного шага. Затем он переходит к общему пути, где допустимый диапазон 0≤αt≤10\leq\alpha_t\leq1 проверяется ещё до чтения параметров. Для каждой координаты оптимизатор один раз получает градиент после масштабирования. Затем, до изменения параметров и состояния оптимизатора, он подготавливает новые значения моментов для каждого стабильного имени, степени β1t\beta_1^t и β2t\beta_2^t, каждый обновлённый тензор параметра и следующую версию каждого значения. Поправка затухания по-прежнему вычисляется из значения параметра до шага и не умножается на αt\alpha_t.

Когда все вычисления и создание тензоров завершены, реализация пытается получить изменяемый доступ к значениям всех существующих параметров. До завершения этой проверки ни один параметр не меняется. Если для значения хотя бы одного параметра всё ещё активно заимствование для чтения, получить изменяемый доступ к этому узлу нельзя. AdamW освобождает все уже полученные доступы, а параметры и состояние оптимизатора не меняются. После получения всех доступов подготовленные тензоры той же формы устанавливаются как новые значения существующих узлов, а версии значений увеличиваются. Эти последние присваивания не могут завершиться с ошибкой. Затем оптимизатор устанавливает подготовленные моменты, степени и номер шага. Градиенты хранятся отдельно от значений параметров, поэтому исходные накопленные градиенты остаются доступными, пока вызывающий код явно их не обнулит.

У каждого узла параметра есть монотонно возрастающая версия значения параметра. После успешного шага AdamW версия этого узла увеличивается на единицу, а после неудачного шага остаётся прежней. Версия отмечает изменение значения внутри узла; это не номер шага оптимизатора и не часть контрольной точки. В главах 37 и 38 кэш KV связывается и с идентичностью узла, и с зафиксированной версией его значения. Ключи и значения, вычисленные до обновления параметров, устаревают, хотя идентичность узла после шага AdamW не меняется.

При прямом проходе на каждом ребре операции также сохраняется текущая версия операнда. Поэтому граф, сохранённый до успешного шага AdamW, нельзя использовать для обратного прохода с обновлёнными параметрами: несовпадение версий отклоняется до изменения градиентов или состояния графа. Для новых значений параметров нужно выполнить новый прямой проход и построить новый граф.

Полная подготовка, проверка изменяемого доступа и запись в существующие листовые узлы показаны в исходном коде:

Подготовьте полное именованное обновление, затем запишите значения в существующие узлы rust/crates/llm-from-scratch/src/training/adamw.rs#transactional-adamw-step 
    fn step_with_config<O: AdamWStepObserver>(
        &mut self,
        parameters: &[NamedParameter],
        step_config: AdamWConfig,
        gradient_scale: f64,
        mut observer: O,
    ) -> Result<O::Output, AdamWError> {
        validate_gradient_scale(gradient_scale)?;
        let actual_names = validate_parameter_names(parameters)?;
        if let Some(groups) = &self.groups {
            let expected_names = groups.parameter_names();
            if expected_names != actual_names {
                return Err(AdamWError::ParameterSetChanged {
                    expected: expected_names,
                    actual: actual_names,
                });
            }
        }
        let next_step = self.step.checked_add(1).ok_or(AdamWError::StepOverflow)?;
        let next_beta1_power = self.beta1_power * step_config.beta1;
        let next_beta2_power = self.beta2_power * step_config.beta2;
        let first_correction = 1.0 - next_beta1_power;
        let second_correction = 1.0 - next_beta2_power;

        let mut candidate_states = if self.step == 0 {
            let mut states = BTreeMap::new();
            for parameter in parameters.iter() {
                let value = parameter.tensor().value();
                states.insert(
                    parameter.name().to_owned(),
                    AdamWMomentState::zeros(value.shape(), value.len()),
                );
            }
            states
        } else {
            let expected_names = self.states.keys().cloned().collect::<Vec<_>>();
            if expected_names != actual_names {
                return Err(AdamWError::ParameterSetChanged {
                    expected: expected_names,
                    actual: actual_names,
                });
            }
            self.states.clone()
        };

        let mut candidate_values = Vec::with_capacity(parameters.len());
        let mut next_revisions = Vec::with_capacity(parameters.len());
        for parameter in parameters.iter() {
            let name = parameter.name();
            let before = parameter.tensor().value();
            let gradient =
                parameter
                    .tensor()
                    .gradient()
                    .ok_or_else(|| AdamWError::MissingGradient {
                        name: name.to_owned(),
                    })?;
            if before.shape() != gradient.shape() {
                return Err(AdamWError::GradientShapeMismatch {
                    name: name.to_owned(),
                    parameter: before.shape().to_vec(),
                    gradient: gradient.shape().to_vec(),
                });
            }

            let state = candidate_states
                .get_mut(name)
                .expect("validated parameter names have candidate state");
            if state.shape != before.shape() {
                return Err(AdamWError::ParameterShapeChanged {
                    name: name.to_owned(),
                    expected: state.shape.clone(),
                    actual: before.shape().to_vec(),
                });
            }

            let after = prepare_parameter_update(
                AdamWPreparation {
                    config: step_config,
                    decay_applied: self
                        .groups
                        .as_ref()
                        .is_none_or(|groups| groups.decays(name)),
                    first_correction,
                    second_correction,
                    name,
                    gradient_scale,
                },
                &before,
                &gradient,
                state,
                &mut observer,
            )?;
            let tensor = Tensor::from_vec(before.shape().to_vec(), after)?;
            let next_revision = parameter.tensor().next_value_revision().ok_or_else(|| {
                AdamWError::ParameterRevisionOverflow {
                    name: name.to_owned(),
                }
            })?;
            candidate_values.push(tensor);
            next_revisions.push(next_revision);
        }

        let observation = observer.finish(
            next_step,
            step_config.learning_rate,
            first_correction,
            second_correction,
        );

        let mut value_writes = Vec::with_capacity(parameters.len());
        for parameter in parameters {
            let write = parameter.tensor().try_value_write().map_err(|_| {
                AdamWError::ParameterValueBorrowed {
                    name: parameter.name().to_owned(),
                }
            })?;
            value_writes.push(write);
        }

        for ((write, value), revision) in value_writes
            .into_iter()
            .zip(candidate_values)
            .zip(next_revisions)
        {
            write.commit(value, revision);
        }
        self.step = next_step;
        self.beta1_power = next_beta1_power;
        self.beta2_power = next_beta2_power;
        self.states = candidate_states;

        Ok(observation)
    }

Типизированные ошибки охватывают недопустимые значения конфигурации и коэффициента масштабирования градиента, пустые наборы и повторы, изменившиеся имена и формы, отсутствующие либо несовпадающие градиенты, переполнение счётчика или версии значения, занятый доступ к значению, неконечные результаты на промежуточных этапах и ошибки создания тензоров:

Отклоняйте недопустимые обновления без частичной записи в существующие значения rust/crates/llm-from-scratch/src/training/adamw.rs#adamw-errors 
/// The arithmetic stage that first produced a non-finite candidate value.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdamWArithmetic {
    FirstMoment,
    SquaredGradient,
    SecondMoment,
    CorrectedFirstMoment,
    CorrectedSecondMoment,
    AdaptiveDirection,
    AdaptiveDelta,
    DecayDelta,
    Parameter,
}

impl fmt::Display for AdamWArithmetic {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::FirstMoment => "first moment",
            Self::SquaredGradient => "squared gradient",
            Self::SecondMoment => "second moment",
            Self::CorrectedFirstMoment => "bias-corrected first moment",
            Self::CorrectedSecondMoment => "bias-corrected second moment",
            Self::AdaptiveDirection => "adaptive direction",
            Self::AdaptiveDelta => "adaptive update",
            Self::DecayDelta => "decoupled decay update",
            Self::Parameter => "updated parameter",
        };
        formatter.write_str(name)
    }
}

/// A deterministic rejection that leaves parameters and optimizer state intact.
#[derive(Clone, Debug, PartialEq)]
pub enum AdamWError {
    InvalidLearningRate {
        value: f64,
    },
    InvalidGradientScale {
        value: f64,
    },
    InvalidBeta1 {
        value: f64,
    },
    InvalidBeta2 {
        value: f64,
    },
    InvalidEpsilon {
        value: f64,
    },
    InvalidWeightDecay {
        value: f64,
    },
    EmptyParameterGroups,
    EmptyGroupedParameterName {
        group: AdamWGroup,
    },
    DuplicateGroupedParameter {
        group: AdamWGroup,
        name: String,
    },
    ParameterInMultipleGroups {
        name: String,
    },
    EmptyParameterSet,
    DuplicateParameterName {
        name: String,
        first: usize,
        repeated: usize,
    },
    ParameterSetChanged {
        expected: Vec<String>,
        actual: Vec<String>,
    },
    ParameterShapeChanged {
        name: String,
        expected: Vec<usize>,
        actual: Vec<usize>,
    },
    MissingGradient {
        name: String,
    },
    GradientShapeMismatch {
        name: String,
        parameter: Vec<usize>,
        gradient: Vec<usize>,
    },
    ParameterRevisionOverflow {
        name: String,
    },
    ParameterValueBorrowed {
        name: String,
    },
    StepOverflow,
    NonFiniteArithmetic {
        name: String,
        index: usize,
        stage: AdamWArithmetic,
        value: f64,
    },
    Tensor(TensorError),
}

impl fmt::Display for AdamWError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidLearningRate { value } => write!(
                formatter,
                "learning rate must be finite and greater than zero, got {value}"
            ),
            Self::InvalidGradientScale { value } => write!(
                formatter,
                "gradient scale must be finite in the closed interval [0,1], got {value}"
            ),
            Self::InvalidBeta1 { value } => write!(
                formatter,
                "beta1 must be finite in the half-open interval [0,1), got {value}"
            ),
            Self::InvalidBeta2 { value } => write!(
                formatter,
                "beta2 must be finite in the half-open interval [0,1), got {value}"
            ),
            Self::InvalidEpsilon { value } => write!(
                formatter,
                "epsilon must be finite and greater than zero, got {value}"
            ),
            Self::InvalidWeightDecay { value } => write!(
                formatter,
                "weight decay must be finite and non-negative, got {value}"
            ),
            Self::EmptyParameterGroups => formatter
                .write_str("explicit AdamW parameter groups must assign at least one stable name"),
            Self::EmptyGroupedParameterName { group } => {
                write!(
                    formatter,
                    "the {group} group contains an empty parameter name"
                )
            }
            Self::DuplicateGroupedParameter { group, name } => write!(
                formatter,
                "parameter name {name:?} repeats inside the {group} group"
            ),
            Self::ParameterInMultipleGroups { name } => write!(
                formatter,
                "parameter name {name:?} appears in both the decay and no-decay groups"
            ),
            Self::EmptyParameterSet => {
                formatter.write_str("AdamW needs at least one named parameter")
            }
            Self::DuplicateParameterName {
                name,
                first,
                repeated,
            } => write!(
                formatter,
                "parameter name {name:?} first appears at index {first} and repeats at index {repeated}"
            ),
            Self::ParameterSetChanged { expected, actual } => write!(
                formatter,
                "parameter-name set changed from {expected:?} to {actual:?}"
            ),
            Self::ParameterShapeChanged {
                name,
                expected,
                actual,
            } => write!(
                formatter,
                "parameter {name:?} changed shape from {expected:?} to {actual:?}"
            ),
            Self::MissingGradient { name } => {
                write!(formatter, "parameter {name:?} has no stored gradient")
            }
            Self::GradientShapeMismatch {
                name,
                parameter,
                gradient,
            } => write!(
                formatter,
                "parameter {name:?} has shape {parameter:?}, but its gradient has shape {gradient:?}"
            ),
            Self::ParameterRevisionOverflow { name } => {
                write!(
                    formatter,
                    "parameter {name:?} value revision overflowed u64"
                )
            }
            Self::ParameterValueBorrowed { name } => write!(
                formatter,
                "parameter {name:?} cannot be updated while its value is borrowed"
            ),
            Self::StepOverflow => formatter.write_str("AdamW step counter overflowed u64"),
            Self::NonFiniteArithmetic {
                name,
                index,
                stage,
                value,
            } => write!(
                formatter,
                "parameter {name:?} produced non-finite {stage} at flat index {index}: {value}"
            ),
            Self::Tensor(error) => error.fmt(formatter),
        }
    }
}

impl Error for AdamWError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Tensor(error) => Some(error),
            _ => None,
        }
    }
}

impl From<TensorError> for AdamWError {
    fn from(error: TensorError) -> Self {
        Self::Tensor(error)
    }
}

В рассматриваемом примере трассировка явно запрашивается для именованного обновления и отдельно для проверки с исходно нулевыми моментами, потому что в главе разбираются точные векторы обоих обновлений. Для проверки отката при изменившемся наборе имён достаточно получить ошибку, а для исторического сопоставления и траектории — итоговые значения параметров. Поэтому во всех трёх случаях вызываются обычные методы:

Подготовьте точные данные первого шага и атомарной операции rust/demos/ch22-adamw/src/lib.rs#chapter-adamw-fixture 
pub fn learner_evidence() -> LearnerEvidence {
    let config = fixture_config();
    let parameters = vec![
        parameter("decoder.output.weight", &[2], &[1.0, -2.0]),
        parameter("decoder.norm.scale", &[1], &[0.5]),
    ];
    seed_gradient(&parameters[0], &[0.2, -0.4]);
    let original_leaves = parameters
        .iter()
        .map(|parameter| parameter.tensor().clone())
        .collect::<Vec<_>>();
    let mut optimizer = AdamW::with_parameter_groups(config, fixture_groups());
    let step = optimizer
        .step_with_trace(&parameters)
        .expect("the complete named set updates atomically");

    let raw_gradients_retained =
        parameters
            .iter()
            .zip(step.updates())
            .all(|(parameter, update)| {
                parameter.tensor().gradient().is_some_and(|gradient| {
                    gradient
                        .as_slice()
                        .iter()
                        .zip(update.gradient())
                        .all(|(actual, expected)| actual.to_bits() == expected.to_bits())
                })
            });
    let parameter_nodes_preserved = parameters
        .iter()
        .zip(&original_leaves)
        .all(|(parameter, original)| parameter.tensor().is_same_node(original));
    let state_names = optimizer.parameter_names().map(str::to_owned).collect();

    let zero_gradient_update = zero_gradient_moment_probe(config);
    let (rejected_error, rejection_rolled_back) = rejected_set_probe(&optimizer, &parameters);
    LearnerEvidence {
        config,
        step,
        state_names,
        raw_gradients_retained,
        parameter_nodes_preserved,
        zero_gradient_update,
        rejected_error,
        rejection_rolled_back,
    }
}

fn zero_gradient_moment_probe(config: AdamWConfig) -> AdamWParameterUpdate {
    let parameters = vec![parameter("probe.weight", &[1], &[3.0])];
    let groups = AdamWParameterGroups::new(["probe.weight"], std::iter::empty::<&str>())
        .expect("the probe weight belongs to the decay group");
    AdamW::with_parameter_groups(config, groups)
        .step_with_trace(&parameters)
        .expect("the live leaf starts with an exact zero gradient")
        .updates()[0]
        .clone()
}

fn rejected_set_probe(
    committed_optimizer: &AdamW,
    committed_parameters: &[NamedParameter],
) -> (AdamWError, bool) {
    let mut optimizer = committed_optimizer.clone();
    let optimizer_before = optimizer.clone();
    let mut parameters = committed_parameters.to_vec();
    parameters.push(parameter("unexpected.weight", &[1], &[1.0]));
    let parameters_before = parameters
        .iter()
        .map(|parameter| {
            (
                parameter.tensor().clone(),
                parameter.tensor().value().as_slice().to_vec(),
                parameter
                    .tensor()
                    .gradient()
                    .map(|gradient| gradient.as_slice().to_vec()),
            )
        })
        .collect::<Vec<_>>();
    let error = optimizer
        .step(&parameters)
        .expect_err("a changed named set must be rejected");
    let parameters_unchanged =
        parameters
            .iter()
            .zip(parameters_before)
            .all(|(parameter, (leaf, values, gradient))| {
                parameter.tensor().is_same_node(&leaf)
                    && parameter.tensor().value().as_slice() == values
                    && parameter
                        .tensor()
                        .gradient()
                        .map(|actual| actual.as_slice().to_vec())
                        == gradient
            });
    (error, parameters_unchanged && optimizer == optimizer_before)
}

Файл rust/demos/ch22-adamw/src/main.rs выводит отчёт со всеми точными векторами и инвариантами:

Выведите детерминированный отчёт для учащегося по главе 22 rust/demos/ch22-adamw/src/main.rs#learner-adamw-output 
fn main() {
    print!("{}", ch22_adamw::learner_report());
}

Выполните cargo run --quiet --locked -p ch22-adamw, чтобы изучить точный первый шаг, состояние, хранящееся по стабильным именам, сохранение узлов параметров и исходных накопленных градиентов, а также откат при изменившемся наборе имён. Та же реализация поддерживает изменённый порядок предъявления параметров, накопленные градиенты, обновление только за счёт затухания при нулевых начальных моментах и движение за счёт сохранённых моментов после последующего нулевого градиента. При недопустимых значениях, идентичностях, формах, арифметических результатах или наборах параметров операция завершается с ошибкой, не изменяя ни одного существующего значения и ни одного поля состояния оптимизатора. Повторный детерминированный запуск 200 шагов на анизотропной квадратичной функции побитно совпадает с первым запуском и в конце даёт значение функции меньше 10−1210^{-12}.

Проследите, как затухание обходит ветвь моментов

Трассировка из четырнадцати строк содержит две записи параметров с указанием групп, по пять точек траекторий SGD и AdamW и одно доказательство свойств полного набора. Первая схема проводит вес из группы с затуханием через все моменты и поправки. В этой зафиксированной записи α1=1\alpha_1=1, поэтому поле входного градиента одновременно показывает исходный g1g_1 и градиент после масштабирования g~1\widetilde g_1. Во второй собраны параметр масштаба нормализации без затухания, обе траектории и доказательство атомарной операции. Все показанные числа сформированы программой на Rust:

Выведите точные данные о моментах, поправках, атомарной записи и откате rust/demos/ch22-adamw/src/diagram_trace.rs#adamw-trace 
pub fn diagram_trace() -> String {
    let evidence = learner_evidence();
    let mut lines = vec![format!(
        "META|step={}|learning_rate={:.6}|beta1={:.6}|beta2={:.6}|epsilon={:.6}|weight_decay={:.6}|first_correction={:.6}|second_correction={:.6}",
        evidence.step.step(),
        evidence.config.learning_rate(),
        evidence.config.beta1(),
        evidence.config.beta2(),
        evidence.config.epsilon(),
        evidence.config.weight_decay(),
        evidence.step.first_correction(),
        evidence.step.second_correction(),
    )];

    for (index, update) in evidence.step.updates().iter().enumerate() {
        lines.push(format!(
            "PARAM|index={index}|name={}|group={}|shape={:?}|before={}|gradient={}",
            update.name(),
            if update.decay_applied() {
                "decay"
            } else {
                "no_decay"
            },
            update.shape(),
            format_vector(update.before()),
            format_vector(update.gradient()),
        ));
        lines.push(format!(
            "MOMENT|index={index}|first={}|second={}|corrected_first={}|corrected_second={}",
            format_vector(update.first_moment()),
            format_vector(update.second_moment()),
            format_vector(update.corrected_first_moment()),
            format_vector(update.corrected_second_moment()),
        ));
        lines.push(format!(
            "DELTA|index={index}|adaptive={}|decay={}|after={}",
            format_vector(update.adaptive_delta()),
            format_vector(update.decay_delta()),
            format_vector(update.after()),
        ));
    }

    let trajectory = anisotropic_trajectory();
    lines.push(format!(
        "QUADRATIC|curvature=[1.000000, 4.000000]|steps={}",
        trajectory.len() - 1,
    ));
    lines.extend(trajectory.iter().map(|point| {
        format!(
            "POINT|step={}|sgd={}|adamw={}",
            point.step,
            format_vector(&point.sgd),
            format_vector(&point.adamw),
        )
    }));

    lines.push(format!(
        "PROOF|state_names={}|raw_gradients={}|parameter_nodes={}|zero_gradient_decay={:.6}|rollback={}|commit=atomic",
        evidence.state_names.join(","),
        if evidence.raw_gradients_retained {
            "retained"
        } else {
            "changed"
        },
        if evidence.parameter_nodes_preserved {
            "preserved"
        } else {
            "replaced"
        },
        evidence.zero_gradient_update.decay_delta()[0],
        if evidence.rejection_rolled_back {
            "unchanged"
        } else {
            "changed"
        },
    ));
    lines.join("\n") + "\n"
}

Проследите раздельные ветви AdamW — адаптивную и ветвь затухания

Проследите, как один именованный параметр из группы с затуханием проходит через моменты градиента и поправку на смещение, а затем получает отдельные адаптивную поправку и поправку затухания перед записью нового значения в существующий узел.

Зафиксированный шаг
t=1t=1
Скорость обучения
η=0.100000\eta=0.100000
Коэффициенты сглаживания
β1=0.500000\beta_1=0.500000 β2=0.500000\beta_2=0.500000
Стабилизатор знаменателя
ε=0.100000\varepsilon=0.100000
Затухание весов
λ=0.100000\lambda=0.100000

Обновите моменты и внесите поправку на смещение

В экспоненциально сглаживаемые моменты поступает только накопленный градиент функции потерь. Поправка на первом шаге точно восстанавливает текущий градиент и его квадрат.

1−β1t=0.5000001-\beta_1^t=0.500000 1−β2t=0.5000001-\beta_2^t=0.500000

Именованный параметр
decoder.output.weight
Группа параметра
Применить затухание decay
Форма
[2]
Значение до обновления θt−1\theta_{t-1}
[1,−2]\left[1,-2\right]
Накопленный градиент gtg_t
[0.2,−0.4]\left[0.2,-0.4\right]
Адаптивная ветвь
Первый момент mtm_t
[0.1,−0.2]\left[0.1,-0.2\right]
Второй нецентрированный момент vtv_t
[0.02,0.08]\left[0.02,0.08\right]
Первый момент с поправкой m^t\hat m_t
[0.2,−0.4]\left[0.2,-0.4\right]
Второй нецентрированный момент с поправкой v^t\hat v_t
[0.04,0.16]\left[0.04,0.16\right]
Адаптивная ветвь

Адаптивная поправка

[0.066667,−0.08]\left[0.066667,-0.08\right]
Значение до обновления θt−1\theta_{t-1} [1,−2]\left[1,-2\right]
Ветвь затухания в обход моментов · Применить затухание

Поправка затухания

[0.01,−0.02]\left[0.01,-0.02\right]
−
Вычесть обе поправки

Новое значение θt\theta_t

[0.923333,−1.9]\left[0.923333,-1.9\right]

Разделите адаптивную поправку и поправку затухания

Адаптивная ветвь использует моменты с поправкой. Ветвь затухания напрямую читает значение параметра до шага и не попадает ни в один момент.

Адаптивная ветвь ηm^t/(v^t+ε)\eta\hat m_t/(\sqrt{\hat v_t}+\varepsilon) − Ветвь затухания в обход моментов ηλθt−1\eta\lambda\theta_{t-1}

Сопоставьте группу без затухания, траектории оптимизаторов и атомарную запись

На параметре масштаба нормализации с нулевым градиентом проверьте, как выбор группы управляет затуханием; затем сравните точные точки SGD и AdamW и проверьте инварианты атомарного обновления всего набора параметров.

Не применять затухание

Именованный параметр
decoder.norm.scale
Группа параметра
Не применять затухание no_decay
Форма
[1]
Значение до обновления
θt−1=[0.5]\theta_{t-1}=\left[0.5\right]
Накопленный градиент
gt=[0]g_t=\left[0\right]
Первый момент с поправкой
m^t=[0]\hat m_t=\left[0\right]
Второй нецентрированный момент с поправкой
v^t=[0]\hat v_t=\left[0\right]
Адаптивная поправка
ηm^tv^t+ε=[0]\frac{\eta\hat m_t}{\sqrt{\hat v_t}+\varepsilon}=\left[0\right]
Поправка затухания
ηλθt−1=[0]\eta\lambda\theta_{t-1}=\left[0\right]
Новое значение
θt=[0.5]\theta_t=\left[0.5\right]

Сравните траектории при неодинаковой кривизне

Обе траектории вычислены на одной заданной квадратичной функции. Точные точки получены программой на Rust, а не восстановлены по рисунку.

Кривизна квадратичной функции q(x,y)=12(x2+4y2)q(x,y)=\frac12(x^2+4y^2) diag⁡(H)=[1,4]\operatorname{diag}(H)=\left[1,4\right]

Траектория SGD
  1. Шаг оптимизатора t=0t=0 [1,1]\left[1,1\right]
  2. Шаг оптимизатора t=1t=1 [0.9,0.6]\left[0.9,0.6\right]
  3. Шаг оптимизатора t=2t=2 [0.81,0.36]\left[0.81,0.36\right]
  4. Шаг оптимизатора t=3t=3 [0.729,0.216]\left[0.729,0.216\right]
  5. Шаг оптимизатора t=4t=4 [0.6561,0.1296]\left[0.6561,0.1296\right]
Траектория AdamW
  1. Шаг оптимизатора t=0t=0 [1,1]\left[1,1\right]
  2. Шаг оптимизатора t=1t=1 [0.899091,0.892439]\left[0.899091,0.892439\right]
  3. Шаг оптимизатора t=2t=2 [0.799889,0.786278]\left[0.799889,0.786278\right]
  4. Шаг оптимизатора t=3t=3 [0.702629,0.681677]\left[0.702629,0.681677\right]
  5. Шаг оптимизатора t=4t=4 [0.60758,0.578823]\left[0.60758,0.578823\right]

Запишите все проверенные значения в существующие узлы

Rust сначала готовит оба новых именованных значения, затем получает доступ на запись ко всему существующему набору и только после этого записывает оба значения в прежние узлы.

Проверьте сохранение узлов и градиентов, а также откат

Моменты связаны со стабильными именами, а не с порядком перечисления параметров. При изменении набора имён значения, градиенты и узлы всех параметров, а также все поля оптимизатора остаются без изменений.

Ключи состояния
decoder.norm.scaledecoder.output.weight
Накопленный градиент после шага AdamW
Без изменений
Тот же узел параметра
Сохранён
Затухание при нулевых моментах и градиенте
ηλθ=0.030000\eta\lambda\theta=0.030000
Попытка с другим набором имён
Без изменений
Успешная фиксация
Атомарно

В трассировке записаны точные имена, формы, векторы, знаменатели поправок и данные атомарной операции. В подробной схеме сплошная адаптивная ветвь и пунктирная ветвь затухания дополнительно различаются подписями и границами. Затем компактная запись группы показывает нулевую поправку затухания для параметра без затухания. Обе траектории используют одни и те же заданные точки функции q(x,y)=12(x2+4y2)q(x,y)=\frac12(x^2+4y^2). Проследите путь градиента функции потерь через моменты, сопоставьте его с отдельным затуханием, пропорциональным параметру, и объедините оба вклада при итоговом вычитании.

Сначала предскажите результат, затем запустите оптимизатор

  1. Выведите m1m_1, v1v_1, m^1\hat m_1 и v^1\hat v_1 для заданного вектора.
  2. До округления до шести знаков предскажите две поправки и θ1\theta_1.
  3. Создайте новое состояние оптимизатора для значения 33 из группы с затуханием при g1=0g_1=0 и предскажите новое значение. Затем объясните, почему после ненулевой предыстории более поздний нулевой градиент всё равно может дать вклад адаптивной ветви.
  4. Объясните, почему при изменении порядка двух записей параметры должны сохранить моменты, связанные с их стабильными именами.
  5. Предскажите результат, если первый кандидат уже успешно вычислен, а при вычислении второго возникает переполнение.
  6. После успешной фиксации предскажите градиент существующего листового узла, его идентичность и версию значения параметра. Решите, можно ли выполнить обратный проход по графу, построенному до обновления.
  7. Расположите по порядку обычный SGD, метод импульса, Adam со штрафом L2L_2, включённым в градиент, и AdamW, а затем укажите, какие слагаемые попадают в память оптимизатора.
  8. Объясните, почему к выходной матрице весов применяется затухание, а к масштабу нормализации — нет.
  9. Сопоставьте заданные траектории SGD и AdamW на квадратичной функции с разной кривизной по координатам.
Сверить предсказания
  1. m1=[0.1,−0.2]m_1=[0.1,-0.2], v1=[0.02,0.08]v_1=[0.02,0.08], m^1=[0.2,−0.4]\hat m_1=[0.2,-0.4], а v^1=[0.04,0.16]\hat v_1=[0.04,0.16].
  2. Адаптивная поправка приблизительно равна [0.066667,−0.08][0.066667,-0.08], поправка затухания — [0.01,−0.02][0.01,-0.02], а θ1≈[0.923333,−1.9]\theta_1\approx[0.923333,-1.9].
  3. При новых нулевых моментах адаптивная поправка остаётся нулевой, поправка затухания равна 0.030.03, а новое значение — 2.972.97. После ненулевой предыстории сохранённые моменты могут поддерживать ненулевую адаптивную поправку даже при нулевом текущем градиенте.
  4. Карта состояния использует стабильные имена; порядок в срезе влияет только на представление.
  5. Ничего не фиксируется: значения и градиенты обоих параметров, оба узла, а также всё состояние оптимизатора остаются без изменений.
  6. В существующем листовом узле остаётся накопленный градиент, который AdamW использовал в расчёте, а идентичность узла не меняется. Монотонно возрастающая версия значения параметра увеличивается на единицу. Граф, на ребре которого записана предыдущая версия, отклоняет обратный проход, поэтому вызывающий код явно обнуляет градиент и выполняет новый прямой проход.
  7. Обычный SGD использует только текущий градиент; метод импульса помнит затухающее направление; при штрафе L2L_2, включённом в градиент, пропорциональная параметру штрафная составляющая попадает в моменты Adam; AdamW оставляет затухание весов вне моментов. Бенжио задаёт отправную точку в виде нейронной языковой модели, а Туврон и соавторы описывают AdamW при предобучении LLaMA.
  8. Это настраиваемое правило группировки, принятое в курсе, а не следствие формулы AdamW. По этому правилу decoder.output.weight относится к группе с затуханием, поэтому AdamW вычитает из него пропорциональную параметру поправку ηλθt−1\eta\lambda\theta_{t-1}. Параметр decoder.norm.scale относится к группе без затухания, и его эффективный коэффициент λ\lambda равен 00: так затухание не создаёт отдельную поправку, напрямую стягивающую обучаемый масштаб нормализации к нулю.
  9. SGD гораздо быстрее уменьшает координату с высокой кривизной; AdamW масштабирует координаты с учётом истории их моментов и дополнительно применяет отдельный вклад затухания.

Далее обучите первую цельную нейронную языковую модель

Теперь общий путь обучения умеет связывать каждый накопленный градиент функции потерь, усреднённой по токенам, со стабильным именем соответствующего параметра, сохранять под этим именем первый и второй моменты между шагами и атомарно обновлять существующие листовые узлы параметров. AdamW оставляет каждый исходный градиент после расчёта обновления. В главе 23 эти градиенты будут явно обнуляться после каждого шага, а затем этот оптимизатор будет применён для обучения нейронной языковой модели с фиксированным контекстом и проверки снижения потерь на валидационной выборке.

AdamW изменяет параметры, но ещё не определяет модель, градиенты которой он получает. В следующей главе этот контур замкнётся, при этом разбиение данных на выборки, автоматическое дифференцирование, пакетная обработка, имена параметров и состояние оптимизатора останутся явными.