Trong các kiến trúc học sâu hiện đại, bên cạnh việc sử dụng các phép biến đổi tuyến tính hay hàm kích hoạt cố định, một hướng tiếp cận mạnh mẽ là định nghĩa một lớp (layer) như là nghiệm của một bài toán tối ưu hóa. Cần phân biệt rõ ràng: "tối ưu hóa" ở đây không ám chỉ quá trình huấn luyện mạng (như cập nhật trọng số qua gradient descent), mà là việc giải các bài toán tối ưu hóa quy mô nhỏ ngay trong lượt truyền xuôi (forward pass) của chính lớp đó.
Bài Toán Tối Ưu Hóa Tổng Quát
Một bài toán tối ưu hóa tổng quát có thể được biểu diễn dưới dạng:
Trong đó \(z \in \mathbb{R}^n\) là biến tối ưu, \(x \in \mathbb{R}^m\) là dữ liệu đầu vào, \(f : \mathbb{R}^n \times \mathbb{R}^m \rightarrow \mathbb{R}\) là hàm mục tiêu và \(\mathcal{C}(x) \subseteq \mathbb{R}^n\) là tập ràng buộc. Lớp tối ưu hóa sẽ ánh xạ đầu vào \(x\) tới nghiệm tối ưu \(z^\star\):
Các lớp tối ưu hóa khả vi có nhiều ứng dụng trong tối ưu hóa hai mức, mạng khai báo sâu (Deep Declarative Networks) và điều khiển tối ưu. Hơn nữa, nhiều hàm phi tuyến phổ biến trong học sâu thực chất có thể được diễn giải dưới dạng nghiệm của các bài toán tối ưu hóa cụ thể.
Tối Ưu Hóa Lồi Và Điều Kiện KKT
Để đảm bảo tính khả thi và hiệu quả tính toán, chúng ta thường giới hạn ở các bài toán tối ưu lồi. Một bài toán lồi chuẩn tắc có dạng:
Để bài toán mang tính lồi, \(f(z)\) và các thành phần của \(g(z)\) phải là hàm lồi, trong khi \(h(z)\) phải là hàm affine. Đối với các bài toán có ràng buộc, điều kiện Karush-Kuhn-Tucker (KKT) cung cấp các điều kiện cần và đủ cho tính tối ưu. Hệ phương trình KKT yêu cầu tồn tại \((z^\star, \lambda^\star, \nu^\star)\) thỏa mãn:
Trong đó \(\lambda\) và \(\nu\) là các biến đối ngẫu, và \(\circ\) ký hiệu phép nhân Hadamard.
Giải Quyết Bài Toán Lồi Trong Thực Tế
Thư viện cvxpy cung cấp giao diện trực quan để mô hình hóa và giải các bài toán lồi. Dưới đây là ví dụ giải bài toán bình phương tối thiểu không âm có chuẩn hóa \(\ell_1\):
import cvxpy as cp
import numpy as np
num_samples, num_features = 60, 15
data_matrix = np.random.randn(num_samples, num_features)
target_vector = np.random.randn(num_samples)
weights = cp.Variable(num_features)
regularization_strength = 0.05
objective_fn = cp.sum_squares(data_matrix @ weights - target_vector) + regularization_strength * cp.norm1(weights)
constraints = [weights >= 0]
problem = cp.Problem(cp.Minimize(objective_fn), constraints)
problem.solve(verbose=False, eps_abs=1e-7, eps_rel=1e-7)
print("Nghiệm tối ưu:", np.round(weights.value, 4))
print("Biến đối ngẫu:", np.round(constraints[0].dual_value, 4))
Lớp Tối Ưu Hóa Lồi Khả Vi
Để làm cho lớp tối ưu hóa có thể lan truyền ngược (backpropagation), ta xem xét bài toán trong đó các hàm \(f, g, h\) phụ thuộc vào cả biến tối ưu \(z\) và đầu vào \(x\). Dựa trên Định lý hàm ẩn, ta có thể vi phân các điều kiện KKT. Xét phần đẳng thức của điều kiện KKT:
Tại điểm tối ưu, ta có thể tính đạo hàm của nghiệm theo \(x\) bằng cách giải hệ phương trình tuyến tính:
Ma trận Jacobian \(\partial_{z, \lambda, \nu} G\) có cấu trúc khối như sau:
Triển Khai Lớp Tối Ưu Hóa Với PyTorch
Chúng ta có thể tự xây dựng một lớp khả vi tùy chỉnh bằng cách kết hợp cvxpy để giải bài toán và torch.autograd để tính toán Jacobian trong quá trình lan truyền ngược.
import torch
import torch.nn as nn
import torch.autograd as autograd
from itertools import accumulate
import cvxpy as cp
class DifferentiableOptLayer(nn.Module):
def __init__(self, opt_vars, input_params, obj_fn, ineq_fns, eq_fns, **solver_kwargs):
super().__init__()
self.opt_vars = opt_vars
self.input_params = input_params
self.obj_fn = obj_fn
self.ineq_fns = ineq_fns
self.eq_fns = eq_fns
self.solver_kwargs = solver_kwargs
self.cp_ineqs = [fn(*opt_vars, *input_params) <= 0 for fn in ineq_fns]
self.cp_eqs = [fn(*opt_vars, *input_params) == 0 for fn in eq_fns]
self.prob = cp.Problem(cp.Minimize(obj_fn(*opt_vars, *input_params)),
self.cp_ineqs + self.cp_eqs)
def forward(self, *batch_inputs):
outputs = []
batch_size = batch_inputs[0].shape[0]
for b_idx in range(batch_size):
current_params = [p[b_idx] for p in batch_inputs]
with torch.no_grad():
for idx, param in enumerate(self.input_params):
param.value = current_params[idx].double().numpy()
self.prob.solve(**self.solver_kwargs)
z_vals = [torch.tensor(v.value).type_as(current_params[0]) for v in self.opt_vars]
lam_vals = [torch.tensor(c.dual_value).type_as(current_params[0]) for c in self.cp_ineqs]
nu_vals = [torch.tensor(c.dual_value).type_as(current_params[0]) for c in self.cp_eqs]
def flatten(*tensors):
return torch.cat([t.reshape(-1) for group in tensors for t in group])
def unflatten(flat_vec):
sizes = [0] + list(accumulate([t.numel() for group in [z_vals, lam_vals, nu_vals] for t in group]))
chunks = [flat_vec[sizes[i]:sizes[i+1]] for i in range(len(sizes)-1)]
z_out = [chunks[i].view_as(z_vals[i]) for i in range(len(z_vals))]
lam_out = [chunks[i+len(z_vals)].view_as(lam_vals[i]) for i in range(len(lam_vals))]
nu_out = [chunks[i+len(z_vals)+len(lam_vals)].view_as(nu_vals[i]) for i in range(len(nu_vals))]
return z_out, lam_out, nu_out
def calc_kkt(z, lam, nu, *params):
g_vals = [fn(*z, *params) for fn in self.ineq_fns]
h_vals = [fn(*z, *params) for fn in self.eq_fns]
lagrangian = (self.obj_fn(*z, *params) +
sum((l * g).sum() for l, g in zip(lam, g_vals)) +
sum((n * h).sum() for n, h in zip(nu, h_vals)))
grad_z = autograd.grad(lagrangian, z, create_graph=True)
grad_lam = [lam[i] * g_vals[i] for i in range(len(lam))]
return grad_z, grad_lam, h_vals
flat_solution = flatten(z_vals, lam_vals, nu_vals)
z_grad = [z_.clone().detach().requires_grad_() for z_ in z_vals]
kkt_res = calc_kkt(z_grad, lam_vals, nu_vals, *current_params)
# Gắn lại vào biểu đồ tính toán
y = flat_solution - flatten(*kkt_res)
# Tính Jacobian cho backward pass
jacobian = autograd.functional.jacobian(
lambda vec: flatten(*calc_kkt(*unflatten(vec), *current_params)), y
)
y.register_hook(lambda grad, J=jacobian: torch.linalg.solve(J.T, grad[:, None])[:, 0])
outputs.append(unflatten(y)[0])
final_out = [torch.stack(o, dim=0) for o in zip(*outputs)]
return final_out[0] if len(final_out) == 1 else tuple(final_out)
Ứng Dụng: Quy Hoạch Toàn Phương (QP)
Lớp trên có thể dễ dàng cấu hình để giải các bài toán QP khả vi, một thành phần cốt lõi trong nhiều mô hình học sâu có ràng buộc.
dim_z, dim_ineq, dim_eq = 12, 5, 6
var_z = cp.Variable(dim_z)
param_P_sqrt = cp.Parameter((dim_z, dim_z))
param_q = cp.Parameter(dim_z)
param_G = cp.Parameter((dim_ineq, dim_z))
param_h = cp.Parameter(dim_ineq)
param_A = cp.Parameter((dim_eq, dim_z))
param_b = cp.Parameter(dim_eq)
def qp_objective(z, P_sqrt, q, G, h, A, b):
if isinstance(z, cp.Variable):
return 0.5 * cp.sum_squares(P_sqrt @ z) + q @ z
return 0.5 * torch.sum((P_sqrt @ z)**2) + q @ z
def qp_ineq(z, P_sqrt, q, G, h, A, b): return G @ z - h
def qp_eq(z, P_sqrt, q, G, h, A, b): return A @ z - b
qp_layer = DifferentiableOptLayer(
opt_vars=[var_z],
input_params=[param_P_sqrt, param_q, param_G, param_h, param_A, param_b],
obj_fn=qp_objective, ineq_fns=[qp_ineq], eq_fns=[qp_eq],
solver=cp.OSQP, eps_abs=1e-7, eps_rel=1e-7
)
Tái Tạo Các Hàm Kích Hoạt Phi Tuyến
Nhiều hàm kích hoạt quen thuộc thực chất là nghiệm của các bài toán tối ưu lồi đơn giản.
Sigmoid và Softmax
Hàm Sigmoid tương đương với việc tối thiểu hóa một hàm tuyến tính cộng với phạt entropy nhị phân. Tương tự, Softmax là nghiệm của bài toán tối ưu với ràng buộc tổng bằng 1 và phạt entropy.
# Sigmoid Layer
var_s = cp.Variable(1)
param_x = cp.Parameter(1)
def sigmoid_obj(s, x):
if isinstance(s, cp.Variable):
return -cp.multiply(s, x) - cp.entr(s) - cp.entr(1 - s)
return -s * x + s * torch.log(s) + (1 - s) * torch.log(1 - s)
sigmoid_layer = DifferentiableOptLayer([var_s], [param_x], sigmoid_obj, [], [])
# Softmax Layer
var_sm = cp.Variable(8)
param_x_sm = cp.Parameter(8)
def softmax_obj(s, x):
return -s @ x - cp.sum(cp.entr(s)) if isinstance(s, cp.Variable) else -s @ x + s @ torch.log(s)
def softmax_eq(s, x):
return cp.sum(s) - 1 if isinstance(s, cp.Variable) else s.sum() - 1
softmax_layer = DifferentiableOptLayer([var_sm], [param_x_sm], softmax_obj, [], [softmax_eq])
Sparsemax
Sparsemax tạo ra các phân phối xác suất thưa (sparse) bằng cách chiếu đầu vào lên đơn hình (simplex). Thay vì viết một solver tùy chỉnh, ta có thể mô hình hóa nó trực tiếp qua lớp tối ưu hóa:
var_sp = cp.Variable(8)
param_x_sp = cp.Parameter(8)
def sparsemax_obj(s, x):
return cp.sum_squares(s - x) if isinstance(s, cp.Variable) else torch.sum((x - s)**2)
def sparsemax_ineq(s, x): return -s
def sparsemax_eq(s, x):
return cp.sum(s) - 1 if isinstance(s, cp.Variable) else s.sum() - 1
sparsemax_layer = DifferentiableOptLayer([var_sp], [param_x_sp], sparsemax_obj, [sparsemax_ineq], [sparsemax_eq])
Học Đa Diện Lồi Từ Dữ Liệu
Một ứng dụng nâng cao là học các ràng buộc hình học (đa diện lồi) trực tiếp từ dữ liệu. Giả sử ta có các điểm \(x\) và hình chiếu \(y\) của chúng lên một tập hợp lồi chưa biết, ta có thể tham số hóa các mặt của đa diện (\(G, h\)) và tối ưu hóa chúng bằng gradient descent.
class PolytopeProjector(nn.Module):
def __init__(self, input_dim, num_faces):
super().__init__()
self.normals = nn.Parameter(torch.randn(num_faces, input_dim))
self.offsets = nn.Parameter(torch.ones(num_faces))
def obj_fn(z, x, N, d):
return cp.sum_squares(z - x) if isinstance(z, cp.Variable) else torch.sum((x - z)**2)
def ineq_fn(z, x, N, d):
return N @ z - d
self.opt_layer = DifferentiableOptLayer(
[cp.Variable(input_dim)],
[cp.Parameter(input_dim), cp.Parameter((num_faces, input_dim)), cp.Parameter(num_faces)],
obj_fn, [ineq_fn], []
)
def forward(self, x):
batch_size = x.shape[0]
N_batch = self.normals.expand(batch_size, *self.normals.shape)
d_batch = self.offsets.expand(batch_size, *self.offsets.shape)
return self.opt_layer(x, N_batch, d_batch)
# Tạo dữ liệu giả lập và huấn luyện
num_points = 150
points_X = torch.randn(num_points, 2)
points_Y = points_X / points_X.norm(dim=1).clamp(min=1)[:, None]
torch.manual_seed(42)
projector = PolytopeProjector(2, 16)
optimizer = torch.optim.Adam(projector.parameters(), lr=0.05)
for epoch in range(50):
if epoch == 25:
optimizer.param_groups[0]["lr"] = 0.01
preds = projector(points_X)
loss = nn.MSELoss()(preds, points_Y)
optimizer.zero_grad()
loss.backward()
optimizer.step()