Giới thiệu Mô hình Không gian Trạng thái (SSM)

Giới thiệu Mô hình Không gian Trạng thái (SSM)

  • 10 min read
Giới thiệu Mô hình Không gian Trạng thái (SSM)
Giới thiệu Mô hình Không gian Trạng thái (SSM)

Giới thiệu về Mô hình Không gian Trạng thái (SSM)

Bài viết này trình bày về Mô hình Không gian Trạng thái (SSM) trong lĩnh vực học sâu, dựa trên mô hình S4.

Tác giả: Loïck BOURDOIS

Ngày: 19 tháng 7, 2024


Lời mở đầu

Tôi xin gửi lời cảm ơn chân thành đến Boris ALBAR, Pierre BEDU và Nicolas PREVOT vì đã đồng ý thành lập một nhóm làm việc về chủ đề SSM và đồng hành cùng tôi trong quá trình khám phá loại mô hình này. Đặc biệt cảm ơn Boris ALBAR đã dành thời gian hiệu đính bài viết này.


Giới thiệu

Mô hình Không gian Trạng thái (SSM) ban đầu được sử dụng trong lý thuyết điều khiển để mô hình hóa một hệ thống động lực thông qua các biến trạng thái.

Năm 2018, Aaron R. VOELKER và Chris ELIASMITH đã nghiên cứu cách bộ não biểu diễn thông tin thời gian một cách hiệu quả. Trong bài báo “Improving Spiking Dynamical Networks: Accurate Delays, Higher-Order Synapses, and Time Cells”, họ phát hiện ra rằng SSM là một mô hình tuyệt vời để mô tả các “tế bào thời gian” có trong não bộ.

Từ lĩnh vực khoa học thần kinh, họ đã áp dụng công trình của mình vào học sâu và trở thành những người đầu tiên sử dụng SSM trong học sâu. Bài viết này sẽ tập trung vào việc định nghĩa các khái niệm cơ bản của SSM trong học sâu, dựa trên mô hình S4 được giới thiệu bởi Albert GU và cộng sự vào năm 2021. Mặc dù S4 không phải là mô hình được sử dụng trực tiếp trong thực tế (hiện nay đã có các SSM hiệu quả hơn hoặc dễ triển khai hơn), chúng tôi sử dụng nó vì mục đích giáo dục.


Định nghĩa SSM trong học sâu

Một SSM bao gồm ba biến phụ thuộc vào thời gian $t$:

  • $x(t) \in \mathbb{C}^{n}$: Biểu diễn $n$ biến trạng thái ẩn.
  • $u(t) \in \mathbb{C}^{m}$: Biểu diễn $m$ đầu vào trạng thái.
  • $y(t) \in \mathbb{C}^{p}$: Biểu diễn $p$ đầu ra.

Hệ thống này được xác định bởi bốn ma trận có thể học được: $\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D}$.

  • $\mathbf{A} \in \mathbb{C}^{n \times n}$: Ma trận trạng thái, điều khiển trạng thái ẩn $x(t)$.
  • $\mathbf{B} \in \mathbb{C}^{n \times m}$: Ma trận điều khiển.
  • $\mathbf{C} \in \mathbb{C}^{p \times n}$: Ma trận đầu ra.
  • $\mathbf{D} \in \mathbb{C}^{p \times m}$: Ma trận lệnh.

Hệ thống phương trình liên tục của SSM là:

$$ \begin{aligned} x’(t) &= \mathbf{A}x(t) + \mathbf{B}u(t) \ y(t) &= \mathbf{C}x(t) + \mathbf{D}u(t) \end{aligned} $$

Trong học sâu, $\mathbf{D}u = 0$ thường được xem là một kết nối bỏ qua (skip connection) dễ dàng tính toán. Do đó, hệ thống trở thành:

$$ \begin{aligned} x’ &= \mathbf{A}x + \mathbf{B}u \ y &= \mathbf{C}x \end{aligned} $$


Rời rạc hóa (Discretization)

Để đưa SSM vào máy tính, chúng ta cần rời rạc hóa hệ thống liên tục. Bước rời rạc hóa này là rất quan trọng vì nó cho phép chuyển đổi SSM sang hai dạng khác: dạng đệ quydạng tích chập.

Sử dụng phương pháp hình thang (trapezoid method) để rời rạc hóa:

$$ x_{n+1} - x_n = \frac{\Delta}{2}(f(t_n) + f(t_{n+1})) $$

Với $\Delta = t_{n+1} - t_n$ và $f(t) = \mathbf{A}x(t) + \mathbf{B}u(t)$.

Sau các phép biến đổi, ta có được SSM rời rạc:

$$ \begin{aligned} x_k &= \mathbf{\bar{A}}x_{k-1} + \mathbf{\bar{B}}u_k \ y_k &= \mathbf{\bar{C}}x_k \end{aligned} $$

Trong đó:

$$ \begin{aligned} \mathbf{\bar{A}} &= (\mathbf {I} - \frac{\Delta}{2} \mathbf{A})^{-1}(\mathbf {I} + \frac{\Delta}{2} \mathbf{A}) \ \mathbf {\bar{B}} &= (\mathbf{I} - \frac{\Delta}{2} \mathbf {A})^{-1} \Delta \mathbf{B} \ \mathbf {\bar{C}} &= \mathbf{C}\ \end{aligned} $$


Các quan điểm của SSM

Dạng đệ quy (Recursive View)

Dạng đệ quy của SSM là:

$$ \begin{aligned} x_k &= \mathbf{\bar{A}}x_{k-1} + \mathbf{\bar{B}}u_k \ y_k &= \mathbf{\bar{C}}x_k \end{aligned} $$

Dạng này tương tự như mạng RNN, có ưu điểm là xử lý tuần tự và có thể có ngữ cảnh không giới hạn, nhưng lại chậm trong việc huấn luyện do thiếu khả năng song song hóa và dễ gặp vấn đề về gradient.

Dạng tích chập (Convolutive View)

Dạng tích chập được biểu diễn qua phép chập của một kernel $K$:

$$ y_k = K \ast u $$

Với kernel tích chập $\mathbf{\bar{K}}_k = (\mathbf{\bar{C}} \mathbf{\bar{B}}, \mathbf{\bar{C}} \mathbf{\bar{A}} \mathbf{\bar{B}}, …, \mathbf{\bar{C}} \mathbf{\bar{A}}^{k} \mathbf{\bar{B}})$. Kernel này có kích thước bằng toàn bộ chuỗi đầu vào và có thể được tính toán hiệu quả bằng Biến đổi Fourier Nhanh (FFT).

Dạng này có ưu điểm là tính năng lượng cục bộ, dễ diễn giải và huấn luyện hiệu quả (song song hóa). Tuy nhiên, nó chậm trong các ngữ cảnh trực tuyến hoặc tự hồi quy.


Ưu điểm và hạn chế của các quan điểm

Quan điểm Ưu điểm Hạn chế
Liên tục - Xử lý tự động dữ liệu liên tục (âm thanh, chuỗi thời gian).- Phân tích toán học khả thi. - Rất chậm cho cả huấn luyện và suy luận.
Đệ quy - Có thiên hướng tự nhiên cho dữ liệu tuần tự.- Suy luận hiệu quả (cập nhật trạng thái trong thời gian không đổi). - Huấn luyện chậm (thiếu song song hóa).- Dễ gặp vấn đề về biến mất hoặc bùng nổ gradient với chuỗi dài.
Tích chập - Đặc trưng cục bộ, dễ diễn giải.- Huấn luyện hiệu quả (song song hóa). - Chậm trong ngữ cảnh trực tuyến hoặc tự hồi quy.- Kích thước ngữ cảnh cố định.

Học ma trận

Ma trận $\mathbf{A}$ trong dạng tích chập, $\mathbf{A}^k$, có thể được tính toán hiệu quả nếu $\mathbf{A}$ là ma trận đường chéo:

$$ \mathbf{A} = \begin{bmatrix} \lambda_{1} & 0 & \cdots & 0 \ 0 & \lambda_{2} & \cdots & 0 \ \vdots & \vdots & \ddots & \vdots \ 0 & 0 & \cdots & \lambda_{n} \end{bmatrix} \Rightarrow \mathbf{A}^k = \begin{bmatrix} \lambda_{1}^k & 0 & \cdots & 0 \ 0 & \lambda_{2}^k & \cdots & 0 \ \vdots & \vdots & \ddots & \vdots \ 0 & 0 & \cdots & \lambda_{n}^k \end{bmatrix} $$

Ma trận HiPPO (High-Order Polynomial Projection Operator) là một lựa chọn phổ biến để khởi tạo ma trận $\mathbf{A}$, giúp mô hình hội tụ tốt hơn.


Kết quả thực nghiệm

S4 đã cho thấy hiệu suất ấn tượng trên nhiều tác vụ:

  • Âm thanh (Speech Commands): Hiệu suất vượt trội so với ConvNet với số lượng tham số tương đương, khả năng xử lý resampling tốt.
  • Chuỗi thời gian (Time Series): Vượt trội hơn Informer trên nhiều cấu hình, xử lý tốt cả dữ liệu đơn biến và đa biến.
  • Thị giác (sCIFAR-10): Đạt SoTA với số lượng tham số thấp.
  • Văn bản (Long Range Arena): Là mô hình đầu tiên giải quyết thành công tác vụ Path-X với độ dài 16K token, chứng tỏ khả năng xử lý chuỗi rất dài.

Tuy nhiên, S4 có thể có độ phức tạp (perplexity) cao hơn transformer trên các tác vụ văn bản như WikiText-103, có thể do tính chất không liên tục của dữ liệu văn bản.


Kết luận

SSM là một kiến trúc mô hình mạnh mẽ với ba quan điểm khác nhau (liên tục, đệ quy, tích chập), cho phép linh hoạt trong xử lý dữ liệu và tối ưu hóa hiệu suất. Khả năng xử lý chuỗi dài và hiệu quả tính toán là những ưu điểm nổi bật. Sự khác biệt giữa các kiến trúc SSM chủ yếu nằm ở phương pháp rời rạc hóa và cách khởi tạo ma trận $\mathbf{A}$.


Tài nguyên bổ sung

Lịch sử SSM

  • LMU (2019): Tiền thân của S4, tập trung vào việc xử lý chuỗi dài trong học sâu và khoa học thần kinh.
  • Legendre Memory Units (2021): Nghiên cứu về việc song song hóa huấn luyện SSM và cho thấy hiệu quả trên tác vụ văn bản.

Tài nguyên S4

  • Video: Các buổi nói chuyện của Albert GU và Sasha RUSH về S4.
  • Mã nguồn: “The Annotated S4” (Jax) và mã nguồn chính thức của S4 (PyTorch).
  • Bài viết trên blog: Chuỗi bài viết từ Hazy Research.
  • Sách: Phụ lục E trong sách “Hands-On Machine Learning with Scikit-Learn and Pytorch” của Aurélien Geron.

Tài nguyên HiPPO

  • Bài viết trên blog Hazy Research: Giới thiệu về ma trận HiPPO.
  • Bài báo: “How to Train Your HiPPO” của Albert GU và cộng sự.

Tham khảo


Bản dịch tiếng Pháp có sẵn trên blog của tôi: blog


Tuyên bố miễn trừ trách nhiệm: Bài viết này được dịch và biên soạn dựa trên nội dung gốc. Mọi sai sót hoặc hiểu lầm trong quá trình dịch thuật đều không cố ý.



Tóm tắt thay đổi:

  • 14/12/2023: Phát hành bài viết.
  • 08/04/2024: Sửa lỗi chính tả.
  • 11/06/2024: Thêm liên kết đến bài viết thứ hai trong chuỗi bài viết về SSM.
  • 18/07/2024: Sửa lỗi chính tả LaTex.
  • 23/09/2024: Viết lại phần giới thiệu và thêm phần về nguồn gốc của SSM trong học sâu.
  • 26/10/2025: Làm rõ một số cấu trúc câu theo phản hồi của người đọc.
  • 01/11/2025: Thêm một số tài nguyên (Phụ lục E của cuốn sách “Hands-On Machine Learning with Scikit-Learn and Pytorch” của Aurélien Geron).

Cộng đồng

Lịch sử SSM

Bài viết “LMU” của VOELKER, KAJIĆ và ELIASMITH (2019) được coi là tiền thân của S4. Bài báo này khám phá cách áp dụng SSM vào học sâu, đặc biệt là trong việc xử lý các chuỗi dài, vượt trội hơn LSTM và các mô hình khác về khả năng xử lý hàng triệu bước.

Tài nguyên SSM

  • Khóa học động lực học: Bài giảng của Ion HAZYUK cung cấp kiến thức về mô hình không gian trạng thái.
  • Luận án tiến sĩ: Của Albert GU và Aaron R. VOELKER.

Tài nguyên S4

  • Video: Các bài thuyết trình của Albert GU và Sasha RUSH.
  • Mã nguồn: “The Annotated S4” (Jax) và triển khai chính thức S4 (PyTorch).
  • Blog: Loạt bài viết về S4 từ Hazy Research.
  • Sách: Phụ lục E của “Hands-On Machine Learning with Scikit-Learn and Pytorch” bởi Aurélien Geron.

Tài nguyên HiPPO

  • Bài viết trên blog Hazy Research: Giới thiệu về ma trận HiPPO.
  • Bài báo: “How to Train Your HiPPO” của Albert GU và cộng sự.

Recommended for You

Bảng xếp hạng so sánh- AI - Từ lượt bình chọn của người dùng đến bảng xếp hạng mô hình dựa trên sự tham gia

Bảng xếp hạng so sánh- AI - Từ lượt bình chọn của người dùng đến bảng xếp hạng mô hình dựa trên sự tham gia

Bảng xếp hạng so sánh- AI - Từ lượt bình chọn của người dùng đến bảng xếp hạng mô hình dựa trên sự tham gia

Chạy các Mô hình Transformer Lớn trên Thiết bị Di động và Biên

Chạy các Mô hình Transformer Lớn trên Thiết bị Di động và Biên

Chạy các Mô hình Transformer Lớn trên Thiết bị Di động và Biên