U-Net: Mạng nơ-ron tích chập cho phân đoạn ảnh y sinh

Hướng dẫn đọc

Vấn đề và động lực

Phân đoạn ảnh y sinh yêu cầu phân loại dày đặc ở cấp độ pixel (định vị) thay vì phân loại ở cấp độ toàn bộ hình ảnh. Lĩnh vực này bị chi phối bởi hai thách thức cơ bản:

1. Dữ liệu huấn luyện có gán nhãn rất khan hiếm, thường chỉ bao gồm vài chục hình ảnh thay vì hàng nghìn hay hàng triệu ảnh.
2. Các phương pháp học sâu trước đây, chẳng hạn như mạng tích chập cửa sổ trượt dựa trên mẩu ảnh (patch-based) (ví dụ: Cireşan et al.), gặp phải sự dư thừa lớn, suy luận chậm và sự đánh đổi cố hữu giữa ngữ cảnh mẩu ảnh và độ chính xác định vị.

Các tác giả đã tìm kiếm một mạng tích chập hoàn toàn (fully convolutional network) huấn luyện đầu-cuối có thể đồng thời nắm bắt ngữ cảnh toàn cục và chi tiết cục bộ tinh xảo, xử lý liền mạch các hình ảnh lớn tùy ý và huấn luyện thành công từ dữ liệu cực kỳ hạn chế.

Ý tưởng cốt lõi

• Kiến trúc chữ U đối xứng: Mở rộng các mạng tích chập hoàn toàn bằng cách kết hợp nhánh encoder co lại (giảm mẫu) với một nhánh decoder mở rộng (tăng mẫu) có năng lực biểu diễn tương đương, được trang bị nhiều kênh đặc trưng để truyền ngữ cảnh lên trên.
• Kết nối tắt (Skip Connections): Các bản đồ đặc trưng độ phân giải cao từ nhánh co lại được sao chép và nối với các bản đồ đặc trưng đã được tăng mẫu ở nhánh mở rộng (sau khi cắt xén), cung cấp các chi tiết định vị không gian chính xác đã bị mất trong quá trình pooling.
• Chiến lược ngói xếp chồng với tích chập Valid (Overlap-Tile Strategy): Mạng sử dụng các phép tích chập không đệm (unpadded). Các vùng biên được dự đoán liền mạch bằng cách chia ngói và ngoại suy ngữ cảnh bị thiếu thông qua lấy đối xứng qua gương, vượt qua giới hạn bộ nhớ GPU trên các ảnh lớn.
• Tăng cường biến dạng đàn hồi (Elastic Deformation Augmentation): Các biến dạng đàn hồi mượt mà và thực tế được tạo ra thông qua các lưới ngẫu nhiên thô giúp mạng học tính bất biến với biến dạng, làm giảm đáng kể số lượng mẫu huấn luyện cần thiết.
• Đánh trọng số mất mát biên phân tách (Separation Border Loss Weighting): Một bản đồ trọng số tính toán trước theo từng pixel sẽ phạt nặng các lỗi ở các đường biên nền nằm giữa các tế bào chạm nhau thuộc cùng một lớp, cho phép phân tách cá thể rõ ràng.

Phương pháp

Cấu trúc mạng

Mạng có 23 lớp tích chập được tổ chức thành hai nhánh:

• Nhánh co lại (Contracting Path): Các khối lặp lại gồm hai phép tích chập $3\times 3$ không đệm (mỗi phép theo sau bởi một ReLU) và một phép max pooling $2\times 2$ với bước trượt stride 2. Số lượng kênh đặc trưng tăng gấp đôi sau mỗi bước pooling.
• Nhánh mở rộng (Expansive Path): Các khối lặp lại gồm một phép tăng mẫu theo sau bởi một tích chập $2\times 2$ ("up-convolution") làm giảm một nửa số kênh đặc trưng, ghép nối với bản đồ đặc trưng đã cắt tương ứng từ nhánh co lại, và hai tích chập $3\times 3$ kèm theo ReLU.
• Lớp cuối cùng: Một tích chập $1\times 1$ ánh xạ các biểu diễn đặc trưng 64 kênh trực tiếp tới $K$ điểm số lớp mong muốn.

Hàm mất mát và Trọng số biên

Mạng được huấn luyện với hàm mất mát softmax cross-entropy theo từng pixel: $$p_k(\mathbf{x})=\frac{\exp (a_k(\mathbf{x}))}{{\sum }_{k^{\prime }=1}^K\exp (a_{k^{\prime }}(\mathbf{x}))}$$ $$E=\sum _{\mathbf{x}\in \Omega }w(\mathbf{x})\log (p_{\ell (\mathbf{x})}(\mathbf{x}))$$ trong đó $a_k(\mathbf{x})$ biểu thị giá trị kích hoạt cho kênh đặc trưng $k$ tại pixel $\mathbf{x}$, $\ell (\mathbf{x})$ là nhãn lớp chân thực, và $w(\mathbf{x})$ là bản đồ trọng số mất mát theo không gian.

Để cân bằng tần suất các lớp và phân tách các đối tượng chạm nhau, $w(\mathbf{x})$ được định nghĩa là: $$w(\mathbf{x})=w_c(\mathbf{x})+w_0\cdot \exp \left(-\frac{(d_1(\mathbf{x})+d_2(\mathbf{x}){)}^2}{2{\sigma }^2}\right)$$ trong đó $w_c(\mathbf{x})$ cân bằng tần suất lớp, $d_1(\mathbf{x})$ là khoảng cách đến biên của tế bào gần nhất, $d_2(\mathbf{x})$ là khoảng cách đến biên của tế bào gần thứ hai, các tác giả chọn $w_0=10$ và $\sigma \approx 5$ pixel.

Tối ưu hóa

• Huấn luyện bằng SGD trong Caffe với kích thước batch bằng 1 (một ngói ảnh lớn duy nhất) và xung lượng (momentum) cao ($0.99$).
• Trọng số được khởi tạo từ phân phối Gaussian với độ lệch chuẩn $\sqrt{2/N}$, trong đó $N$ là số lượng nút đi vào trên mỗi nơ-ron.

Kết quả

• Thử thách ISBI EM Segmentation (Cấu trúc nơ-ron):
  • Đạt thứ hạng cao nhất với lỗi biến dạng (warping error) là $0.000353$, Rand error là $0.0382$, và lỗi pixel là $0.0611$ (trung bình trên 7 phiên bản xoay, không có hậu xử lý bổ sung).
  • Vượt trội hơn phương pháp tốt nhất trước đó của Cireşan et al. (lỗi biến dạng $0.000420$, Rand error $0.0504$).
• Thử thách ISBI Cell Tracking Challenge 2015:
  • PhC-U373 (Độ tương phản pha): Đạt Intersection over Union (IOU) trung bình là $0.9203$, so với $0.83$ của thuật toán đứng thứ hai.
  • DIC-HeLa (Độ tương phản giao thoa vi sai): Đạt IOU trung bình là $0.7756$, so với $0.46$ của thuật toán đứng thứ hai.
• Hiệu quả: Phân đoạn một hình ảnh $512\times 512$ mất chưa đầy một giây trên GPU hiện đại, với tổng thời gian huấn luyện mất khoảng 10 giờ trên GPU NVIDIA Titan (6 GB).

Hạn chế và các câu hỏi mở

• Tích chập không đệm và cắt xén: Do các phép tích chập ở chế độ valid (không đệm), ngói đầu ra nhỏ hơn ngói đầu vào, đòi hỏi phải lựa chọn cẩn thận kích thước ngói đầu vào chia hết cho $2$ qua tất cả các bước pooling và cần lấy đối xứng tại viền ảnh.
• Phân tách cá thể qua trọng số kinh nghiệm: Việc phân tách các cá thể chạm nhau dựa vào bản đồ trọng số hình thái học được thiết kế thủ công thay vì các nhánh nhận diện hay embedding cá thể đầu-cuối.
• Thiết kế dành riêng cho 2D: Bài báo tập trung vào các lát cắt 2D, để ngỏ việc điều chỉnh cho các tập dữ liệu thể tích 3D y sinh cho các nghiên cứu tiếp theo.