Zygote.jl là thư viện vi phân nguồn-tới-nguồn (source-to-source automatic differentiation) dành cho ngôn ngữ Julia. Nó đóng vai trò là backend tính toán cho Flux.jl, xử lý các tác vụ liên quan đến gradient, ma trận Jacobi và các quy trình tối ưu hóa trong huấn luyện mô hình học máy.

Để bắt đầu, người dùng cần cài đặt gói thông qua trình quản lý gói của Julia:

using Pkg
Pkg.add("Zygote")

Tính toán Gradient cơ bản

Để tính đạo hàm của hàm số $f(x) = 3x^3 + 2x$ tại điểm $x = 2$:

using Zygote
gradient(x -> 3*x^3 + 2*x, 2)

Kết quả trả về là một tuple chứa giá trị đạo hàm: (38.0,).

Đối với hàm nhiều biến $f(u, v) = u^2 + v^2$, ta có thể truyền các tham số trực tiếp:

gradient((u, v) -> u^2 + v^2, 3.0, 4.0)

Đầu ra sẽ là (6.0, 8.0), tương ứng với các phần tử của vector gradient.

Zygote cũng hỗ trợ các phép toán trên ma trận. Giả sử bạn có ma trận trọng số $W$ và vector đầu vào $x$:

W = rand(3, 4)
x = rand(4)

# Hàm tính toán mất mát đơn giản
compute_loss(W) = sum((W * x) .^ 2)

# Tính gradient theo W
gradient(compute_loss, W)[1]

Lệnh trên trả về một ma trận kích thước $3 \times 4$ chứa các giá trị đạo hàm riêng.

Xử lý cấu trúc điều khiển

Thư viện có khả năng phân tích gradient qua các vòng lặp. Xét hàm tính tổng cộng dồn:

function accumulate_value(val, n)
    total = 0
    for i = 1:n
        total += val
    end
    return total
end

gradient(v -> accumulate_value(v, 5), 10)

Kết quả là (5.0,) vì hàm thực hiện phép cộng biến số 5 lần.

Sử dụng Dictionary

Gradient cũng có thể được tính toán trong các khối lệnh sử dụng Dict:

data_store = Dict()

gradient(5) do input
    data_store[:key] = input
    data_store[:key] * 3
end

Kết quả trả về là (3.0,).

Định nghĩa kiểu dữ liệu tùy chỉnh (Custom Types)

Dưới đây là ví dụ định nghĩa cấu trúc Vector2D và nạp chồng các toán tử:

import Base: +, -

struct Vector2D
    u::Float64
    v::Float64
end

+(a::Vector2D, b::Vector2D) = Vector2D(a.u + b.u, a.v + b.v)
-(a::Vector2D, b::Vector2D) = Vector2D(a.u - b.u, a.v - b.v)
calc_norm(p::Vector2D) = sqrt(p.u^2 + p.v^2)

p1 = Vector2D(1.0, 2.0)
p2 = Vector2D(3.0, 4.0)

# Tính gradient của hàm chuẩn vector
gradient(p -> calc_norm(p + p2), p1)[1]

Ma trận Jacobi

Đối với các ánh xạ từ $\mathbb{R}^n \to \mathbb{R}^m$, hàm jacobian được sử dụng. Ví dụ với ánh xạ $f(z) = [z_1^2, z_2^2]$:

jacobian(z -> [z[1]^2, z[2]^2], [2, 3])[1]

Kết quả là ma trận $2 \times 2$:
2.0 0.0
0.0 6.0

Ma trận Hessian

Để tính toán ma trận Hessian của hàm hai biến $f(x, y) = x^2 + y^2$:

hessian(v -> v[1]^2 + v[2]^2, randn(2))

Kết quả trả về ma trận đơn vị nhân với 2:
2.0 0.0
0.0 2.0

Lưu ý: Zygote được tối ưu hóa cho các kiểu dữ liệu số và mảng. Việc tính toán vi phân trên các kiểu dữ liệu khác như String hoặc Tuple sẽ thường xuyên trả về kết quả là Nothing.